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,
    Myc.Ast.Binding.Nodes;

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 TBoundFunctionCallNode.
    TAstTCO = class(TAstTransformer, IAstTCO)
    private
        FIsTailStack: TStack<Boolean>;
        FNextIsTail: Boolean;
    protected
        function Accept(const Node: IAstNode): TDataValue; override;
        // Overrides for TCO propagation
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; override;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; override;
        function VisitRecurNode(const Node: IRecurNode): TDataValue; override;
        function VisitMacroExpansionNode(const Node: IMacroExpansionNode): TDataValue; override;
        // The core TCO logic
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; 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
    var transformedValue := Accept(RootNode);
    if transformedValue.IsVoid then
        Result := TAst.Block([])
    else
        Result := transformedValue.AsIntf<IAstNode>;
end;

function TAstTCO.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;

function TAstTCO.VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue;
var
    i: Integer;
    isContextTail: Boolean;
begin
    isContextTail := FIsTailStack.Peek;
    for i := 0 to High(Node.Expressions) do
    begin
        // Only the last expression in a block is in tail position
        FNextIsTail := isContextTail and (i = High(Node.Expressions));
        Accept(Node.Expressions[i]);
    end;
    Result := TDataValue.FromIntf<IBlockExpressionNode>(Node);
end;

function TAstTCO.VisitIfExpression(const Node: IIfExpressionNode): TDataValue;
var
    isContextTail: Boolean;
begin
    isContextTail := FIsTailStack.Peek;

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

    // Then/Else branches ARE in tail position if the IfExpr is
    FNextIsTail := isContextTail;
    Accept(Node.ThenBranch);
    if Assigned(Node.ElseBranch) then
        Accept(Node.ElseBranch);

    Result := TDataValue.FromIntf<IIfExpressionNode>(Node);
end;

function TAstTCO.VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue;
var
    isContextTail: Boolean;
begin
    isContextTail := FIsTailStack.Peek;

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

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

    Result := TDataValue.FromIntf<ITernaryExpressionNode>(Node);
end;

function TAstTCO.VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue;
begin
    // Parameters are never in tail position
    FNextIsTail := False;
    AcceptNodes<IIdentifierNode>(Node.Parameters);

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

    Result := TDataValue.FromIntf<ILambdaExpressionNode>(Node);
end;

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

    // Arguments are not in tail position
    FNextIsTail := False;
    AcceptNodes<IAstNode>(Node.Arguments);

    Result := TDataValue.FromIntf<IRecurNode>(Node);
end;

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

    // Also visit the original call nodes, though they are not in tail pos
    FNextIsTail := False;
    Accept(Node.Callee);
    AcceptNodes<IAstNode>(Node.Arguments);

    Result := TDataValue.FromIntf<IMacroExpansionNode>(Node);
end;

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

    // Arguments are not in tail position
    FNextIsTail := False;
    callee := Accept(Node.Callee).AsIntf<IAstNode>;
    args := AcceptNodes<IAstNode>(Node.Arguments);

    // If TCO status changed, we MUST rebuild the node.
    var boundNode := (Node as TBoundFunctionCallNode);
    if (boundNode.IsTailCall <> isTailCall) or (callee <> Node.Callee) or (args <> Node.Arguments) then
    begin
        var newBoundCall := TBoundFunctionCallNode.Create(Node, callee, args, isTailCall);
        (newBoundCall as TAstNode).StaticType := (Node as TAstNode).StaticType;
        Result := TDataValue.FromIntf<IFunctionCallNode>(newBoundCall);
    end
    else
        Result := TDataValue.FromIntf<IFunctionCallNode>(Node); // No change
end;

end.