Refine FE owner metadata lookups

Agent-Logs-Url: https://github.com/dadhi/FastExpressionCompiler/sessions/be12d6ed-b5f1-46be-828a-a7673aba6a54

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

Author: Copilot

src/FastExpressionCompiler.LightExpression/FlatExpression.cs

@@ -217,32 +217,23 @@ public struct ExprTree
     /// <summary>Gets or sets owner-side metadata for block nodes, populated during FE construction.</summary>
     public SmallList<BlockInfo, Stack16<BlockInfo>, NoArrayPool<BlockInfo>> Blocks;
 
+    private SmallMap16<int, int, IntEq> _lambdaInfoIndexesByNodeIndex;
+    private SmallMap16<int, int, IntEq> _blockInfoIndexesByNodeIndex;
+
     /// <summary>Returns the owner-side metadata for the specified lambda node, if present.</summary>
     public bool TryGetLambdaInfo(int lambdaIndex, out LambdaInfo info)
     {
-        for (var i = 0; i < Lambdas.Count; ++i)
-        {
-            info = Lambdas.GetSurePresentRef(i);
-            if (info.LambdaIndex == lambdaIndex)
-                return true;
-        }
-
-        info = default;
-        return false;
+        ref var infoIndex = ref _lambdaInfoIndexesByNodeIndex.Map.TryGetValueRef(lambdaIndex, out var found);
+        info = found ? Lambdas.GetSurePresentRef(infoIndex) : default;
+        return found;
     }
 
     /// <summary>Returns the owner-side metadata for the specified block node, if present.</summary>
     public bool TryGetBlockInfo(int blockIndex, out BlockInfo info)
     {
-        for (var i = 0; i < Blocks.Count; ++i)
-        {
-            info = Blocks.GetSurePresentRef(i);
-            if (info.BlockIndex == blockIndex)
-                return true;
-        }
-
-        info = default;
-        return false;
+        ref var infoIndex = ref _blockInfoIndexesByNodeIndex.Map.TryGetValueRef(blockIndex, out var found);
+        info = found ? Blocks.GetSurePresentRef(infoIndex) : default;
+        return found;
     }
 
     /// <summary>Adds a parameter node and returns its index.</summary>
@@ -1368,12 +1359,32 @@ private void BindParameterDeclaration(int ownerIndex, int declarationIndex, int
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private void AddLambdaInfo(int lambdaIndex, int bodyIndex, in ChildList declarations) =>
-        Lambdas.Add(new LambdaInfo(lambdaIndex, bodyIndex, declarations));
+    private void AddLambdaInfo(int lambdaIndex, int bodyIndex, in ChildList declarations)
+    {
+        ref var infoIndex = ref _lambdaInfoIndexesByNodeIndex.Map.AddOrGetValueRef(lambdaIndex, out var found);
+        var info = new LambdaInfo(lambdaIndex, bodyIndex, declarations);
+        if (found)
+            Lambdas.GetSurePresentRef(infoIndex) = info;
+        else
+        {
+            infoIndex = Lambdas.Count;
+            Lambdas.Add(info);
+        }
+    }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private void AddBlockInfo(int blockIndex, int bodyExpressionsIndex, in ChildList declarations) =>
-        Blocks.Add(new BlockInfo(blockIndex, bodyExpressionsIndex, declarations));
+    private void AddBlockInfo(int blockIndex, int bodyExpressionsIndex, in ChildList declarations)
+    {
+        ref var infoIndex = ref _blockInfoIndexesByNodeIndex.Map.AddOrGetValueRef(blockIndex, out var found);
+        var info = new BlockInfo(blockIndex, bodyExpressionsIndex, declarations);
+        if (found)
+            Blocks.GetSurePresentRef(infoIndex) = info;
+        else
+        {
+            infoIndex = Blocks.Count;
+            Blocks.Add(info);
+        }
+    }
 
     private int AddParameterUsageNode(in ExprNode node, int declarationIndex) =>
         AddRawLeafExpressionNode(node.Type, node.Obj, ExpressionType.Parameter,

test/FastExpressionCompiler.LightExpression.UnitTests/LightExpressionTests.cs

@@ -31,11 +31,12 @@ public int Run()
             Can_roundtrip_light_expression_through_flat_expression();
             Flat_expression_preserves_parameter_and_label_identity_and_collects_closure_constants();
             Flat_expression_uses_parameter_declarations_parent_links_and_inline_constants();
+            Flat_expression_collects_lambda_and_block_owner_metadata_when_flattening_system_expression();
             Can_convert_dynamic_runtime_variables_and_debug_info_to_light_expression_and_flat_expression();
             Can_build_flat_expression_directly_with_light_expression_like_api();
             Can_build_flat_expression_control_flow_directly();
             Can_property_test_generated_flat_expression_roundtrip_structurally();
-            return 18;
+            return 19;
         }