MycLib/Src/AST/Myc.Ast.Compiler.TCO.pas

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;

        // 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.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 := TAst.Block(newExprs, Node.StaticType);
end;

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

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

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

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

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

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

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

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

function TAstTCO.VisitLambdaExpression(const Node: ILambdaExpressionNode): IAstNode;
var
    newParams: TArray<IIdentifierNode>;
    newBody: IAstNode;
begin
    // Parameters are not in tail position (handled by AcceptParameters)
    FNextIsTail := False;
    newParams := AcceptParameters(Node.Parameters);

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

    if (newParams = Node.Parameters) and (newBody = Node.Body) then
        Result := Node
    else
    begin
        // Rebuild using factory and copy properties via interface
        Result := TAst.LambdaExpr(newParams, newBody, Node.ScopeDescriptor, Node.Upvalues, Node.HasNestedLambdas, Node.StaticType);
    end;
end;

function TAstTCO.VisitRecurNode(const Node: IRecurNode): IAstNode;
var
    newArgs: TArray<IAstNode>;
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;
    newArgs := AcceptNodes(Node.Arguments);

    if newArgs = Node.Arguments then
        Result := Node
    else
        Result := TAst.Recur(newArgs, Node.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 := TAst.MacroExpansionNode(Node.CallNode, newBody);
end;

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

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

    // CoW check: Create a new node only if children changed OR IsTailCall needs update
    // OR StaticTarget is missing (which it shouldn't be, but good to check)
    if (newCallee = Node.Callee) and (newArgs = Node.Arguments) and (isTailCall = Node.IsTailCall) then
    begin
        Result := Node;
        exit;
    end;

    // Use factory to create new node, passing the *new* TCO status
    Result :=
        TAst.FunctionCall(
            newCallee,
            newArgs,
            Node.StaticType,
            isTailCall,
            Node.StaticTarget // Preserve the static target
        );
end;

end.