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;

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>;
        FCurrentScope: 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): TDataValue; override;
        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue; override;
    public
        constructor Create(const AParent: IScopeDescriptor);
        destructor Destroy; override;
        class function Analyze(const ARootNode: IAstNode; const AParent: IScopeDescriptor): THashSet<IVariableDeclarationNode>; static;
    end;

implementation

uses
    Myc.Ast.Types;

{ 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
    for var dict in FDeclarationMap.Values do
        dict.Free;
    FDeclarationMap.Free;
    FBoxedDeclarations.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
    symbol: TResolvedSymbol;
    declarationScope: IScopeDescriptor;
    i: Integer;
begin
    symbol := FCurrentScope.FindSymbol(AName);
    if symbol.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 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): TDataValue;
var
    symbol: TResolvedSymbol;
begin
    if Assigned(FCurrentScope) then
    begin
        symbol := FCurrentScope.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;

    // As a traverser, return the original node wrapped in a TDataValue.
    Result := TDataValue.FromIntf<IIdentifierNode>(Node);
end;

function TUpvalueAnalyzer.VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue;
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.
        // We use TTypes.Unknown as type inference hasn't run yet.
        for var param in Node.Parameters do
            FCurrentScope.Define(param.Name, TTypes.Unknown);

        // Traverse the lambda body within the new scope context.
        Node.Body.Accept(Self);

        // We do not transform, just analyze. Return the original node wrapped.
        Result := TDataValue.FromIntf<ILambdaExpressionNode>(Node);
    finally
        // Restore the parent scope after leaving the lambda.
        FCurrentScope := FCurrentScope.Parent;
    end;
end;

function TUpvalueAnalyzer.VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue;
var
    scopeDeclarations: TDictionary<string, IVariableDeclarationNode>;
begin
    // Traverse the initializer first. It's evaluated in the current scope
    // before the new variable is defined.
    if Assigned(Node.Initializer) then
        Node.Initializer.Accept(Self);

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

    // 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);

    // As a traverser, return the original node wrapped in a TDataValue.
    Result := TDataValue.FromIntf<IVariableDeclarationNode>(Node);
end;

end.