Fixed closure upvalue scoping

Src/AST/Myc.Ast.Analyzer.pas

@@ -0,0 +1,155 @@
+unit Myc.Ast.Analyzer;
+
+interface
+
+uses
+    System.SysUtils,
+    System.Classes,
+    System.Generics.Collections,
+    Myc.Ast.Nodes,
+    Myc.Ast.Visitor,
+    Myc.Ast.Scope;
+
+type
+    // This visitor analyzes the AST to find all variables that need to be "lifted" or "boxed"
+    // because they are captured by a nested lambda.
+    TUpvalueAnalyzer = class(TAstTraverser)
+    private
+        FBoxedDeclarations: THashSet<IVariableDeclarationNode>;
+        FCurrentScope: IScopeDescriptor;
+        FDeclarationMap: TDictionary<IScopeDescriptor, TDictionary<string, IVariableDeclarationNode>>;
+        procedure MarkDeclarationForBoxing(const AName: string);
+    protected
+        function TransformLambdaExpression(const Node: ILambdaExpressionNode): ILambdaExpressionNode; override;
+        function TransformIdentifier(const Node: IIdentifierNode): IIdentifierNode; override;
+        function TransformVariableDeclaration(const Node: IVariableDeclarationNode): IVariableDeclarationNode; override;
+    public
+        constructor Create(const AParent: IScopeDescriptor);
+        destructor Destroy; override;
+        class function Analyze(const ARootNode: IAstNode; const AParent: IScopeDescriptor): THashSet<IVariableDeclarationNode>; static;
+    end;
+
+implementation
+
+{ TUpvalueAnalyzer }
+
+constructor TUpvalueAnalyzer.Create(const AParent: IScopeDescriptor);
+begin
+    inherited Create;
+    FBoxedDeclarations := THashSet<IVariableDeclarationNode>.Create;
+    FDeclarationMap := TDictionary<IScopeDescriptor, TDictionary<string, IVariableDeclarationNode>>.Create;
+    FCurrentScope := TScope.CreateDescriptor(AParent);
+end;
+
+destructor TUpvalueAnalyzer.Destroy;
+begin
+    FBoxedDeclarations.Free;
+    for var dict in FDeclarationMap.Values do
+        dict.Free;
+    FDeclarationMap.Free;
+    inherited Destroy;
+end;
+
+class function TUpvalueAnalyzer.Analyze(const ARootNode: IAstNode; const AParent: IScopeDescriptor): THashSet<IVariableDeclarationNode>;
+var
+    analyzer: TUpvalueAnalyzer;
+begin
+    if not Assigned(ARootNode) then
+        exit(THashSet<IVariableDeclarationNode>.Create);
+
+    analyzer := TUpvalueAnalyzer.Create(AParent);
+    try
+        analyzer.Execute(ARootNode);
+        Result := analyzer.FBoxedDeclarations;
+        analyzer.FBoxedDeclarations := nil; // Transfer ownership
+    finally
+        analyzer.Free;
+    end;
+end;
+
+procedure TUpvalueAnalyzer.MarkDeclarationForBoxing(const AName: string);
+var
+    address: TResolvedAddress;
+    declarationScope: IScopeDescriptor;
+    i: Integer;
+begin
+    address := FCurrentScope.FindSymbol(AName);
+    if address.Kind <> akLocalOrParent then
+        exit;
+
+    // Walk up the scope chain to find the scope where the variable was declared.
+    declarationScope := FCurrentScope;
+    for i := 1 to address.ScopeDepth do
+    begin
+        if not Assigned(declarationScope.Parent) then
+            exit; // Should not happen in a correctly bound tree
+        declarationScope := declarationScope.Parent;
+    end;
+
+    // Find the declaration node in our map and add it to the set.
+    var dict: TDictionary<string, IVariableDeclarationNode>;
+    if FDeclarationMap.TryGetValue(declarationScope, dict) then
+    begin
+        var declNode: IVariableDeclarationNode;
+        if dict.TryGetValue(AName, declNode) then
+            FBoxedDeclarations.Add(declNode);
+    end;
+end;
+
+function TUpvalueAnalyzer.TransformIdentifier(const Node: IIdentifierNode): IIdentifierNode;
+var
+    address: TResolvedAddress;
+begin
+    Result := Node;
+    if not Assigned(FCurrentScope) then
+        exit;
+
+    address := FCurrentScope.FindSymbol(Node.Name);
+
+    if (address.Kind = akLocalOrParent) and (address.ScopeDepth > 0) then
+    begin
+        // This is an upvalue. Mark its original declaration for boxing.
+        MarkDeclarationForBoxing(Node.Name);
+    end;
+end;
+
+function TUpvalueAnalyzer.TransformLambdaExpression(const Node: ILambdaExpressionNode): ILambdaExpressionNode;
+begin
+    // A lambda creates a new lexical scope, inheriting from the current one.
+    FCurrentScope := TScope.CreateDescriptor(FCurrentScope);
+    try
+        // Define the lambda's parameters within its new scope.
+        for var param in Node.Parameters do
+            FCurrentScope.Define(param.Name);
+
+        // Traverse the lambda body within the new scope context.
+        Node.Body.Accept(Self);
+        Result := Node; // We do not transform, just analyze.
+    finally
+        // Restore the parent scope after leaving the lambda.
+        FCurrentScope := FCurrentScope.Parent;
+    end;
+end;
+
+function TUpvalueAnalyzer.TransformVariableDeclaration(const Node: IVariableDeclarationNode): IVariableDeclarationNode;
+var
+    scopeDeclarations: TDictionary<string, IVariableDeclarationNode>;
+begin
+    Result := Node;
+
+    if Assigned(Node.Initializer) then
+        Node.Initializer.Accept(Self);
+
+    // After processing the initializer, define the variable in the current scope.
+    FCurrentScope.Define(Node.Identifier.Name);
+
+    // Map this declaration node to its scope and name for later lookup.
+    if not FDeclarationMap.TryGetValue(FCurrentScope, scopeDeclarations) then
+    begin
+        scopeDeclarations := TDictionary<string, IVariableDeclarationNode>.Create;
+        FDeclarationMap.Add(FCurrentScope, scopeDeclarations);
+    end;
+    scopeDeclarations.Add(Node.Identifier.Name, Node);
+end;
+
+end.