From 05090b450071fec6dee092a0e180cfcce67dc3cb Mon Sep 17 00:00:00 2001
From: Fati Iseni <fati.iseni@pozitrongroup.com>
Date: Wed, 23 Apr 2025 14:08:35 +0200
Subject: [PATCH] Fix type inference for ThenInclude with casting.

---
 .../Evaluators/IncludeEvaluator.cs            |  7 ++---
 .../Builders/Builder_Include.cs               | 12 ++++----
 .../Expressions/IncludeExpressionInfo.cs      | 20 +++++++++++--
 .../Evaluators/IncludeEvaluatorTests.cs       |  2 +-
 .../Expressions/IncludeExpressionInfoTests.cs | 30 +++++++++++++++++--
 5 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/src/Ardalis.Specification.EntityFrameworkCore/Evaluators/IncludeEvaluator.cs b/src/Ardalis.Specification.EntityFrameworkCore/Evaluators/IncludeEvaluator.cs
index ce8b0ba4..87925228 100644
--- a/src/Ardalis.Specification.EntityFrameworkCore/Evaluators/IncludeEvaluator.cs
+++ b/src/Ardalis.Specification.EntityFrameworkCore/Evaluators/IncludeEvaluator.cs
@@ -41,7 +41,6 @@ private IncludeEvaluator() { }
     /// <inheritdoc/>
     public IQueryable<T> GetQuery<T>(IQueryable<T> query, ISpecification<T> specification) where T : class
     {
-        Type? previousReturnType = null;
         foreach (var includeExpression in specification.IncludeExpressions)
         {
             var lambdaExpr = includeExpression.LambdaExpression;
@@ -49,14 +48,12 @@ public IQueryable<T> GetQuery<T>(IQueryable<T> query, ISpecification<T> specific
             if (includeExpression.Type == IncludeTypeEnum.Include)
             {
                 var key = new CacheKey(typeof(T), lambdaExpr.ReturnType, null);
-                previousReturnType = lambdaExpr.ReturnType;
                 var include = _cache.GetOrAdd(key, CreateIncludeDelegate);
                 query = (IQueryable<T>)include(query, lambdaExpr);
             }
             else if (includeExpression.Type == IncludeTypeEnum.ThenInclude)
             {
-                var key = new CacheKey(typeof(T), lambdaExpr.ReturnType, previousReturnType);
-                previousReturnType = lambdaExpr.ReturnType;
+                var key = new CacheKey(typeof(T), lambdaExpr.ReturnType, includeExpression.PreviousPropertyType);
                 var include = _cache.GetOrAdd(key, CreateThenIncludeDelegate);
                 query = (IQueryable<T>)include(query, lambdaExpr);
             }
@@ -104,7 +101,7 @@ private static Func<IQueryable, LambdaExpression, IQueryable> CreateThenIncludeD
 
     private static bool IsGenericEnumerable(Type type, out Type propertyType)
     {
-        if (type.IsGenericType && typeof(IEnumerable).IsAssignableFrom(type))
+        if (type.IsGenericType && type.GetGenericTypeDefinition() == typeof(IEnumerable<>))
         {
             propertyType = type.GenericTypeArguments[0];
             return true;
diff --git a/src/Ardalis.Specification/Builders/Builder_Include.cs b/src/Ardalis.Specification/Builders/Builder_Include.cs
index ee7ec96e..918ef314 100644
--- a/src/Ardalis.Specification/Builders/Builder_Include.cs
+++ b/src/Ardalis.Specification/Builders/Builder_Include.cs
@@ -101,7 +101,7 @@ public static IIncludableSpecificationBuilder<T, TResult, TProperty> Include<T,
     {
         if (condition)
         {
-            var expr = new IncludeExpressionInfo(navigationSelector, IncludeTypeEnum.Include);
+            var expr = new IncludeExpressionInfo(navigationSelector);
             builder.Specification.Add(expr);
         }
 
@@ -139,7 +139,7 @@ public static IIncludableSpecificationBuilder<T, TProperty> Include<T, TProperty
     {
         if (condition)
         {
-            var expr = new IncludeExpressionInfo(navigationSelector, IncludeTypeEnum.Include);
+            var expr = new IncludeExpressionInfo(navigationSelector);
             builder.Specification.Add(expr);
         }
 
@@ -183,7 +183,7 @@ public static IIncludableSpecificationBuilder<TEntity, TResult, TProperty> ThenI
     {
         if (condition && !Specification<TEntity>.IsChainDiscarded)
         {
-            var expr = new IncludeExpressionInfo(navigationSelector, IncludeTypeEnum.ThenInclude);
+            var expr = new IncludeExpressionInfo(navigationSelector, typeof(TPreviousProperty));
             builder.Specification.Add(expr);
         }
         else
@@ -228,7 +228,7 @@ public static IIncludableSpecificationBuilder<TEntity, TProperty> ThenInclude<TE
     {
         if (condition && !Specification<TEntity>.IsChainDiscarded)
         {
-            var expr = new IncludeExpressionInfo(navigationSelector, IncludeTypeEnum.ThenInclude);
+            var expr = new IncludeExpressionInfo(navigationSelector, typeof(TPreviousProperty));
             builder.Specification.Add(expr);
         }
         else
@@ -275,7 +275,7 @@ public static IIncludableSpecificationBuilder<TEntity, TResult, TProperty> ThenI
     {
         if (condition && !Specification<TEntity>.IsChainDiscarded)
         {
-            var expr = new IncludeExpressionInfo(navigationSelector, IncludeTypeEnum.ThenInclude);
+            var expr = new IncludeExpressionInfo(navigationSelector, typeof(IEnumerable<TPreviousProperty>));
             builder.Specification.Add(expr);
         }
         else
@@ -320,7 +320,7 @@ public static IIncludableSpecificationBuilder<TEntity, TProperty> ThenInclude<TE
     {
         if (condition && !Specification<TEntity>.IsChainDiscarded)
         {
-            var expr = new IncludeExpressionInfo(navigationSelector, IncludeTypeEnum.ThenInclude);
+            var expr = new IncludeExpressionInfo(navigationSelector, typeof(IEnumerable<TPreviousProperty>));
             builder.Specification.Add(expr);
         }
         else
diff --git a/src/Ardalis.Specification/Expressions/IncludeExpressionInfo.cs b/src/Ardalis.Specification/Expressions/IncludeExpressionInfo.cs
index 0a5db08b..12816151 100644
--- a/src/Ardalis.Specification/Expressions/IncludeExpressionInfo.cs
+++ b/src/Ardalis.Specification/Expressions/IncludeExpressionInfo.cs
@@ -11,16 +11,32 @@ public class IncludeExpressionInfo
     /// </summary>
     public LambdaExpression LambdaExpression { get; }
 
+    /// <summary>
+    /// The type of the previously included entity.
+    /// </summary>
+    public Type? PreviousPropertyType { get; }
+
     /// <summary>
     /// The include type.
     /// </summary>
     public IncludeTypeEnum Type { get; }
 
-    public IncludeExpressionInfo(LambdaExpression expression, IncludeTypeEnum includeType)
+    public IncludeExpressionInfo(LambdaExpression expression)
+    {
+        _ = expression ?? throw new ArgumentNullException(nameof(expression));
+
+        LambdaExpression = expression;
+        PreviousPropertyType = null;
+        Type = IncludeTypeEnum.Include;
+    }
+
+    public IncludeExpressionInfo(LambdaExpression expression, Type previousPropertyType)
     {
         _ = expression ?? throw new ArgumentNullException(nameof(expression));
+        _ = previousPropertyType ?? throw new ArgumentNullException(nameof(previousPropertyType));
 
         LambdaExpression = expression;
-        Type = includeType;
+        PreviousPropertyType = previousPropertyType;
+        Type = IncludeTypeEnum.ThenInclude;
     }
 }
diff --git a/tests/Ardalis.Specification.EntityFrameworkCore.Tests/Evaluators/IncludeEvaluatorTests.cs b/tests/Ardalis.Specification.EntityFrameworkCore.Tests/Evaluators/IncludeEvaluatorTests.cs
index f13f02f2..57c5e90b 100644
--- a/tests/Ardalis.Specification.EntityFrameworkCore.Tests/Evaluators/IncludeEvaluatorTests.cs
+++ b/tests/Ardalis.Specification.EntityFrameworkCore.Tests/Evaluators/IncludeEvaluatorTests.cs
@@ -35,7 +35,7 @@ public void QueriesMatch_GivenInheritanceModel()
         var spec = new Specification<Bar>();
         spec.Query
             .Include(x => x.BarChildren)
-            .ThenInclude<Bar, BarChild, BarDerivedInfo>(x => (x as BarDerived)!.BarDerivedInfo);
+            .ThenInclude(x => (x as BarDerived)!.BarDerivedInfo);
 
         var actual = _evaluator
             .GetQuery(DbContext.Bars, spec)
diff --git a/tests/Ardalis.Specification.Tests/Expressions/IncludeExpressionInfoTests.cs b/tests/Ardalis.Specification.Tests/Expressions/IncludeExpressionInfoTests.cs
index 9ecab725..949ac89a 100644
--- a/tests/Ardalis.Specification.Tests/Expressions/IncludeExpressionInfoTests.cs
+++ b/tests/Ardalis.Specification.Tests/Expressions/IncludeExpressionInfoTests.cs
@@ -9,18 +9,44 @@ public record City(int Id);
     [Fact]
     public void Constructor_ThrowsArgumentNullException_GivenNullForLambdaExpression()
     {
-        var sut = () => new IncludeExpressionInfo(null!, IncludeTypeEnum.Include);
+        var sut = () => new IncludeExpressionInfo(null!);
 
         sut.Should().Throw<ArgumentNullException>().WithParameterName("expression");
+
+
+        sut = () => new IncludeExpressionInfo(null!, typeof(Customer));
+
+        sut.Should().Throw<ArgumentNullException>().WithParameterName("expression");
+    }
+
+    [Fact]
+    public void Constructor_ThrowsArgumentNullException_GivenNullForPreviousPropertyType()
+    {
+        Expression<Func<Customer, Address>> expr = x => x.Address;
+        var sut = () => new IncludeExpressionInfo(expr, null!);
+
+        sut.Should().Throw<ArgumentNullException>().WithParameterName("previousPropertyType");
     }
 
     [Fact]
     public void Constructor_GivenIncludeExpression()
     {
         Expression<Func<Customer, Address>> expr = x => x.Address;
-        var sut = new IncludeExpressionInfo(expr, IncludeTypeEnum.Include);
+        var sut = new IncludeExpressionInfo(expr);
 
         sut.Type.Should().Be(IncludeTypeEnum.Include);
         sut.LambdaExpression.Should().Be(expr);
     }
+
+    [Fact]
+    public void Constructor_GivenThenIncludeExpressionAndPreviousPropertyType()
+    {
+        Expression<Func<Address, City>> expr = x => x.City;
+        var previousPropertyType = typeof(Customer);
+        var sut = new IncludeExpressionInfo(expr, previousPropertyType);
+
+        sut.Type.Should().Be(IncludeTypeEnum.ThenInclude);
+        sut.LambdaExpression.Should().Be(expr);
+        sut.PreviousPropertyType.Should().Be(previousPropertyType);
+    }
 }
