unit Myc.Ast.Analyzer;

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
        FBoxedDeclarations: THashSet<IVariableDeclarationNode>;
        FCurrentDescriptor: IScopeDescriptor;
        FDeclarationMap: TDictionary<IScopeDescriptor, TDictionary<string, IVariableDeclarationNode>>;
        procedure MarkDeclarationForBoxing(const AName: string);
    protected
        // Overridden Visit methods to perform analysis during traversal.
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode; override;
        function VisitIdentifier(const Node: IIdentifierNode): IAstNode; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode; override;
    public
        constructor Create(const AParent: IScopeDescriptor);
        destructor Destroy; override;

        // Added Execute method
        function Execute(const ARootNode: IAstNode): IAstNode;

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

implementation

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

{ TUpvalueAnalyzer }

constructor TUpvalueAnalyzer.Create(const AParent: IScopeDescriptor);
begin
    inherited Create;
    FBoxedDeclarations := THashSet<IVariableDeclarationNode>.Create;
    FDeclarationMap :=
        TObjectDictionary<IScopeDescriptor, TDictionary<string, IVariableDeclarationNode>>
            .Create([doOwnsValues], TEqualityComparer<IScopeDescriptor>.Default);
    FCurrentDescriptor := TScope.CreateDescriptor(AParent);
end;

destructor TUpvalueAnalyzer.Destroy;
begin
    FDeclarationMap.Free;
    FBoxedDeclarations.Free;
    inherited Destroy;
end;

function TUpvalueAnalyzer.Execute(const ARootNode: IAstNode): IAstNode;
begin
    // Accept will call the Visit... methods and traverse the tree
    Result := Accept(ARootNode);
end;

class function TUpvalueAnalyzer.Analyze(const ARootNode: IAstNode; const AParent: IScopeDescriptor): THashSet<IVariableDeclarationNode>;
var
    analyzer: TUpvalueAnalyzer; // Changed to concrete type
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
    symbol: TResolvedSymbol;
    declarationScope: IScopeDescriptor;
    i: Integer;
begin
    symbol := FCurrentDescriptor.FindSymbol(AName);
    if symbol.Address.Kind <> akLocalOrParent then
        exit;

    // Walk up the scope chain to find the scope where the variable was declared.
    declarationScope := FCurrentDescriptor;
    for i := 1 to symbol.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.VisitIdentifier(const Node: IIdentifierNode): IAstNode;
var
    symbol: TResolvedSymbol;
begin
    if Assigned(FCurrentDescriptor) then
    begin
        symbol := FCurrentDescriptor.FindSymbol(Node.Name);

        if (symbol.Address.Kind = akLocalOrParent) and (symbol.Address.ScopeDepth > 0) then
        begin
            // This is an upvalue. Mark its original declaration for boxing.
            MarkDeclarationForBoxing(Node.Name);
        end;
    end;

    // This is a leaf node, do not call inherited.
    Result := Node;
end;

function TUpvalueAnalyzer.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    N: TLambdaExpressionNode;
begin
    N := (Node as TLambdaExpressionNode);

    // A lambda creates a new lexical scope, inheriting from the current one.
    FCurrentDescriptor := TScope.CreateDescriptor(FCurrentDescriptor);
    try
        // Define the lambda's parameters within its new scope.
        // We use TTypes.Unknown as type inference hasn't run yet.
        for var param in N.Parameters do
            FCurrentDescriptor.Define(Accept(param).AsIdentifier.Name, TTypes.Unknown);

        // Traverse the lambda body within the new scope context.
        N.Body := Accept(N.Body); // Manual traversal

        Result := N;
    finally
        // Restore the parent scope after leaving the lambda.
        FCurrentDescriptor := FCurrentDescriptor.Parent;
    end;
end;

function TUpvalueAnalyzer.VisitVariableDeclaration(const Node: IVariableDeclarationNode): IAstNode;
var
    scopeDeclarations: TDictionary<string, IVariableDeclarationNode>;
    N: TVariableDeclarationNode;
begin
    N := (Node as TVariableDeclarationNode);

    // Traverse the initializer first. It's evaluated in the current scope
    // before the new variable is defined.
    if Assigned(N.Initializer) then
        N.Initializer := Accept(N.Initializer);

    // Traverse the identifier
    Accept(N.Identifier);

    // After processing the initializer, define the variable in the current scope.
    // We use TTypes.Unknown as type inference hasn't run yet.
    FCurrentDescriptor.Define(N.Identifier.Name, TTypes.Unknown);

    // Map this declaration node to its scope and name for later lookup.
    if not FDeclarationMap.TryGetValue(FCurrentDescriptor, scopeDeclarations) then
    begin
        scopeDeclarations := TDictionary<string, IVariableDeclarationNode>.Create;
        FDeclarationMap.Add(FCurrentDescriptor, scopeDeclarations);
    end;
    scopeDeclarations.Add(N.Identifier.Name, N);

    Result := N;
end;

end.