unit Myc.Ast.Binding;

interface

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

type
    TAstBinder = class(TAstTraverser)
        type
            TUpvalueMapping = class
                Map: TDictionary<TResolvedAddress, Integer>;
                Nodes: TList<IIdentifierNode>;
            public
                constructor Create;
                destructor Destroy; override;
            end;
    private
        FCurrentDescriptor: IScopeDescriptor;
        FUpvalueStack: TStack<TUpvalueMapping>;
        FNestedLambdaCount: Integer;
        FIsTailStack: TStack<Boolean>;
        FNextIsTail: Boolean;
        procedure EnterScope;
        procedure ExitScope;

    protected
        function Accept(const Node: IAstNode): TDataValue; override;

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

        class function Bind(const RootNode: IAstNode; const ParentScope: IExecutionScope): IScopeDescriptor;
        class function CreateDescriptor(const Scope: IExecutionScope): IScopeDescriptor; static;

        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue; override;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; override;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
        function VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue; override;
        function VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue; override;
        function VisitAssignment(const Node: IAssignmentNode): TDataValue; override;

        property CurrentDescriptor: IScopeDescriptor read FCurrentDescriptor;
    end;

implementation

uses
    System.Generics.Defaults,
    Myc.Ast;

type
    TScopeDescriptor = class(TInterfacedObject, IScopeDescriptor)
    private
        FParent: IScopeDescriptor;
        FSymbols: TDictionary<string, Integer>;
        function GetParent: IScopeDescriptor;
        function GetSlotCount: Integer;
        function GetSymbols: TDictionary<string, Integer>;
    public
        constructor Create(const AParent: IScopeDescriptor);
        destructor Destroy; override;
        function Define(const Name: string): Integer;
        function FindSymbol(const Name: string; out Depth, Index: Integer): Boolean;
        function CreateScope(const Parent: IExecutionScope): IExecutionScope;
        procedure PopulateFromScope(Scope: TExecutionScope);
        property Symbols: TDictionary<string, Integer> read FSymbols;
    end;

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

{ TResolvedAddressComparer }

function TResolvedAddressComparer.Equals(const Left, Right: TResolvedAddress): Boolean;
begin
    // Use the existing equality operator for the record.
    Result := (Left = Right);
end;

function TResolvedAddressComparer.GetHashCode(const Value: TResolvedAddress): Integer;
begin
    // Classic hash combining algorithm using prime numbers.
    Result := 17;
    Result := Result * 23 + Ord(Value.Kind);
    Result := Result * 23 + Value.ScopeDepth;
    Result := Result * 23 + Value.SlotIndex;
end;

{ TAstBinder }

constructor TAstBinder.Create(const AInitialScope: IExecutionScope);
begin
    inherited Create;
    FCurrentDescriptor := CreateDescriptor(AInitialScope);
    FUpvalueStack := TObjectStack<TUpvalueMapping>.Create(true);
    FNestedLambdaCount := 0;
    FIsTailStack := TStack<Boolean>.Create;
    // The content of the root node is in a tail position.
    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;

class function TAstBinder.Bind(const RootNode: IAstNode; const ParentScope: IExecutionScope): IScopeDescriptor;
var
    binder: TAstBinder;
begin
    binder := TAstBinder.Create(ParentScope);
    try
        binder.EnterScope;
        try
            // Start the traversal
            binder.Accept(RootNode);
            Result := binder.CurrentDescriptor;
        finally
            binder.ExitScope;
        end;
    finally
        binder.Free;
    end;
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;

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

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

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

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

function TAstBinder.VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue;
begin
    FNextIsTail := False;

    var n := Node.Expressions.Count - 1;
    for var i := 0 to n do
    begin
        // The last expression is in a tail position IF the block itself is.
        if i = n then
            FNextIsTail := FIsTailStack.Peek;

        Accept(Node.Expressions[i]);
    end;
end;

function TAstBinder.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
begin
    // Annotate this node based on its context, which is on top of the stack.
    (Node as TFunctionCallNode).IsTailCall := FIsTailStack.Peek;

    // Let the default traverser visit children (callee, args), but ensure
    // their context is non-tail.
    FNextIsTail := False;
    inherited;
end;

function TAstBinder.VisitIdentifier(const Node: IIdentifierNode): TDataValue;
var
    depth, idx: Integer;
    identNode: TIdentifierNode;
    upvalue: TUpvalueMapping;
    originalAddress: TResolvedAddress;
    upvalueIndex: Integer;
begin
    identNode := Node as TIdentifierNode;
    if identNode.Address.Kind <> akUnresolved then
        exit;

    if FCurrentDescriptor.FindSymbol(identNode.Name, depth, idx) then
    begin
        if (depth > 0) and (FUpvalueStack.Count > 0) then
        begin
            upvalue := FUpvalueStack.Peek;

            // Address is relative to the lambda's parent scope.
            dec(depth);
            originalAddress := TResolvedAddress.Create(akLocalOrParent, depth, idx);

            if not upvalue.Map.TryGetValue(originalAddress, upvalueIndex) then
            begin
                upvalueIndex := upvalue.Map.Count;
                upvalue.Map.Add(originalAddress, upvalueIndex);
            end;

            (Node as TIdentifierNode).Address := TResolvedAddress.Create(akUpvalue, 0, upvalueIndex);
        end
        else
            (Node as TIdentifierNode).Address := TResolvedAddress.Create(akLocalOrParent, depth, idx);
    end
    else
        raise Exception.CreateFmt('Undefined identifier: "%s"', [identNode.Name]);
end;

(*
function TAstBinder.VisitIdentifier(const Node: IIdentifierNode): TDataValue;
var
    depth, idx: Integer;
    identNode: TIdentifierNode;
    upvalue: TUpvalueMapping;
    originalAddress: TResolvedAddress;
    upvalueIndex: Integer;
begin
    identNode := Node as TIdentifierNode;
    if identNode.Address.Kind <> akUnresolved then
        Exit;

    if FCurrentDescriptor.FindSymbol(identNode.Name, depth, idx) then
    begin
        if (depth > 0) and (FUpvalueStack.Count > 0) then
        begin
            upvalue := FUpvalueStack.Peek;

            // Address is relative to the lambda's parent scope.
            dec(depth);
            originalAddress.Create(akLocalOrParent, depth, idx);

            if not upvalue.Map.TryGetValue(originalAddress, upvalueIndex) then
            begin
                upvalueIndex := upvalue.Nodes.Count;
                upvalue.Map.Add(originalAddress, upvalueIndex);
                upvalue.Nodes.Add(identNode);
            end;

            (Node as TIdentifierNode).Address.Create(akUpvalue, 0, upvalueIndex);
        end
        else
            (Node as TIdentifierNode).Address.Create(akLocalOrParent, depth, idx);
    end
    else
        raise Exception.CreateFmt('Undefined identifier: "%s"', [identNode.Name]);
end;
*)
function TAstBinder.VisitIfExpression(const Node: IIfExpressionNode): TDataValue;
begin
    // The condition is never in a tail position.
    FNextIsTail := False;
    Accept(Node.Condition);

    // The branches are in a tail position if the if-expression itself is.
    FNextIsTail := FIsTailStack.Peek;
    Accept(Node.ThenBranch);
    if Assigned(Node.ElseBranch) then
        Accept(Node.ElseBranch);
end;

function TAstBinder.VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue;
var
    param: IIdentifierNode;
    sourceAddresses: TArray<TResolvedAddress>;
    sortedPairs: TArray<TPair<TResolvedAddress, Integer>>;
begin
    FUpvalueStack.Push(TUpvalueMapping.Create);
    try
        EnterScope;
        try
            FCurrentDescriptor.Define('Self');

            for param in Node.Parameters do
                FCurrentDescriptor.Define(param.Name);

            var lastNestedLambdaCount := FNestedLambdaCount;

            // Parameters are never in a tail position.
            FNextIsTail := False;
            for param in Node.Parameters do
                Accept(param);

            // The body of a lambda is always in a tail position.
            FNextIsTail := True;
            Accept(Node.Body);

            (Node as TLambdaExpressionNode).HasNestedLambdas := FNestedLambdaCount > lastNestedLambdaCount;
            (Node as TLambdaExpressionNode).ScopeDescriptor := FCurrentDescriptor;
        finally
            ExitScope;
        end;
    finally
        var upvalue := FUpvalueStack.Peek;
        try
            sortedPairs := upvalue.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(sourceAddresses, Length(sortedPairs));
            for var i := 0 to High(sortedPairs) do
                sourceAddresses[i] := sortedPairs[i].Key;

            (Node as TLambdaExpressionNode).Upvalues := sourceAddresses;
        finally
            FUpvalueStack.Pop;
        end;

        inc(FNestedLambdaCount);
    end;
end;

function TAstBinder.VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue;
begin
    // The condition is never in a tail position.
    FNextIsTail := False;
    Accept(Node.Condition);

    // The branches are in a tail position if the ternary expression itself is.
    FNextIsTail := FIsTailStack.Peek;
    Accept(Node.ThenBranch);
    Accept(Node.ElseBranch);
end;

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

function TAstBinder.VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue;
var
    slotIndex: Integer;
begin
    // The initializer expression is never in a tail position.
    FNextIsTail := False;
    if Assigned(Node.Initializer) then
        Accept(Node.Initializer);

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

    FNextIsTail := False;
    Accept(Node.Identifier);
end;

{ TScopeDescriptor }

constructor TScopeDescriptor.Create(const AParent: IScopeDescriptor);
begin
    inherited Create;
    FParent := AParent;
    FSymbols := TDictionary<string, Integer>.Create;
end;

destructor TScopeDescriptor.Destroy;
begin
    FSymbols.Free;
    inherited;
end;

function TScopeDescriptor.Define(const Name: string): Integer;
begin
    Result := FSymbols.Count;
    FSymbols.Add(Name, Result);
end;

function TScopeDescriptor.FindSymbol(const Name: string; out Depth, Index: Integer): Boolean;
var
    currentDescriptor: TScopeDescriptor;
begin
    Depth := 0;
    currentDescriptor := Self;
    while currentDescriptor <> nil do
    begin
        if currentDescriptor.FSymbols.TryGetValue(Name, Index) then
            exit(true);
        inc(Depth);
        currentDescriptor := currentDescriptor.FParent as TScopeDescriptor;
    end;
    Result := False;
end;

function TScopeDescriptor.GetParent: IScopeDescriptor;
begin
    Result := FParent;
end;

function TScopeDescriptor.GetSlotCount: Integer;
begin
    Result := FSymbols.Count;
end;

function TScopeDescriptor.GetSymbols: TDictionary<string, Integer>;
begin
    Result := FSymbols;
end;

function TScopeDescriptor.CreateScope(const Parent: IExecutionScope): IExecutionScope;
begin
    Result := TExecutionScope.Create(Parent, Self, nil);
end;

procedure TScopeDescriptor.PopulateFromScope(Scope: TExecutionScope);
begin
    for var pair in Scope.NameToIndex do
    begin
        var name := Scope.NameStrings[pair.Key];
        if not FSymbols.ContainsKey(name) then
            FSymbols.Add(name, pair.Value);
    end;
end;

constructor TAstBinder.TUpvalueMapping.Create;
begin
    inherited Create;
    // Use the custom equality comparer to ensure correct dictionary behavior.
    Map := TDictionary<TResolvedAddress, Integer>.Create(TResolvedAddressComparer.Create);
    Nodes := TList<IIdentifierNode>.Create();
end;

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

end.