Compile reducible custom expression node kinds

src/FastExpressionCompiler/FastExpressionCompiler.cs

@@ -649,6 +649,8 @@ public struct CompilerContext
             /// The map of inlined invocations collected in TryCollect and then used in TryEmit
             internal SmallMap4<InvocationExpression, Expression, RefEq<InvocationExpression>> InlinedLambdaInvocation;
 
+            internal ExpressionReductionCache ReducedExpressions;
+
             /// New or Call expressions containing the complex expression, e.g. inlined Lambda Invoke or Try with Finally
             internal SmallSet4<Expression, RefEq<Expression>> ArgsContainingComplexExpression;
 
@@ -1699,8 +1701,8 @@ public static Result TryCollectInfo(ref CompilerContext context, Expression expr
                         return r;
 
                     case ExpressionType.Extension:
-                        var reduced = expr.Reduce();
-                        if (reduced == null || ReferenceEquals(reduced, expr))
+                        var reduced = TryReduceNode(expr, ref context);
+                        if (reduced == null)
                             return Result.NotSupported_UnknownExpression;
                         expr = reduced;
                         enterable = false;
@@ -1752,11 +1754,55 @@ public static Result TryCollectInfo(ref CompilerContext context, Expression expr
                             continue;
                         }
 
+#if !LIGHT_EXPRESSION
+                        var customReduced = TryReduceNode(expr, ref context);
+                        if (customReduced != null)
+                        {
+                            expr = customReduced;
+                            enterable = false;
+                            continue;
+                        }
+#endif
                         return Result.NotSupported_UnknownExpression;
                 }
             }
         }
 
+        internal sealed class ExpressionReductionCache
+        {
+            internal SmallMap4<Expression, Expression, RefEq<Expression>> Reductions;
+        }
+
+        // Collection and emission must use the same reduced tree: Reduce may create fresh locals,
+        // constants, labels, or nested lambdas. Allocate the cache only when reducing a node.
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private static Expression TryReduceNode(Expression expression, ref CompilerContext context)
+        {
+            var cache = context.ReducedExpressions;
+            if (cache != null)
+            {
+                ref var cached = ref cache.Reductions.Map.TryGetValueRef(expression, out var found);
+                if (found)
+                    return cached;
+            }
+
+#if !LIGHT_EXPRESSION
+            // System's unknown node kinds use its validation contract. Extension nodes keep
+            // the existing Reduce/null/self contract shared with LightExpression.
+            var reduced = expression.NodeType != ExpressionType.Extension
+                ? expression.CanReduce ? expression.ReduceAndCheck() : null
+                : expression.Reduce();
+#else
+            var reduced = expression.Reduce();
+#endif
+            if (reduced == null || ReferenceEquals(reduced, expression))
+                return null;
+            if (cache == null)
+                context.ReducedExpressions = cache = new ExpressionReductionCache();
+            cache.Reductions.Map.AddOrGetValueRef(expression, out _) = reduced;
+            return reduced;
+        }
+
         private static Expression CreateInlinedLambdaInvocationExpression(ArgExprs invokeArgs, LambdaExpression lambdaExpr)
         {
             var invokeArgCount = invokeArgs.GetArgCount();
@@ -2681,8 +2727,8 @@ public static bool TryEmit(Expression expr, ILGenerator il, ref CompilerContext
                             return TryEmitSwitch((SwitchExpression)expr, il, ref context, parent);
 
                         case ExpressionType.Extension:
-                            var reduced = expr.Reduce();
-                            if (reduced == null || ReferenceEquals(reduced, expr))
+                            var reduced = TryReduceNode(expr, ref context);
+                            if (reduced == null)
                                 return false;
                             expr = reduced;
                             continue;
@@ -2692,6 +2738,14 @@ public static bool TryEmit(Expression expr, ILGenerator il, ref CompilerContext
 
                         case ExpressionType.Quote:     // todo: @feature - is not supported yet
                         default:
+#if !LIGHT_EXPRESSION
+                            var customReduced = TryReduceNode(expr, ref context);
+                            if (customReduced != null)
+                            {
+                                expr = customReduced;
+                                continue;
+                            }
+#endif
                             return false;
 
                     }

test/FastExpressionCompiler.TestsRunner.Net472/Program.cs

@@ -44,6 +44,7 @@ static void RunLightExpressionTests(object state)
                 t.Run(new LightExpression.IssueTests.Issue498_InvalidProgramException_when_using_loop());
                 t.Run(new LightExpression.IssueTests.Issue499_InvalidProgramException_for_Sorting_and_comparison_function());
                 t.Run(new LightExpression.UnitTests.RuntimeTypeConstantRegressionTests());
+                t.Run(new LightExpression.UnitTests.ExtensionReductionIdentityRegressionTests());
                 t.Run(new LightExpression.UnitTests.ParityTests());
 
                 Console.WriteLine($"Just LightExpression tests are passing in {justLightTestsStopwatch.ElapsedMilliseconds} ms.");
@@ -66,6 +67,8 @@ static void RunLightExpressionTests(object state)
             fecTests.Run(new Issue499_InvalidProgramException_for_Sorting_and_comparison_function());
             fecTests.Run(new ReductionProgressRegressionTests());
             fecTests.Run(new RuntimeTypeConstantRegressionTests());
+            fecTests.Run(new CustomReducibleNodeRegressionTests());
+            fecTests.Run(new ExtensionReductionIdentityRegressionTests());
             fecTests.Run(new ParityTests());
 
 

test/FastExpressionCompiler.TestsRunner/Program.cs

@@ -76,6 +76,7 @@ public static void Main(string[] args)
                 lt.Run(new LightExpression.IssueTests.Issue499_InvalidProgramException_for_Sorting_and_comparison_function());
                 lt.Run(new LightExpression.IssueTests.Issue561_NonBooleanEqualityOperator());
                 lt.Run(new LightExpression.UnitTests.RuntimeTypeConstantRegressionTests());
+                lt.Run(new LightExpression.UnitTests.ExtensionReductionIdentityRegressionTests());
                 lt.Run(new LightExpression.UnitTests.ParityTests());
                 ltTime = ltStopwatch.ElapsedMilliseconds;
             });
@@ -85,6 +86,8 @@ public static void Main(string[] args)
             st.Run(new Issue561_NonBooleanEqualityOperator());
             st.Run(new ReductionProgressRegressionTests());
             st.Run(new RuntimeTypeConstantRegressionTests());
+            st.Run(new CustomReducibleNodeRegressionTests());
+            st.Run(new ExtensionReductionIdentityRegressionTests());
             st.Run(new ParityTests());
             st.Run(new Issue498_InvalidProgramException_when_using_loop());
             st.Run(new Issue495_Incomplete_pattern_detection_for_NotSupported_1007_Return_goto_from_TryCatch_with_Assign_generates_invalid_IL());

test/FastExpressionCompiler.UnitTests/CustomReducibleNodeRegressionTests.cs

@@ -0,0 +1,151 @@
+using System;
+#if !LIGHT_EXPRESSION
+using System.Linq.Expressions;
+namespace FastExpressionCompiler.UnitTests;
+public struct CustomReducibleNodeRegressionTests : ITestX
+{
+    public void Run(TestRun t)
+    {
+        Custom_node_type_reduces_before_emission(t);
+        Reduction_collects_captured_constants_and_chains(t);
+        Invalid_reductions_use_system_validation(t);
+        Fresh_block_locals_and_constants_are_shared_between_passes(t);
+        Fresh_nested_lambdas_are_shared_between_passes(t);
+        Shared_nodes_and_distinct_nodes_preserve_reference_identity(t);
+        Reduction_cache_is_local_to_each_compilation(t);
+        Emission_without_collection_can_reduce_custom_nodes(t);
+    }
+    public void Custom_node_type_reduces_before_emission(TestContext t)
+    {
+        var expression = Expression.Lambda<Func<int>>(new Reducible(Expression.Constant(3)));
+        var fast = expression.CompileFast(true);
+        if (t.IsNotNull(fast))
+        {
+#if NETFRAMEWORK
+            t.AreEqual(3, fast()); // Framework System does not reduce unknown node kinds.
+#else
+            t.AreEqual(expression.Compile()(), fast());
+#endif
+        }
+    }
+    public void Reduction_collects_captured_constants_and_chains(TestContext t)
+    {
+        var marker = new object();
+        var expression = Expression.Lambda<Func<object>>(new Reducible(new Reducible(Expression.Constant(marker))));
+        var fast = expression.CompileFast(true);
+        if (t.IsNotNull(fast)) t.AreEqual(true, ReferenceEquals(marker, fast()));
+    }
+    public void Invalid_reductions_use_system_validation(TestContext t)
+    {
+        var source = Expression.Lambda<Func<int>>(new InvalidReduction());
+        t.Throws<ArgumentException>(() => source.CompileFast(true));
+    }
+    private static readonly CompilerFlags[] Flags =
+    {
+        CompilerFlags.Default, CompilerFlags.NoInvocationLambdaInlining,
+        CompilerFlags.EnableDelegateDebugInfo | CompilerFlags.ThrowOnNotSupportedExpression
+    };
+    public void Fresh_block_locals_and_constants_are_shared_between_passes(TestContext t)
+    {
+        foreach (var flags in Flags)
+        {
+            var node = new FreshBlock(41);
+            var fast = Expression.Lambda<Func<int>>(node).CompileFast(true, flags);
+            if (t.IsNotNull(fast)) t.AreEqual(42, fast());
+            t.AreEqual(1, node.Reductions);
+        }
+    }
+    public void Fresh_nested_lambdas_are_shared_between_passes(TestContext t)
+    {
+        foreach (var flags in Flags)
+        {
+            var node = new FreshNestedLambda();
+            var fast = Expression.Lambda<Func<Func<int>>>(node).CompileFast(true, flags);
+            if (t.IsNotNull(fast)) t.AreEqual(42, fast()());
+            t.AreEqual(1, node.Reductions);
+        }
+    }
+    public void Shared_nodes_and_distinct_nodes_preserve_reference_identity(TestContext t)
+    {
+        foreach (var flags in Flags)
+        {
+            var shared = new FreshBlock(20);
+            var same = Expression.Lambda<Func<int>>(Expression.Add(shared, shared)).CompileFast(true, flags);
+            if (t.IsNotNull(same)) t.AreEqual(42, same());
+            t.AreEqual(1, shared.Reductions);
+            // User equality must not make two different custom nodes share a reduction.
+            var left = new FreshBlock(20); var right = new FreshBlock(21);
+            var distinct = Expression.Lambda<Func<int>>(Expression.Add(left, right)).CompileFast(true, flags);
+            if (t.IsNotNull(distinct)) t.AreEqual(43, distinct());
+            t.AreEqual(1, left.Reductions); t.AreEqual(1, right.Reductions);
+        }
+    }
+    public void Reduction_cache_is_local_to_each_compilation(TestContext t)
+    {
+        var node = new FreshBlock(41);
+        var source = Expression.Lambda<Func<int>>(node);
+        for (var i = 1; i <= 2; ++i)
+        {
+            var fast = source.CompileFast(true);
+            if (t.IsNotNull(fast)) t.AreEqual(42, fast());
+            t.AreEqual(i, node.Reductions);
+        }
+    }
+    public void Emission_without_collection_can_reduce_custom_nodes(TestContext t)
+    {
+        var source = Expression.Lambda<Func<int>>(new Reducible(Expression.Constant(42)));
+        var fast = source.TryCompileWithoutClosure<Func<int>>();
+        if (t.IsNotNull(fast)) t.AreEqual(42, fast());
+    }
+    public sealed class Marker { public int Value; }
+    private sealed class FreshBlock : Expression
+    {
+        private readonly int _value;
+        public int Reductions;
+        public FreshBlock(int value) => _value = value;
+        public override Type Type => typeof(int);
+        public override ExpressionType NodeType => (ExpressionType)(-1);
+        public override bool CanReduce => true;
+        public override Expression Reduce()
+        {
+            ++Reductions;
+            var local = Expression.Variable(typeof(Marker), "fresh");
+            return Expression.Block(new[] { local },
+                Expression.Assign(local, Expression.Constant(new Marker { Value = _value })),
+                Expression.Add(Expression.Field(local, nameof(Marker.Value)), Expression.Constant(1)));
+        }
+        public override bool Equals(object obj) => obj is FreshBlock;
+        public override int GetHashCode() => 0;
+    }
+    private sealed class FreshNestedLambda : Expression
+    {
+        public int Reductions;
+        public override Type Type => typeof(Func<int>);
+        public override ExpressionType NodeType => (ExpressionType)(-1);
+        public override bool CanReduce => true;
+        public override Expression Reduce()
+        {
+            ++Reductions;
+            var local = Expression.Variable(typeof(int), "fresh");
+            return Expression.Block(new[] { local }, Expression.Assign(local, Expression.Constant(41)),
+                Expression.Lambda<Func<int>>(Expression.Add(local, Expression.Constant(1))));
+        }
+    }
+    private sealed class InvalidReduction : Expression
+    {
+        public override Type Type => typeof(int);
+        public override ExpressionType NodeType => (ExpressionType)(-1);
+        public override bool CanReduce => true;
+        public override Expression Reduce() => Expression.Constant("wrong type");
+    }
+    public class Reducible : Expression
+    {
+        private readonly Expression _reduced;
+        public Reducible(Expression reduced) => _reduced = reduced;
+        public override Type Type => _reduced.Type;
+        public override ExpressionType NodeType => (ExpressionType)(-1);
+        public override bool CanReduce => true;
+        public override Expression Reduce() => _reduced;
+    }
+}
+#endif