unit Myc.Ast.Compiler.Binder.Upvalues;

interface

uses
    System.SysUtils,
    System.Classes,
    System.Generics.Collections,
    Myc.Ast.Nodes,
    Myc.Ast.Visitor,
    Myc.Ast.Scope,
    Myc.Data.Value,
    Myc.Ast;

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(TAstTransformer)
    private
        type
            // Internal helper to track scopes during analysis without using IScopeDescriptor
            TAnalysisScope = class
            private
                FParent: TAnalysisScope;
                // Maps Name -> DeclarationNode.
                // If Value is nil, it means the name is defined (e.g. parameter) but not a candidate for boxing.
                FDeclarations: TDictionary<string, IVariableDeclarationNode>;
            public
                constructor Create(AParent: TAnalysisScope);
                destructor Destroy; override;

                procedure Define(const Name: string; const Node: IVariableDeclarationNode);

                // Returns the scope where the symbol is defined and the declaration node (if any)
                function Resolve(const Name: string; out ScopeDepth: Integer; out Node: IVariableDeclarationNode): Boolean;
            end;

    private
        FBoxedDeclarations: THashSet<IVariableDeclarationNode>;
        FCurrentScope: TAnalysisScope;

        procedure MarkDeclarationForBoxing(const AName: string);

    strict private
        // Analysis Handlers (IAstNode signature)
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;
        function VisitIdentifier(const Node: IAstNode): IAstNode;
        function VisitVariableDeclaration(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;

    public
        constructor Create;
        destructor Destroy; override;

        function Execute(const ARootNode: IAstNode): IAstNode;

        class function Analyze(const ARootNode: IAstNode): THashSet<IVariableDeclarationNode>; static;
    end;

implementation

uses
    System.Generics.Defaults,
    Myc.Ast.Types;

{ TUpvalueAnalyzer.TAnalysisScope }

constructor TUpvalueAnalyzer.TAnalysisScope.Create(AParent: TAnalysisScope);
begin
    inherited Create;
    FParent := AParent;
    FDeclarations := TDictionary<string, IVariableDeclarationNode>.Create;
end;

destructor TUpvalueAnalyzer.TAnalysisScope.Destroy;
begin
    FDeclarations.Free;
    inherited;
end;

procedure TUpvalueAnalyzer.TAnalysisScope.Define(const Name: string; const Node: IVariableDeclarationNode);
begin
    FDeclarations.AddOrSetValue(Name, Node);
end;

function TUpvalueAnalyzer.TAnalysisScope.Resolve(const Name: string; out ScopeDepth: Integer; out Node: IVariableDeclarationNode): Boolean;
var
    current: TAnalysisScope;
begin
    ScopeDepth := 0;
    current := Self;
    while Assigned(current) do
    begin
        if current.FDeclarations.TryGetValue(Name, Node) then
        begin
            Result := True;
            exit;
        end;
        Inc(ScopeDepth);
        current := current.FParent;
    end;

    Node := nil;
    Result := False;
end;

{ TUpvalueAnalyzer }

constructor TUpvalueAnalyzer.Create;
begin
    inherited Create;
    FBoxedDeclarations := THashSet<IVariableDeclarationNode>.Create;
    // Create a root scope to handle top-level definitions cleanly
    FCurrentScope := TAnalysisScope.Create(nil);
end;

destructor TUpvalueAnalyzer.Destroy;
begin
    // Unwind scope stack if exception occurred or Execute wasn't fully clean
    while Assigned(FCurrentScope.FParent) do
    begin
        var temp := FCurrentScope;
        FCurrentScope := FCurrentScope.FParent;
        temp.Free;
    end;
    FCurrentScope.Free;

    FBoxedDeclarations.Free;
    inherited Destroy;
end;

procedure TUpvalueAnalyzer.SetupHandlers;
begin
    inherited SetupHandlers; // Load default transformer logic

    // Override specific handlers for analysis
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akIdentifier, VisitIdentifier);
    Register(akVariableDeclaration, VisitVariableDeclaration);
end;

function TUpvalueAnalyzer.Execute(const ARootNode: IAstNode): IAstNode;
begin
    Result := Accept(ARootNode);
end;

class function TUpvalueAnalyzer.Analyze(const ARootNode: IAstNode): THashSet<IVariableDeclarationNode>;
var
    analyzer: TUpvalueAnalyzer;
begin
    // Note: AParent (IScopeDescriptor) removed from arguments as this analyzer
    // builds its own structural view and only cares about boxing internal variables.

    if not Assigned(ARootNode) then
        exit(THashSet<IVariableDeclarationNode>.Create);

    analyzer := TUpvalueAnalyzer.Create;
    try
        analyzer.Execute(ARootNode);
        Result := analyzer.FBoxedDeclarations;
        analyzer.FBoxedDeclarations := nil; // Transfer ownership
    finally
        analyzer.Free;
    end;
end;

procedure TUpvalueAnalyzer.MarkDeclarationForBoxing(const AName: string);
var
    depth: Integer;
    declNode: IVariableDeclarationNode;
begin
    // Check if the symbol exists in our analysis scopes
    if FCurrentScope.Resolve(AName, depth, declNode) then
    begin
        // If it is found in a parent scope (Depth > 0) AND we have a trackable declaration node for it
        if (depth > 0) and Assigned(declNode) then
        begin
            FBoxedDeclarations.Add(declNode);
        end;
    end;
end;

function TUpvalueAnalyzer.VisitIdentifier(const Node: IAstNode): IAstNode;
begin
    // Check if this identifier refers to a variable from an outer scope
    MarkDeclarationForBoxing(Node.AsIdentifier.Name);

    // Return original node (Analysis pass only)
    Result := Node;
end;

function TUpvalueAnalyzer.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    i: Integer;
begin
    L := Node.AsLambdaExpression;

    // 1. Enter new analysis scope
    FCurrentScope := TAnalysisScope.Create(FCurrentScope);
    try
        // 2. Register parameters (they mask outer variables)
        // We pass 'nil' as the node because we currently don't box parameters,
        // but we must ensure Resolve() finds them so we don't accidentally box a shadowed variable.

        // Optimized: Access Elements array directly
        var paramElements := L.Parameters.Elements;
        for i := 0 to High(paramElements) do
        begin
            if paramElements[i].Kind = akIdentifier then
                FCurrentScope.Define(paramElements[i].AsIdentifier.Name, nil);
        end;

        // 3. Visit Body
        Accept(L.Body); // Recursive call

        // Rebuild if needed (default CoW behavior)
        Result := Node;
    finally
        // 4. Exit scope
        var temp := FCurrentScope;
        FCurrentScope := FCurrentScope.FParent;
        temp.Free;
    end;
end;

function TUpvalueAnalyzer.VisitVariableDeclaration(const Node: IAstNode): IAstNode;
var
    V: IVariableDeclarationNode;
begin
    V := Node.AsVariableDeclaration;

    // 1. Visit initializer first (it executes in the CURRENT scope)
    if Assigned(V.Initializer) then
        Accept(V.Initializer);

    // 2. Define the variable in the CURRENT scope
    // Store the Node reference so we can add it to FBoxedDeclarations if captured.
    FCurrentScope.Define(V.Target.AsIdentifier.Name, V);

    Result := Node;
end;

end.