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
        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
    newParams: TArray<IIdentifierNode>;
    newBody: IAstNode;
    i: Integer;
begin
    // A lambda creates a new lexical scope, inheriting from the current one.
    FCurrentDescriptor := TScope.CreateDescriptor(FCurrentDescriptor);
    try
        // Manually visit parameters to define them
        var oldParams := Node.Parameters;
        SetLength(newParams, Length(oldParams));
        for i := 0 to High(oldParams) do
        begin
            // Accept(param) will call VisitIdentifier, which does its job.
            // We *must* use the result of Accept.
            newParams[i] := Accept(oldParams[i]).AsIdentifier;
            FCurrentDescriptor.Define(newParams[i].Name, TTypes.Unknown);
        end;

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

        // Rebuild node if changed
        if (newParams = Node.Parameters) and (newBody = Node.Body) then
            Result := Node
        else
            // Use factory, pass through other properties (which are nil/false at this stage)
            Result := TAst.LambdaExpr(newParams, newBody, Node.ScopeDescriptor, Node.Upvalues, Node.HasNestedLambdas, Node.StaticType);
    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>;
    newInitializer: IAstNode;
    newIdentifier: IIdentifierNode;
begin
    // Traverse the initializer first. It's evaluated in the current scope
    // before the new variable is defined.
    if Assigned(Node.Initializer) then
        newInitializer := Accept(Node.Initializer)
    else
        newInitializer := nil;

    // Traverse the identifier
    // This calls VisitIdentifier, but VisitIdentifier is just an analyzer, it returns the same node.
    newIdentifier := Accept(Node.Identifier).AsIdentifier;

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

    // --- Rebuild Node ---
    // Use CoW, check if anything changed
    if (newInitializer = Node.Initializer) and (newIdentifier = Node.Identifier) then
        Result := Node
    else
        // Use factory, pass through other properties (which are nil/false at this stage)
        Result := TAst.VarDecl(newIdentifier, newInitializer, Node.StaticType, Node.IsBoxed);

    // Map this *new* 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;
    // IMPORTANT: Store the *new* node (Result)
    scopeDeclarations.Add(newIdentifier.Name, Result.AsVariableDeclaration);
end;

end.