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<Boolean>)
        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;
        FNextIsTail: Boolean;
        procedure EnterScope;
        procedure ExitScope;

    protected
        function EnterNode(const Node: IAstNode): Boolean; 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;
        property NextIsTail: Boolean write FNextIsTail;
    end;

implementation

uses
    System.Generics.Defaults,
    Myc.Ast;

type
    TResolvedAddressComparer = class(TEqualityComparer<TResolvedAddress>)
    public
        function Equals(const Left, Right: TResolvedAddress): Boolean; override;
        function GetHashCode(const Value: TResolvedAddress): Integer; override;
    end;

    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;

{ TResolvedAddressComparer }

function TResolvedAddressComparer.Equals(const Left, Right: TResolvedAddress): Boolean;
begin
    Result := (Left.Kind = Right.Kind) and (Left.ScopeDepth = Right.ScopeDepth) and (Left.SlotIndex = Right.SlotIndex);
end;

function TResolvedAddressComparer.GetHashCode(const Value: TResolvedAddress): Integer;
begin
    // Simple combining hash function
    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;
    FNextIsTail := False;
end;

destructor TAstBinder.Destroy;
begin
    FUpvalueStack.Free;
    inherited;
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. The content of the root node is in a tail position.
            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.EnterNode(const Node: IAstNode): Boolean;
begin
    // The context for the current node is whatever the parent Visit... method prepared.
    Result := FNextIsTail;
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;
var
    i: Integer;
begin
    for i := 0 to Node.Expressions.Count - 2 do
    begin
        FNextIsTail := False;
        Accept(Node.Expressions[i]);
    end;

    if Node.Expressions.Count > 0 then
    begin
        // The last expression is in a tail position IF the block itself is.
        FNextIsTail := Data.Peek;
        Accept(Node.Expressions.Last);
    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 := Data.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.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 := Data.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 as TScopeDescriptor).Define('Self');

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

            var lastNestedLambdaCount := FNestedLambdaCount;

            // Manually traverse children since we can't call 'inherited' from the simple traverser.
            // 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;

        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;

        FUpvalueStack.Pop;

        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 := Data.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 as TScopeDescriptor).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
        if not FSymbols.ContainsKey(pair.Key) then
            FSymbols.Add(pair.Key, pair.Value);
end;

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

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

end.