MycLib/Src/AST/Myc.Ast.Analyzer.pas

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.