FlatExpression: also track BlocksWithVariables, GotoNodes, LabelNodes and TryCatchNodes during construction

Agent-Logs-Url: https://github.com/dadhi/FastExpressionCompiler/sessions/fbbd4fe6-1e27-488e-b48b-b5ff5dbc8626

Co-authored-by: dadhi <39516+dadhi@users.noreply.github.com>

Author: Copilot

src/FastExpressionCompiler.LightExpression/FlatExpression.cs

@@ -173,6 +173,30 @@ public struct ExprTree
     /// enabling callers to discover nested lambdas without a full tree traversal.</summary>
     public SmallList<int, Stack16<int>, NoArrayPool<int>> LambdaNodes;
 
+    /// <summary>Gets or sets the indices of all block nodes that carry explicit variable declarations.
+    /// These are the block nodes where <c>children.Count == 2</c> (variable list + expression list).
+    /// Populated automatically by <see cref="Block(Type,IEnumerable{int},int[])"/> and <see cref="ExprTree.FromExpression"/>,
+    /// enabling callers to enumerate block-scoped variables without a full tree traversal.</summary>
+    public SmallList<int, Stack16<int>, NoArrayPool<int>> BlocksWithVariables;
+
+    /// <summary>Gets or sets the indices of all <see cref="ExpressionType.Goto"/> nodes
+    /// (including <c>return</c> and <c>break</c>/<c>continue</c> goto-family nodes).
+    /// Populated automatically by <see cref="MakeGoto"/> and <see cref="ExprTree.FromExpression"/>,
+    /// enabling callers to link gotos to their label targets without a full tree traversal.</summary>
+    public SmallList<int, Stack16<int>, NoArrayPool<int>> GotoNodes;
+
+    /// <summary>Gets or sets the indices of all <see cref="ExpressionType.Label"/> expression nodes.
+    /// Populated automatically by <see cref="Label(int,int?)"/> and <see cref="ExprTree.FromExpression"/>,
+    /// enabling callers to link label expressions to their targets without a full tree traversal.</summary>
+    public SmallList<int, Stack16<int>, NoArrayPool<int>> LabelNodes;
+
+    /// <summary>Gets or sets the indices of all <see cref="ExpressionType.Try"/> nodes
+    /// (try/catch, try/finally, try/fault, and combined forms).
+    /// Populated automatically by <see cref="TryCatch"/>, <see cref="TryFinally"/>,
+    /// <see cref="TryFault"/>, <see cref="TryCatchFinally"/> and <see cref="ExprTree.FromExpression"/>,
+    /// enabling callers to locate all try regions without a full tree traversal.</summary>
+    public SmallList<int, Stack16<int>, NoArrayPool<int>> TryCatchNodes;
+
     /// <summary>Adds a parameter node and returns its index.</summary>
     public int Parameter(Type type, string name = null)
     {
@@ -380,26 +404,36 @@ public int Block(params int[] expressions) =>
     /// Variable parameter nodes share the same id-slot as the refs used inside the body
     /// (out-of-order: the variable decl nodes appear in children[0] before the body expressions
     /// that reference them in children[1]).
+    /// <para>When the block has explicit variable declarations its node index is recorded in
+    /// <see cref="BlocksWithVariables"/>, enabling callers to enumerate block-scoped variables
+    /// without a full tree traversal.</para>
     /// </remarks>
     public int Block(Type type, IEnumerable<int> variables, params int[] expressions)
     {
         if (expressions == null || expressions.Length == 0)
             throw new ArgumentException("Block should contain at least one expression.", nameof(expressions));
 
         ChildList children = default;
+        var hasVariables = false;
         if (variables != null)
         {
             ChildList variableChildren = default;
             foreach (var variable in variables)
                 variableChildren.Add(variable);
             if (variableChildren.Count != 0)
+            {
                 children.Add(AddChildListNode(in variableChildren));
+                hasVariables = true;
+            }
         }
         ChildList bodyChildren = default;
         for (var i = 0; i < expressions.Length; ++i)
             bodyChildren.Add(expressions[i]);
         children.Add(AddChildListNode(in bodyChildren));
-        return AddFactoryExpressionNode(type ?? Nodes[expressions[expressions.Length - 1]].Type, null, ExpressionType.Block, in children);
+        var index = AddFactoryExpressionNode(type ?? Nodes[expressions[expressions.Length - 1]].Type, null, ExpressionType.Block, in children);
+        if (hasVariables)
+            BlocksWithVariables.Add(index);
+        return index;
     }
 
     /// <summary>Adds a typed lambda node.</summary>
@@ -468,18 +502,26 @@ public int Label(Type type = null, string name = null)
     }
 
     /// <summary>Adds a label-expression node.</summary>
-    public int Label(int target, int? defaultValue = null) =>
-        defaultValue.HasValue
+    /// <remarks>The node index is recorded in <see cref="LabelNodes"/>.</remarks>
+    public int Label(int target, int? defaultValue = null)
+    {
+        var index = defaultValue.HasValue
             ? AddFactoryExpressionNode(Nodes[target].Type, null, ExpressionType.Label, 0, target, defaultValue.Value)
             : AddFactoryExpressionNode(Nodes[target].Type, null, ExpressionType.Label, 0, target);
+        LabelNodes.Add(index);
+        return index;
+    }
 
     /// <summary>Adds a goto-family node.</summary>
+    /// <remarks>The node index is recorded in <see cref="GotoNodes"/>.</remarks>
     public int MakeGoto(GotoExpressionKind kind, int target, int? value = null, Type type = null)
     {
         var resultType = type ?? (value.HasValue ? Nodes[value.Value].Type : typeof(void));
-        return value.HasValue
+        var index = value.HasValue
             ? AddFactoryExpressionNode(resultType, kind, ExpressionType.Goto, 0, target, value.Value)
             : AddFactoryExpressionNode(resultType, kind, ExpressionType.Goto, 0, target);
+        GotoNodes.Add(index);
+        return index;
     }
 
     /// <summary>Adds a goto node.</summary>
@@ -557,29 +599,48 @@ public int MakeCatchBlock(Type test, int? variable, int body, int? filter)
     }
 
     /// <summary>Adds a try/catch node.</summary>
+    /// <remarks>The node index is recorded in <see cref="TryCatchNodes"/>.</remarks>
     public int TryCatch(int body, params int[] handlers)
     {
+        int index;
         if (handlers == null || handlers.Length == 0)
-            return AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, 0, body);
-
-        ChildList handlerChildren = default;
-        for (var i = 0; i < handlers.Length; ++i)
-            handlerChildren.Add(handlers[i]);
-        ChildList children = default;
-        children.Add(body);
-        children.Add(AddChildListNode(in handlerChildren));
-        return AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, in children);
+        {
+            index = AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, 0, body);
+        }
+        else
+        {
+            ChildList handlerChildren = default;
+            for (var i = 0; i < handlers.Length; ++i)
+                handlerChildren.Add(handlers[i]);
+            ChildList children = default;
+            children.Add(body);
+            children.Add(AddChildListNode(in handlerChildren));
+            index = AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, in children);
+        }
+        TryCatchNodes.Add(index);
+        return index;
     }
 
     /// <summary>Adds a try/finally node.</summary>
-    public int TryFinally(int body, int @finally) =>
-        AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, 0, body, @finally);
+    /// <remarks>The node index is recorded in <see cref="TryCatchNodes"/>.</remarks>
+    public int TryFinally(int body, int @finally)
+    {
+        var index = AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, 0, body, @finally);
+        TryCatchNodes.Add(index);
+        return index;
+    }
 
     /// <summary>Adds a try/fault node.</summary>
-    public int TryFault(int body, int fault) =>
-        AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, TryFaultFlag, body, fault);
+    /// <remarks>The node index is recorded in <see cref="TryCatchNodes"/>.</remarks>
+    public int TryFault(int body, int fault)
+    {
+        var index = AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, TryFaultFlag, body, fault);
+        TryCatchNodes.Add(index);
+        return index;
+    }
 
     /// <summary>Adds a try node with optional finally block and catch handlers.</summary>
+    /// <remarks>The node index is recorded in <see cref="TryCatchNodes"/>.</remarks>
     public int TryCatchFinally(int body, int? @finally, params int[] handlers)
     {
         ChildList children = default;
@@ -593,7 +654,9 @@ public int TryCatchFinally(int body, int? @finally, params int[] handlers)
                 handlerChildren.Add(handlers[i]);
             children.Add(AddChildListNode(in handlerChildren));
         }
-        return AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, 0, in children);
+        var index = AddFactoryExpressionNode(Nodes[body].Type, null, ExpressionType.Try, 0, in children);
+        TryCatchNodes.Add(index);
+        return index;
     }
 
     /// <summary>Adds a type-test node.</summary>
@@ -847,7 +910,8 @@ private int AddExpression(SysExpr expression)
                         // children.Count == 2 is the canonical test for the presence of variables.
                         var block = (System.Linq.Expressions.BlockExpression)expression;
                         ChildList children = default;
-                        if (block.Variables.Count != 0)
+                        var hasVariables = block.Variables.Count != 0;
+                        if (hasVariables)
                         {
                             ChildList variables = default;
                             for (var i = 0; i < block.Variables.Count; ++i)
@@ -858,7 +922,10 @@ private int AddExpression(SysExpr expression)
                         for (var i = 0; i < block.Expressions.Count; ++i)
                             expressions.Add(AddExpression(block.Expressions[i]));
                         children.Add(_tree.AddChildListNode(in expressions));
-                        return _tree.AddRawExpressionNode(expression.Type, null, expression.NodeType, in children);
+                        var blockIndex = _tree.AddRawExpressionNode(expression.Type, null, expression.NodeType, in children);
+                        if (hasVariables)
+                            _tree.BlocksWithVariables.Add(blockIndex);
+                        return blockIndex;
                     }
                 case ExpressionType.MemberAccess:
                     {
@@ -944,7 +1011,9 @@ private int AddExpression(SysExpr expression)
                         children.Add(AddLabelTarget(@goto.Target));
                         if (@goto.Value != null)
                             children.Add(AddExpression(@goto.Value));
-                        return _tree.AddRawExpressionNode(expression.Type, @goto.Kind, expression.NodeType, children);
+                        var gotoIndex = _tree.AddRawExpressionNode(expression.Type, @goto.Kind, expression.NodeType, children);
+                        _tree.GotoNodes.Add(gotoIndex);
+                        return gotoIndex;
                     }
                 case ExpressionType.Label:
                     {
@@ -953,7 +1022,9 @@ private int AddExpression(SysExpr expression)
                         children.Add(AddLabelTarget(label.Target));
                         if (label.DefaultValue != null)
                             children.Add(AddExpression(label.DefaultValue));
-                        return _tree.AddRawExpressionNode(expression.Type, null, expression.NodeType, children);
+                        var labelIndex = _tree.AddRawExpressionNode(expression.Type, null, expression.NodeType, children);
+                        _tree.LabelNodes.Add(labelIndex);
+                        return labelIndex;
                     }
                 case ExpressionType.Switch:
                     {
@@ -991,7 +1062,9 @@ private int AddExpression(SysExpr expression)
                                 handlers.Add(AddCatchBlock(@try.Handlers[i]));
                             children.Add(_tree.AddChildListNode(in handlers));
                         }
-                        return _tree.AddNode(expression.Type, null, expression.NodeType, ExprNodeKind.Expression, flags, in children);
+                        var tryIndex = _tree.AddNode(expression.Type, null, expression.NodeType, ExprNodeKind.Expression, flags, in children);
+                        _tree.TryCatchNodes.Add(tryIndex);
+                        return tryIndex;
                     }
                 case ExpressionType.MemberInit:
                     {

test/FastExpressionCompiler.LightExpression.UnitTests/LightExpressionTests.cs

@@ -44,7 +44,13 @@ public int Run()
             Flat_lambda_nodes_tracks_deeply_nested_lambdas_during_direct_construction();
             Flat_lambda_nodes_tracks_lambdas_from_expression_conversion();
             Flat_lambda_nodes_has_single_entry_for_root_only_lambda();
-            return 27;
+            Flat_blocks_with_variables_tracked_during_direct_construction();
+            Flat_goto_and_label_nodes_tracked_during_direct_construction();
+            Flat_try_catch_nodes_tracked_during_direct_construction();
+            Flat_blocks_with_variables_tracked_from_expression_conversion();
+            Flat_goto_and_label_nodes_tracked_from_expression_conversion();
+            Flat_try_catch_nodes_tracked_from_expression_conversion();
+            return 33;
         }
 
 
@@ -796,5 +802,134 @@ public void Flat_lambda_nodes_has_single_entry_for_root_only_lambda()
             Asserts.AreEqual(1, fe.LambdaNodes.Count);
             Asserts.AreEqual(fe.RootIndex, fe.LambdaNodes[0]);
         }
+
+        /// <summary>
+        /// Block nodes with explicit variable declarations are recorded in BlocksWithVariables;
+        /// blocks without variables produce no entry.
+        /// </summary>
+        public void Flat_blocks_with_variables_tracked_during_direct_construction()
+        {
+            var fe = default(ExprTree);
+            var p = fe.ParameterOf<int>("p");
+            var v = fe.Variable(typeof(int), "v");
+
+            // Block with one variable: should be tracked
+            var blockWithVar = fe.Block(typeof(int), new[] { v }, fe.Assign(v, p), v);
+            // Block without variables: should NOT be tracked
+            var blockNoVar = fe.Block(fe.Add(p, fe.ConstantInt(1)));
+
+            fe.RootIndex = fe.Lambda<Func<int, int>>(fe.Block(blockWithVar, blockNoVar), p);
+
+            Asserts.AreEqual(1, fe.BlocksWithVariables.Count);
+            Asserts.AreEqual(blockWithVar, fe.BlocksWithVariables[0]);
+        }
+
+        /// <summary>
+        /// Goto and label expression nodes are recorded in GotoNodes and LabelNodes respectively.
+        /// </summary>
+        public void Flat_goto_and_label_nodes_tracked_during_direct_construction()
+        {
+            var fe = default(ExprTree);
+            var p = fe.ParameterOf<int>("p");
+            var target = fe.Label(typeof(int), "done");
+
+            var gotoNode = fe.Goto(target, p, typeof(int));
+            var labelNode = fe.Label(target, fe.ConstantInt(0));
+
+            fe.RootIndex = fe.Lambda<Func<int, int>>(fe.Block(gotoNode, labelNode), p);
+
+            Asserts.AreEqual(1, fe.GotoNodes.Count);
+            Asserts.AreEqual(gotoNode, fe.GotoNodes[0]);
+
+            Asserts.AreEqual(1, fe.LabelNodes.Count);
+            Asserts.AreEqual(labelNode, fe.LabelNodes[0]);
+        }
+
+        /// <summary>
+        /// Try/catch, try/finally and try/fault node indices are all recorded in TryCatchNodes.
+        /// </summary>
+        public void Flat_try_catch_nodes_tracked_during_direct_construction()
+        {
+            var fe = default(ExprTree);
+            var p = fe.ParameterOf<int>("p");
+
+            var tryCatchNode = fe.TryCatch(
+                fe.Add(p, fe.ConstantInt(1)),
+                fe.Catch(typeof(Exception), fe.ConstantInt(-1)));
+
+            var tryFinallyNode = fe.TryFinally(
+                fe.Add(p, fe.ConstantInt(2)),
+                fe.Default(typeof(void)));
+
+            fe.RootIndex = fe.Lambda<Func<int, int>>(
+                fe.Block(tryCatchNode, tryFinallyNode), p);
+
+            Asserts.AreEqual(2, fe.TryCatchNodes.Count);
+
+            var foundTryCatch = false;
+            var foundTryFinally = false;
+            for (var i = 0; i < fe.TryCatchNodes.Count; i++)
+            {
+                if (fe.TryCatchNodes[i] == tryCatchNode) foundTryCatch = true;
+                if (fe.TryCatchNodes[i] == tryFinallyNode) foundTryFinally = true;
+            }
+            Asserts.IsTrue(foundTryCatch);
+            Asserts.IsTrue(foundTryFinally);
+        }
+
+        /// <summary>
+        /// When converting a System.Linq expression tree, blocks with variables are
+        /// recorded in BlocksWithVariables; plain blocks are not.
+        /// </summary>
+        public void Flat_blocks_with_variables_tracked_from_expression_conversion()
+        {
+            var p = SysExpr.Parameter(typeof(int), "p");
+            var v = SysExpr.Variable(typeof(int), "v");
+            // block with variable
+            var sysBlock = SysExpr.Block(new[] { v }, SysExpr.Assign(v, p), v);
+            var sysLambda = SysExpr.Lambda<Func<int, int>>(sysBlock, p);
+
+            var fe = sysLambda.ToFlatExpression();
+
+            Asserts.AreEqual(1, fe.BlocksWithVariables.Count);
+        }
+
+        /// <summary>
+        /// When converting a System.Linq expression tree with goto/label, both
+        /// GotoNodes and LabelNodes are populated.
+        /// </summary>
+        public void Flat_goto_and_label_nodes_tracked_from_expression_conversion()
+        {
+            var p = SysExpr.Parameter(typeof(int), "p");
+            var target = SysExpr.Label(typeof(int), "done");
+            var sysLambda = SysExpr.Lambda<Func<int, int>>(
+                SysExpr.Block(
+                    SysExpr.Goto(target, p, typeof(int)),
+                    SysExpr.Label(target, SysExpr.Constant(0))),
+                p);
+
+            var fe = sysLambda.ToFlatExpression();
+
+            Asserts.AreEqual(1, fe.GotoNodes.Count);
+            Asserts.AreEqual(1, fe.LabelNodes.Count);
+        }
+
+        /// <summary>
+        /// When converting a System.Linq expression tree with a try/catch,
+        /// TryCatchNodes is populated.
+        /// </summary>
+        public void Flat_try_catch_nodes_tracked_from_expression_conversion()
+        {
+            var p = SysExpr.Parameter(typeof(int), "p");
+            var sysLambda = SysExpr.Lambda<Func<int, int>>(
+                SysExpr.TryCatch(
+                    SysExpr.Add(p, SysExpr.Constant(1)),
+                    SysExpr.Catch(typeof(Exception), SysExpr.Constant(-1))),
+                p);
+
+            var fe = sysLambda.ToFlatExpression();
+
+            Asserts.AreEqual(1, fe.TryCatchNodes.Count);
+        }
     }
 }