unit Myc.Ast.Compiler.TCO;

interface

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

type
    IAstTCO = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode): IAstNode;
    end;

    // This transformer runs *after* the Lowerer (Phase 4).
    // Its sole responsibility is to identify tail calls (TCO)
    // and set the `IsTailCall` flag on TFunctionCallNode.
    TAstTCO = class(TAstTransformer, IAstTCO)
    private
        FIsTailStack: TStack<Boolean>;
        FNextIsTail: Boolean;
    protected
        function Accept(const Node: IAstNode): IAstNode; override;
        // Overrides for TCO propagation
        function VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode; override;
        function VisitIfExpression(const Node: IIfExpressionNode): IAstNode; override;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): IAstNode; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode; override;
        function VisitRecurNode(const Node: IRecurNode): IAstNode; override;
        function VisitMacroExpansionNode(const Node: IMacroExpansionNode): IAstNode; override;

        // Operands are never in tail position.
        function VisitBinaryExpression(const Node: IBinaryExpressionNode): IAstNode; override;
        function VisitUnaryExpression(const Node: IUnaryExpressionNode): IAstNode; override;

        // The core TCO logic
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;
    public
        constructor Create;
        destructor Destroy; override;
        function Execute(const RootNode: IAstNode): IAstNode;

        class function Optimize(const RootNode: IAstNode): IAstNode; static;
    end;

implementation

{ TAstTCO }

constructor TAstTCO.Create;
begin
    inherited Create;
    FIsTailStack := TStack<Boolean>.Create;
    FNextIsTail := True; // The root expression is in tail position
end;

destructor TAstTCO.Destroy;
begin
    FIsTailStack.Free;
    inherited;
end;

class function TAstTCO.Optimize(const RootNode: IAstNode): IAstNode;
begin
    var optimizer := TAstTCO.Create as IAstTCO;
    Result := optimizer.Execute(RootNode);
end;

function TAstTCO.Execute(const RootNode: IAstNode): IAstNode;
begin
    Result := Accept(RootNode); // Use IAstNode-returning Accept
    if not Assigned(Result) then
        Result := TAst.Block([]);
end;

function TAstTCO.Accept(const Node: IAstNode): IAstNode;
begin
    if (not Assigned(Node)) then
    begin
        Result := nil;
        exit;
    end;

    FIsTailStack.Push(FNextIsTail);
    try
        // Call inherited Accept, which handles the data unwrapping
        Result := inherited Accept(Node);
    finally
        FNextIsTail := FIsTailStack.Pop;
    end;
end;

function TAstTCO.VisitBinaryExpression(const Node: IBinaryExpressionNode): IAstNode;
begin
    // Operands are never in tail position.
    FNextIsTail := False;

    // Call inherited, which will visit Left and Right with FNextIsTail = False
    Result := inherited VisitBinaryExpression(Node);
end;

function TAstTCO.VisitUnaryExpression(const Node: IUnaryExpressionNode): IAstNode;
begin
    // Operand is never in tail position.
    FNextIsTail := False;

    // Call inherited, which will visit Right with FNextIsTail = False
    Result := inherited VisitUnaryExpression(Node);
end;

function TAstTCO.VisitBlockExpression(const Node: IBlockExpressionNode): IAstNode;
var
    isContextTail: Boolean;
    newExprs: TArray<IAstNode>;
begin
    isContextTail := FIsTailStack.Peek;

    var nTail := High(Node.Expressions);
    newExprs :=
        AcceptNodes(
            Node.Expressions,
            function(idx: Integer; Node: IAstNode): IAstNode
            begin
                FNextIsTail := isContextTail and (idx = nTail);
                Result := Accept(Node);
            end
        );

    if newExprs = Node.Expressions then
        Result := Node
    else
        Result := TBlockExpressionNode.Create(newExprs, Node.StaticType);
end;

function TAstTCO.VisitIfExpression(const Node: IIfExpressionNode): IAstNode;
var
    isContextTail: Boolean;
    N: TIfExpressionNode;
    newCond, newThen, newElse: IAstNode;
begin
    N := (Node as TIfExpressionNode);
    isContextTail := FIsTailStack.Peek;

    // Condition is never in tail position
    FNextIsTail := False;
    newCond := Accept(N.Condition);

    // Then/Else branches ARE in tail position if the IfExpr is
    FNextIsTail := isContextTail;
    newThen := Accept(N.ThenBranch);
    newElse := Accept(N.ElseBranch);

    if (newCond = N.Condition) and (newThen = N.ThenBranch) and (newElse = N.ElseBranch) then
        Result := Node
    else
        Result := TIfExpressionNode.Create(newCond, newThen, newElse, N.StaticType);
end;

function TAstTCO.VisitTernaryExpression(const Node: ITernaryExpressionNode): IAstNode;
var
    isContextTail: Boolean;
    N: TTernaryExpressionNode;
    newCond, newThen, newElse: IAstNode;
begin
    N := (Node as TTernaryExpressionNode);
    isContextTail := FIsTailStack.Peek;

    // Condition is never in tail position
    FNextIsTail := False;
    newCond := Accept(N.Condition);

    // Then/Else branches ARE in tail position if the TernaryExpr is
    FNextIsTail := isContextTail;
    newThen := Accept(N.ThenBranch);
    newElse := Accept(N.ElseBranch);

    if (newCond = N.Condition) and (newThen = N.ThenBranch) and (newElse = N.ElseBranch) then
        Result := Node
    else
        Result := TTernaryExpressionNode.Create(newCond, newThen, newElse, N.StaticType);
end;

function TAstTCO.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    N: TLambdaExpressionNode;
    newParams: TArray<IIdentifierNode>;
    newBody: IAstNode;
begin
    N := (Node as TLambdaExpressionNode);

    // Parameters are not in tail position (handled by AcceptParameters)
    FNextIsTail := False;
    newParams := AcceptParameters(N.Parameters);

    // The body of a lambda is *always* a tail position (relative to the lambda)
    FNextIsTail := True;
    newBody := Accept(N.Body);

    if (newParams = N.Parameters) and (newBody = N.Body) then
        Result := Node
    else
    begin
        Result := TLambdaExpressionNode.Create(newParams, newBody, N.StaticType);
        // Copy runtime properties
        (Result as TLambdaExpressionNode).ScopeDescriptor := N.ScopeDescriptor;
        (Result as TLambdaExpressionNode).Upvalues := N.Upvalues;
        (Result as TLambdaExpressionNode).HasNestedLambdas := N.HasNestedLambdas;
    end;
end;

function TAstTCO.VisitRecurNode(const Node: IRecurNode): IAstNode;
var
    N: TRecurNode;
    newArgs: TArray<IAstNode>;
begin
    if not FIsTailStack.Peek then
        raise Exception.Create('''recur'' can only be used in a tail position.');

    N := (Node as TRecurNode);

    // Arguments are not in tail position
    FNextIsTail := False;
    newArgs := AcceptNodes(N.Arguments);

    if newArgs = N.Arguments then
        Result := Node
    else
        Result := TRecurNode.Create(newArgs, N.StaticType);
end;

function TAstTCO.VisitMacroExpansionNode(const Node: IMacroExpansionNode): IAstNode;
var
    newBody: IAstNode;
begin
    // Propagate tail call status to the expanded body
    FNextIsTail := FIsTailStack.Peek;
    newBody := Accept(Node.ExpandedBody);

    if newBody = Node.ExpandedBody then
        Result := Node
    else
        // Rebuild, preserving the original CallNode metadata
        Result := TMacroExpansionNode.Create(Node.CallNode, newBody.AsTypedNode);
end;

function TAstTCO.VisitFunctionCall(const Node: IFunctionCallNode): IAstNode;
var
    isTailCall: Boolean;
    N: TFunctionCallNode;
    newCallee: IAstNode;
    newArgs: TArray<IAstNode>;
begin
    isTailCall := FIsTailStack.Peek;
    N := (Node as TFunctionCallNode);

    // Arguments are not in tail position
    FNextIsTail := False;
    newCallee := Accept(N.Callee);
    newArgs := AcceptNodes(N.Arguments);

    // CoW check: Create a new node only if children changed OR IsTailCall needs update
    if (newCallee = N.Callee) and (newArgs = N.Arguments) and (isTailCall = N.IsTailCall) then
        Result := Node
    else
    begin
        Result := TFunctionCallNode.Create(newCallee, newArgs, N.StaticType);
        // Mutate the *new* node with the TCO status
        (Result as TFunctionCallNode).IsTailCall := isTailCall;
    end;
end;

end.