Support ThenBy/ThenByDescending after Include/ThenInclude

src/ExpressiveSharp.EntityFrameworkCore/Infrastructure/Internal/ExpressiveOptionsExtension.cs

@@ -94,7 +94,8 @@ public void ApplyServices(IServiceCollection services)
                 new RemoveNullConditionalPatterns(),
                 new FlattenTupleComparisons(),
                 new FlattenConcatArrayCalls(),
-                new FlattenBlockExpressions());
+                new FlattenBlockExpressions(),
+                new Transformers.RewriteThenByAfterInclude());
             if (extraTransformers.Length > 0)
                 options.AddTransformers(extraTransformers);
             return options;

src/ExpressiveSharp.EntityFrameworkCore/Transformers/RewriteThenByAfterInclude.cs

@@ -0,0 +1,71 @@
+using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore;
+
+namespace ExpressiveSharp.EntityFrameworkCore.Transformers;
+
+/// <summary>
+/// Rewrites <c>ThenBy</c>/<c>ThenByDescending</c> applied after <c>Include</c>/<c>ThenInclude</c>
+/// into the equivalent tree EF Core can translate: the ordering is applied to the ordered source
+/// beneath the include chain. The C# type system cannot express this shape directly, so the
+/// generated interceptors produce a cast node (<see cref="ExpressiveQueryableExtensions.AsOrdered{T}"/>)
+/// that this transformer resolves before EF Core sees the query.
+/// </summary>
+public sealed class RewriteThenByAfterInclude : ExpressionVisitor, IExpressionTreeTransformer
+{
+    public Expression Transform(Expression expression) => Visit(expression);
+
+    protected override Expression VisitMethodCall(MethodCallExpression node)
+    {
+        if (node.Method.DeclaringType != typeof(Queryable)
+            || node.Method.Name is not (nameof(Queryable.ThenBy) or nameof(Queryable.ThenByDescending)))
+            return base.VisitMethodCall(node);
+
+        var arguments = new Expression[node.Arguments.Count];
+        for (var i = 0; i < node.Arguments.Count; i++)
+            arguments[i] = Visit(node.Arguments[i]);
+
+        var source = arguments[0];
+        if (source is UnaryExpression { NodeType: ExpressionType.Convert } convert
+            && typeof(IOrderedQueryable).IsAssignableFrom(convert.Type))
+        {
+            source = convert.Operand;
+        }
+
+        var includes = new List<MethodCallExpression>();
+        var baseSource = source;
+        while (baseSource is MethodCallExpression call && IsIncludeCall(call))
+        {
+            includes.Add(call);
+            baseSource = call.Arguments[0];
+        }
+
+        if (includes.Count == 0)
+        {
+            return node.Update(node.Object, arguments);
+        }
+
+        if (!typeof(IOrderedQueryable).IsAssignableFrom(baseSource.Type))
+        {
+            throw new InvalidOperationException(
+                $"'{node.Method.Name}' was called after 'Include'/'ThenInclude' on a source that " +
+                "is not ordered. Call 'OrderBy'/'OrderByDescending' before 'Include', or apply " +
+                "the complete ordering after the includes.");
+        }
+
+        arguments[0] = baseSource;
+        Expression current = Expression.Call(node.Method, arguments);
+        for (var i = includes.Count - 1; i >= 0; i--)
+        {
+            var includeArguments = includes[i].Arguments.ToArray();
+            includeArguments[0] = current;
+            current = includes[i].Update(null, includeArguments);
+        }
+
+        return current;
+    }
+
+    private static bool IsIncludeCall(MethodCallExpression call)
+        => call.Method.DeclaringType == typeof(EntityFrameworkQueryableExtensions)
+           && call.Method.Name is nameof(EntityFrameworkQueryableExtensions.Include)
+               or nameof(EntityFrameworkQueryableExtensions.ThenInclude);
+}

src/ExpressiveSharp.Generator/PolyfillInterceptorGenerator.cs

@@ -619,7 +619,7 @@ private static string MethodId(string op, string fileTag, int line, int col)
         var typeAliases = new Dictionary<ITypeSymbol, string>(SymbolEqualityComparer.Default);
         var delegateFqns = new string[funcParamIndices.Count];
         string elemRef;
-        string castFqn;
+        string sourceRef;
         string typeParams;
         string returnRef;
         string interceptorParamList;
@@ -667,9 +667,11 @@ private static string MethodId(string op, string fileTag, int line, int col)
                 delegateFqns[fi] = funcFqnGenerics[fi];
             }
 
-            castFqn = isOrdered
-                ? $"global::System.Linq.IOrderedQueryable<{elemRef}>"
-                : $"global::System.Linq.IQueryable<{elemRef}>";
+            // ThenBy/ThenByDescending go through AsOrdered: a plain cast fails when the source's
+            // expression is not statically ordered (e.g. an EF Core Include after OrderBy).
+            sourceRef = isOrdered
+                ? $"global::ExpressiveSharp.ExpressiveQueryableExtensions.AsOrdered<{elemRef}>((global::System.Linq.IQueryable<{elemRef}>)(object)source)"
+                : $"(global::System.Linq.IQueryable<{elemRef}>)(object)source";
 
             if (isRewritableReturn)
             {
@@ -718,9 +720,9 @@ private static string MethodId(string op, string fileTag, int line, int col)
                         t.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat))) + ">";
             }
 
-            castFqn = isOrdered
-                ? $"global::System.Linq.IOrderedQueryable<{elemFqn}>"
-                : $"global::System.Linq.IQueryable<{elemFqn}>";
+            sourceRef = isOrdered
+                ? $"global::ExpressiveSharp.ExpressiveQueryableExtensions.AsOrdered<{elemFqn}>(source)"
+                : $"(global::System.Linq.IQueryable<{elemFqn}>)source";
 
             returnRef = isRewritableReturn
                 ? returnElemType!.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)
@@ -781,7 +783,7 @@ private static string MethodId(string op, string fileTag, int line, int col)
                         {
                 {{allBodies}}            return global::ExpressiveSharp.ExpressiveQueryableExtensions.AsExpressive(
                                 {{targetTypeFqn}}.{{methodName}}(
-                                    ({{castFqn}})source,
+                                    {{sourceRef}},
                                     {{queryableArgList}}));
                         }
 
@@ -795,7 +797,7 @@ private static string MethodId(string op, string fileTag, int line, int col)
                         {{interceptorParamList}})
                     {
             {{allBodies}}            return {{targetTypeFqn}}.{{methodName}}(
-                                ({{castFqn}})source,
+                                {{sourceRef}},
                                 {{queryableArgList}});
                     }
 
@@ -813,7 +815,7 @@ private static string MethodId(string op, string fileTag, int line, int col)
             {{allBodies}}            return (global::ExpressiveSharp.IExpressiveQueryable<{{returnRef}}>)(object)
                             global::ExpressiveSharp.ExpressiveQueryableExtensions.AsExpressive(
                                 {{targetTypeFqn}}.{{methodName}}(
-                                    ({{castFqn}})(object)source,
+                                    {{sourceRef}},
                                     {{queryableArgList}}));
                     }
 
@@ -827,7 +829,7 @@ private static string MethodId(string op, string fileTag, int line, int col)
                     {{interceptorParamList}})
                 {
         {{allBodies}}            return {{targetTypeFqn}}.{{methodName}}(
-                            ({{castFqn}})(object)source,
+                            {{sourceRef}},
                             {{queryableArgList}});
                 }
 

src/ExpressiveSharp/Extensions/ExpressiveQueryableExtensions.cs

@@ -1,4 +1,7 @@
+using System.Collections;
+using System.Collections.Generic;
 using System.Linq;
+using System.Linq.Expressions;
 
 namespace ExpressiveSharp
 {
@@ -20,5 +23,37 @@ public static IExpressiveQueryable<T> AsExpressive<T>(
             // etc.) remain observable through the returned reference.
             => source as IExpressiveQueryable<T>
                ?? new ExpressiveQueryableWrapper<T>(source);
+
+        /// <summary>
+        /// Presents a queryable as <see cref="IOrderedQueryable{T}"/> for the generated
+        /// <c>ThenBy</c>/<c>ThenByDescending</c> interceptors. When the underlying expression's
+        /// static type is not ordered (e.g. an EF Core <c>Include</c> call following
+        /// <c>OrderBy</c>), the expression is wrapped in a cast node so
+        /// <see cref="Queryable.ThenBy{TSource,TKey}"/> can compose a valid tree.
+        /// </summary>
+        public static IOrderedQueryable<T> AsOrdered<T>(this IQueryable<T> source)
+            => source is IOrderedQueryable<T> ordered
+               && typeof(IOrderedQueryable<T>).IsAssignableFrom(source.Expression.Type)
+                ? ordered
+                : new OrderedQueryableAdapter<T>(source);
+
+        private sealed class OrderedQueryableAdapter<T> : IOrderedQueryable<T>
+        {
+            private readonly IQueryable<T> _source;
+
+            public OrderedQueryableAdapter(IQueryable<T> source)
+            {
+                _source = source;
+                Expression = typeof(IOrderedQueryable<T>).IsAssignableFrom(source.Expression.Type)
+                    ? source.Expression
+                    : Expression.Convert(source.Expression, typeof(IOrderedQueryable<T>));
+            }
+
+            public Expression Expression { get; }
+            public Type ElementType => _source.ElementType;
+            public IQueryProvider Provider => _source.Provider;
+            public IEnumerator<T> GetEnumerator() => _source.GetEnumerator();
+            IEnumerator IEnumerable.GetEnumerator() => ((IEnumerable)_source).GetEnumerator();
+        }
     }
 }

tests/ExpressiveSharp.EntityFrameworkCore.IntegrationTests/Infrastructure/IncludeTestBase.cs

@@ -1,3 +1,4 @@
+using ExpressiveSharp.IntegrationTests.Scenarios.Store;
 using ExpressiveSharp.IntegrationTests.Scenarios.Store.Models;
 using Microsoft.EntityFrameworkCore;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -105,4 +106,48 @@ public async Task ChainedModifiers_Include_Execute()
         Assert.AreEqual(1, results.Count);
         Assert.IsNotNull(results[0].Customer);
     }
+
+    [TestMethod]
+    public async Task ThenBy_AfterInclude_ExecutesWithComposedOrdering()
+    {
+        var results = await Context.Orders.AsExpressiveDbSet()
+            .OrderBy(o => o.Status)
+            .Include(o => o.Customer)
+            .ThenBy(o => o.Id)
+            .ToListAsync();
+
+        var expectedIds = SeedData.Orders
+            .OrderBy(o => o.Status).ThenBy(o => o.Id)
+            .Select(o => o.Id).ToList();
+        CollectionAssert.AreEqual(expectedIds, results.Select(r => r.Id).ToList());
+        Assert.IsNotNull(results.Single(r => r.Id == 1).Customer);
+    }
+
+    [TestMethod]
+    public async Task ThenBy_AfterIncludeThenInclude_ExecutesWithComposedOrdering()
+    {
+        var results = await Context.Orders.AsExpressiveDbSet()
+            .OrderBy(o => o.Status)
+            .Include(o => o.Customer)
+            .ThenInclude(c => c!.Address)
+            .ThenBy(o => o.Id)
+            .ToListAsync();
+
+        var expectedIds = SeedData.Orders
+            .OrderBy(o => o.Status).ThenBy(o => o.Id)
+            .Select(o => o.Id).ToList();
+        CollectionAssert.AreEqual(expectedIds, results.Select(r => r.Id).ToList());
+        Assert.IsNotNull(results.Single(r => r.Id == 1).Customer!.Address);
+    }
+
+    [TestMethod]
+    public async Task ThenBy_AfterInclude_WithoutOrderBy_ThrowsActionableError()
+    {
+        var query = Context.Orders.AsExpressiveDbSet()
+            .Include(o => o.Customer)
+            .ThenBy(o => o.Id);
+
+        var ex = await Assert.ThrowsExactlyAsync<InvalidOperationException>(() => query.ToListAsync());
+        StringAssert.Contains(ex.Message, "OrderBy");
+    }
 }

tests/ExpressiveSharp.Generator.Tests/PolyfillInterceptorGenerator/OrderByTests.ThenBy_CastsToIOrderedQueryable.verified.txt

@@ -16,7 +16,7 @@ namespace ExpressiveSharp.Generated.Interceptors
             var __lambda = global::System.Linq.Expressions.Expression.Lambda<global::System.Func<global::TestNs.Order, string>>(i361d11c18_expr_0, i361d11c18_p_o);
             return global::ExpressiveSharp.ExpressiveQueryableExtensions.AsExpressive(
                 global::System.Linq.Queryable.ThenBy(
-                    (global::System.Linq.IOrderedQueryable<global::TestNs.Order>)source,
+                    global::ExpressiveSharp.ExpressiveQueryableExtensions.AsOrdered<global::TestNs.Order>(source),
                     __lambda));
         }
         [global::System.Runtime.CompilerServices.InterceptsLocationAttribute(/* scrubbed */)]

tests/ExpressiveSharp.Generator.Tests/PolyfillInterceptorGenerator/SingleLambdaQueryableTests.ThenByDescending_GeneratesInterceptor.verified.txt

@@ -16,7 +16,7 @@ namespace ExpressiveSharp.Generated.Interceptors
             var __lambda = global::System.Linq.Expressions.Expression.Lambda<global::System.Func<global::TestNs.Order, string>>(i361d11c18_expr_0, i361d11c18_p_o);
             return global::ExpressiveSharp.ExpressiveQueryableExtensions.AsExpressive(
                 global::System.Linq.Queryable.ThenByDescending(
-                    (global::System.Linq.IOrderedQueryable<global::TestNs.Order>)source,
+                    global::ExpressiveSharp.ExpressiveQueryableExtensions.AsOrdered<global::TestNs.Order>(source),
                     __lambda));
         }
         [global::System.Runtime.CompilerServices.InterceptsLocationAttribute(/* scrubbed */)]