unit Myc.Ast.Binding;

interface

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

type
    // The binder is a transformer that enriches the AST with semantic information
    // like resolved addresses, scopes, and tail-call annotations.
    IAstBinder = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;
    end;

    TAstBinder = class(TAstTransformer, IAstBinder)
    private
        type
            // Helper class to track upvalues for a lambda expression.
            TUpvalueMapping = class
            public
                Map: TDictionary<TResolvedAddress, Integer>;
                Nodes: TList<IIdentifierNode>;
                constructor Create;
                destructor Destroy; override;
            end;
    private
        FCurrentDescriptor: IScopeDescriptor;
        FUpvalueStack: TStack<TUpvalueMapping>;
        FNestedLambdaCount: Integer;
        FIsTailStack: TStack<Boolean>;
        FNextIsTail: Boolean;

        procedure EnterScope;
        procedure ExitScope;
        function IsValidIdentifier(const Name: string): Boolean;

    protected
        // Stack management for tail-call state is centralized here.
        function Accept(const Node: IAstNode): TDataValue; override;

    public
        constructor Create(const AInitialScope: IExecutionScope);
        destructor Destroy; override;

        function Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;

        class function CreateDescriptor(const Scope: IExecutionScope): IScopeDescriptor; static;

        // The binder overrides specific transform methods to enrich the AST.
        function TransformIdentifier(const Node: IIdentifierNode): IIdentifierNode; override;
        function TransformVariableDeclaration(const Node: IVariableDeclarationNode): IVariableDeclarationNode; override;
        function TransformAssignment(const Node: IAssignmentNode): IAssignmentNode; override;
        function TransformLambdaExpression(const Node: ILambdaExpressionNode): ILambdaExpressionNode; override;
        function TransformFunctionCall(const Node: IFunctionCallNode): IFunctionCallNode; override;
        function TransformRecur(const Node: IRecurNode): IRecurNode; override;
        function TransformBlockExpression(const Node: IBlockExpressionNode): IBlockExpressionNode; override;
        function TransformIfExpression(const Node: IIfExpressionNode): IIfExpressionNode; override;
        function TransformTernaryExpression(const Node: ITernaryExpressionNode): ITernaryExpressionNode; override;
        function TransformBinaryExpression(const Node: IBinaryExpressionNode): IBinaryExpressionNode; override;
        function TransformUnaryExpression(const Node: IUnaryExpressionNode): IUnaryExpressionNode; override;
    end;

    TBoundIdentifierNode = class(TIdentifierNode)
    private
        FAddress: TResolvedAddress;
    public
        constructor Create(const AUnboundNode: IIdentifierNode; const AAddress: TResolvedAddress);
        property Address: TResolvedAddress read FAddress;
    end;

    TBoundLambdaExpressionNode = class(TLambdaExpressionNode)
    private
        FScopeDescriptor: IScopeDescriptor;
        FUpvalues: TArray<TResolvedAddress>;
        FHasNestedLambdas: Boolean;
    public
        constructor Create(
            const AUnboundNode: ILambdaExpressionNode;
            const ABody: IAstNode;
            const AParameters: TArray<IIdentifierNode>;
            const AScopeDescriptor: IScopeDescriptor;
            const AUpvalues: TArray<TResolvedAddress>;
            AHasNestedLambdas: Boolean
        );

        property ScopeDescriptor: IScopeDescriptor read FScopeDescriptor;
        property Upvalues: TArray<TResolvedAddress> read FUpvalues;
        property HasNestedLambdas: Boolean read FHasNestedLambdas;
    end;

    TBoundFunctionCallNode = class(TFunctionCallNode)
    private
        FIsTailCall: Boolean;
    public
        constructor Create(
            const AUnboundNode: IFunctionCallNode;
            const ACallee: IAstNode;
            const AArguments: TArray<IAstNode>;
            AIsTailCall: Boolean
        );
        property IsTailCall: Boolean read FIsTailCall;
    end;

implementation

uses
    System.Generics.Defaults,
    System.Character;

type
    // A custom equality comparer for TResolvedAddress to ensure correct behavior in TDictionary.
    TResolvedAddressComparer = class(TEqualityComparer<TResolvedAddress>)
    public
        function Equals(const Left, Right: TResolvedAddress): Boolean; override;
        function GetHashCode(const Value: TResolvedAddress): Integer; override;
    end;

{ TBoundIdentifierNode }
constructor TBoundIdentifierNode.Create(const AUnboundNode: IIdentifierNode; const AAddress: TResolvedAddress);
begin
    inherited Create(AUnboundNode.Name);
    FAddress := AAddress;
end;

{ TBoundLambdaExpressionNode }
constructor TBoundLambdaExpressionNode.Create(
    const AUnboundNode: ILambdaExpressionNode;
    const ABody: IAstNode;
    const AParameters: TArray<IIdentifierNode>;
    const AScopeDescriptor: IScopeDescriptor;
    const AUpvalues: TArray<TResolvedAddress>;
    AHasNestedLambdas: Boolean
);
begin
    inherited Create(AParameters, ABody);
    FScopeDescriptor := AScopeDescriptor;
    FUpvalues := AUpvalues;
    FHasNestedLambdas := AHasNestedLambdas;
end;

{ TBoundFunctionCallNode }
constructor TBoundFunctionCallNode.Create(
    const AUnboundNode: IFunctionCallNode;
    const ACallee: IAstNode;
    const AArguments: TArray<IAstNode>;
    AIsTailCall: Boolean
);
begin
    inherited Create(ACallee, AArguments);
    FIsTailCall := AIsTailCall;
end;

{ TResolvedAddressComparer }
function TResolvedAddressComparer.Equals(const Left, Right: TResolvedAddress): Boolean;
begin
    Result := (Left = Right);
end;

function TResolvedAddressComparer.GetHashCode(const Value: TResolvedAddress): Integer;
begin
    Result := 17;
    Result := Result * 23 + Ord(Value.Kind);
    Result := Result * 23 + Value.ScopeDepth;
    Result := Result * 23 + Value.SlotIndex;
end;

{ TAstBinder.TUpvalueMapping }
constructor TAstBinder.TUpvalueMapping.Create;
begin
    inherited Create;
    Map := TDictionary<TResolvedAddress, Integer>.Create(TResolvedAddressComparer.Create);
    Nodes := TList<IIdentifierNode>.Create();
end;

destructor TAstBinder.TUpvalueMapping.Destroy;
begin
    Nodes.Free;
    Map.Free;
    inherited Destroy;
end;

class function TAstBinder.CreateDescriptor(const Scope: IExecutionScope): IScopeDescriptor;
begin
    if Scope is TExecutionScope then
    begin
        var res := TScopeDescriptor.Create(CreateDescriptor(Scope.Parent));
        res.PopulateFromScope(Scope as TExecutionScope);
        Result := res;
    end
    else
        Result := TScopeDescriptor.Create(nil);
end;

constructor TAstBinder.Create(const AInitialScope: IExecutionScope);
begin
    inherited Create;
    FCurrentDescriptor := CreateDescriptor(AInitialScope);
    FUpvalueStack := TObjectStack<TUpvalueMapping>.Create(True);
    FNestedLambdaCount := 0;
    FIsTailStack := TStack<Boolean>.Create;
    FNextIsTail := True;
end;

destructor TAstBinder.Destroy;
begin
    FIsTailStack.Free;
    FUpvalueStack.Free;
    inherited;
end;

function TAstBinder.Accept(const Node: IAstNode): TDataValue;
begin
    if (not Assigned(Node)) or Done then
        exit;

    FIsTailStack.Push(FNextIsTail);
    try
        Result := inherited Accept(Node);
    finally
        FNextIsTail := FIsTailStack.Pop;
    end;
end;

procedure TAstBinder.EnterScope;
begin
    FCurrentDescriptor := TScopeDescriptor.Create(FCurrentDescriptor);
end;

function TAstBinder.Execute(const RootNode: IAstNode; out Descriptor: IScopeDescriptor): IAstNode;
begin
    EnterScope;
    try
        Result := Accept(RootNode).AsIntf<IAstNode>;
        Descriptor := FCurrentDescriptor;
    finally
        ExitScope;
    end;
end;

procedure TAstBinder.ExitScope;
begin
    FCurrentDescriptor := FCurrentDescriptor.Parent;
end;

function TAstBinder.IsValidIdentifier(const Name: string): Boolean;
var
    c: Char;
begin
    if Name.IsEmpty then
        exit(False);

    c := Name[1];
    if not (c.IsLetter or (c = '_')) then
        exit(False);

    for c in Name do
    begin
        if not (c.IsLetterOrDigit or (c = '_') or (c = '-')) then
            exit(False);
    end;
    Result := True;
end;

function TAstBinder.TransformIdentifier(const Node: IIdentifierNode): IIdentifierNode;
var
    depth, idx: Integer;
begin
    if FCurrentDescriptor.FindSymbol(Node.Name, depth, idx) then
    begin
        if (depth > 0) and (FUpvalueStack.Count > 0) then
        begin
            var upvalue := FUpvalueStack.Peek;
            dec(depth);
            var originalAddress := TResolvedAddress.Create(akLocalOrParent, depth, idx);

            var upvalueIndex: Integer;
            if not upvalue.Map.TryGetValue(originalAddress, upvalueIndex) then
            begin
                upvalueIndex := upvalue.Map.Count;
                upvalue.Map.Add(originalAddress, upvalueIndex);
            end;
            var address := TResolvedAddress.Create(akUpvalue, 0, upvalueIndex);
            Result := TBoundIdentifierNode.Create(Node, address);
        end
        else
        begin
            var address := TResolvedAddress.Create(akLocalOrParent, depth, idx);
            Result := TBoundIdentifierNode.Create(Node, address);
        end
    end
    else
        raise Exception.CreateFmt('Undefined identifier: "%s"', [Node.Name]);
end;

function TAstBinder.TransformVariableDeclaration(const Node: IVariableDeclarationNode): IVariableDeclarationNode;
var
    initializer: IAstNode;
    slotIndex: Integer;
    address: TResolvedAddress;
    boundIdentifier: IIdentifierNode;
begin
    if not IsValidIdentifier(Node.Identifier.Name) then
        raise Exception.CreateFmt('Invalid identifier name: "%s".', [Node.Identifier.Name]);

    FNextIsTail := False;
    if Node.Initializer <> nil then
        initializer := Accept(Node.Initializer).AsIntf<IAstNode>;

    slotIndex := FCurrentDescriptor.Define(Node.Identifier.Name);
    address := TResolvedAddress.Create(akLocalOrParent, 0, slotIndex);

    boundIdentifier := TBoundIdentifierNode.Create(Node.Identifier, address);

    Result := TAst.VarDecl(boundIdentifier, initializer);
end;

function TAstBinder.TransformAssignment(const Node: IAssignmentNode): IAssignmentNode;
begin
    FNextIsTail := False;
    Result := inherited TransformAssignment(Node);
end;

function TAstBinder.TransformBinaryExpression(const Node: IBinaryExpressionNode): IBinaryExpressionNode;
begin
    FNextIsTail := False;
    Result := inherited TransformBinaryExpression(Node);
end;

function TAstBinder.TransformUnaryExpression(const Node: IUnaryExpressionNode): IUnaryExpressionNode;
begin
    FNextIsTail := False;
    Result := inherited TransformUnaryExpression(Node);
end;

function TAstBinder.TransformLambdaExpression(const Node: ILambdaExpressionNode): ILambdaExpressionNode;
var
    i: integer;
    boundParams: TArray<IIdentifierNode>;
    boundBody: IAstNode;
    lambdaScope: IScopeDescriptor;
    upvalues: TArray<TResolvedAddress>;
    hasNestedLambdas: Boolean;
    lastNestedLambdaCount: Integer;
begin
    FUpvalueStack.Push(TUpvalueMapping.Create);
    try
        EnterScope;
        try
            FCurrentDescriptor.Define('<self>');

            SetLength(boundParams, Length(Node.Parameters));
            for i := 0 to High(Node.Parameters) do
            begin
                var paramNode := Node.Parameters[i];
                var slotIndex := FCurrentDescriptor.Define(paramNode.Name);
                var address := TResolvedAddress.Create(akLocalOrParent, 0, slotIndex);
                boundParams[i] := TBoundIdentifierNode.Create(paramNode, address);
            end;

            lastNestedLambdaCount := FNestedLambdaCount;

            FNextIsTail := True;
            boundBody := Accept(Node.Body).AsIntf<IAstNode>;

            hasNestedLambdas := FNestedLambdaCount > lastNestedLambdaCount;
            lambdaScope := FCurrentDescriptor;
        finally
            ExitScope;
        end;

        var upvalueMapping := FUpvalueStack.Peek;
        var sortedPairs := upvalueMapping.Map.ToArray;
        TArray.Sort<TPair<TResolvedAddress, Integer>>(
            sortedPairs,
            TComparer<TPair<TResolvedAddress, Integer>>.Construct(
                function(const Left, Right: TPair<TResolvedAddress, Integer>): Integer begin Result := Left.Value - Right.Value; end
            )
        );

        SetLength(upvalues, Length(sortedPairs));
        for i := 0 to High(sortedPairs) do
            upvalues[i] := sortedPairs[i].Key;

    finally
        FUpvalueStack.Pop;
    end;

    inc(FNestedLambdaCount);
    Result := TBoundLambdaExpressionNode.Create(Node, boundBody, boundParams, lambdaScope, upvalues, hasNestedLambdas);
end;

function TAstBinder.TransformFunctionCall(const Node: IFunctionCallNode): IFunctionCallNode;
var
    isTailCall: Boolean;
    callee: IAstNode;
    args: TArray<IAstNode>;
begin
    isTailCall := FIsTailStack.Peek;

    FNextIsTail := False;
    callee := Accept(Node.Callee).AsIntf<IAstNode>;
    args := TransformNodes<IAstNode>(Node.Arguments);

    Result := TBoundFunctionCallNode.Create(Node, callee, args, isTailCall);
end;

function TAstBinder.TransformRecur(const Node: IRecurNode): IRecurNode;
begin
    if not FIsTailStack.Peek then
        raise Exception.Create('''recur'' can only be used in a tail position.');

    FNextIsTail := False;
    Result := inherited TransformRecur(Node);
end;

function TAstBinder.TransformBlockExpression(const Node: IBlockExpressionNode): IBlockExpressionNode;
var
    exprs: TArray<IAstNode>;
    i: Integer;
    isContextTail: Boolean;
begin
    isContextTail := FIsTailStack.Peek;

    SetLength(exprs, Node.Expressions.Count);
    for i := 0 to Node.Expressions.Count - 1 do
    begin
        FNextIsTail := isContextTail and (i = Node.Expressions.Count - 1);
        exprs[i] := Accept(Node.Expressions[i]).AsIntf<IAstNode>;
    end;
    Result := TAst.Block(exprs);
end;

function TAstBinder.TransformIfExpression(const Node: IIfExpressionNode): IIfExpressionNode;
var
    isContextTail: Boolean;
    condition, thenBranch, elseBranch: IAstNode;
begin
    isContextTail := FIsTailStack.Peek;

    FNextIsTail := False;
    condition := Accept(Node.Condition).AsIntf<IAstNode>;

    FNextIsTail := isContextTail;
    thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    elseBranch := Accept(Node.ElseBranch).AsIntf<IAstNode>;

    if (condition <> Node.Condition) or (thenBranch <> Node.ThenBranch) or (elseBranch <> Node.ElseBranch) then
        Result := TAst.IfExpr(condition, thenBranch, elseBranch)
    else
        Result := Node;
end;

function TAstBinder.TransformTernaryExpression(const Node: ITernaryExpressionNode): ITernaryExpressionNode;
var
    isContextTail: Boolean;
    condition, thenBranch, elseBranch: IAstNode;
begin
    isContextTail := FIsTailStack.Peek;

    FNextIsTail := False;
    condition := Accept(Node.Condition).AsIntf<IAstNode>;

    FNextIsTail := isContextTail;
    thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    elseBranch := Accept(Node.ElseBranch).AsIntf<IAstNode>;

    if (condition <> Node.Condition) or (thenBranch <> Node.ThenBranch) or (elseBranch <> Node.ElseBranch) then
        Result := TAst.TernaryExpr(condition, thenBranch, elseBranch)
    else
        Result := Node;
end;

end.