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
    // Exception specific to optimization/TCO errors
    EOptimizerException = class(EAstException);

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

    // This transformer runs *after* the Lowerer (Phase 4/5) and Specializer.
    // 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;

    strict private
        // Typed Handlers (IAstNode signature)
        function VisitBlockExpression(const Node: IAstNode): IAstNode;
        function VisitIfExpression(const Node: IAstNode): IAstNode;
        function VisitCondExpression(const Node: IAstNode): IAstNode;
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;
        function VisitRecurNode(const Node: IAstNode): IAstNode;
        function VisitMacroExpansionNode(const Node: IAstNode): IAstNode;
        function VisitFunctionCall(const Node: IAstNode): IAstNode;
        function VisitTuple(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;

        // Intercept dispatch to manage the TCO Stack state
        function Accept(const Node: IAstNode): 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;

procedure TAstTCO.SetupHandlers;
begin
    inherited SetupHandlers; // Defaults

    Register(akBlockExpression, VisitBlockExpression);
    Register(akIfExpression, VisitIfExpression);
    Register(akCondExpression, VisitCondExpression);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akRecur, VisitRecurNode);
    Register(akMacroExpansion, VisitMacroExpansionNode);
    Register(akFunctionCall, VisitFunctionCall);
    Register(akTuple, VisitTuple);
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);
    if not Assigned(Result) then
        Result := TAst.Block([], nil);
end;

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

    // Push current context state before visiting children
    FIsTailStack.Push(FNextIsTail);
    try
        // Dispatch to handler (which will call VisitXyz)
        Result := inherited Accept(Node);
    finally
        // Restore context state
        FNextIsTail := FIsTailStack.Pop;
    end;
end;

function TAstTCO.VisitBlockExpression(const Node: IAstNode): IAstNode;
var
    B: IBlockExpressionNode;
    isContextTail: Boolean;
    newExprs: TArray<IAstNode>;
    exprsTuple: ITupleNode;
    i: Integer;
    item, newItem: IAstNode;
    hasChanged: Boolean;
    elements: TArray<IAstNode>;
begin
    B := Node.AsBlockExpression;
    isContextTail := FIsTailStack.Peek;
    exprsTuple := B.Expressions;
    elements := exprsTuple.Elements;

    var count := Length(elements);
    SetLength(newExprs, count);
    hasChanged := False;

    var nTail := count - 1;

    for i := 0 to nTail do
    begin
        item := elements[i];

        // Only the last expression in the block inherits the tail position status
        FNextIsTail := isContextTail and (i = nTail);

        newItem := Accept(item);
        newExprs[i] := newItem;

        if item <> newItem then
            hasChanged := True;
    end;

    if not hasChanged then
        Result := Node
    else
    begin
        // Create new Tuple for expressions
        Result := TAst.Block(Node.Identity, TAst.Tuple(B.Expressions.Identity, newExprs), B.StaticType);
    end;
end;

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

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

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

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

function TAstTCO.VisitCondExpression(const Node: IAstNode): IAstNode;
var
    C: ICondExpressionNode;
    isContextTail: Boolean;
    hasChanged: Boolean;
    i: Integer;
    newPairs: TArray<TCondPair>;
    newElse: IAstNode;
    newCond, newBranch: IAstNode;
begin
    C := Node.AsCondExpression;
    isContextTail := FIsTailStack.Peek;
    hasChanged := False;
    SetLength(newPairs, Length(C.Pairs));

    for i := 0 to High(C.Pairs) do
    begin
        // 1. Condition is never in tail position
        FNextIsTail := False;
        newCond := Accept(C.Pairs[i].Condition);

        // 2. Branch IS in tail position (if CondExpr is)
        FNextIsTail := isContextTail;
        newBranch := Accept(C.Pairs[i].Branch);

        newPairs[i] := TCondPair.Create(newCond, newBranch);

        if (newCond <> C.Pairs[i].Condition) or (newBranch <> C.Pairs[i].Branch) then
            hasChanged := True;
    end;

    // 3. Else Branch IS in tail position
    FNextIsTail := isContextTail;
    newElse := Accept(C.ElseBranch);

    if newElse <> C.ElseBranch then
        hasChanged := True;

    if not hasChanged then
        Result := Node
    else
        Result := TAst.CondExpr(Node.Identity, newPairs, newElse, C.StaticType);
end;

function TAstTCO.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    newParams: ITupleNode;
    newBody: IAstNode;
begin
    L := Node.AsLambdaExpression;
    // Parameters are not in tail position
    FNextIsTail := False;
    newParams := Accept(L.Parameters).AsTuple;

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

    if (newParams = L.Parameters) and (newBody = L.Body) then
        Result := Node
    else
    begin
        Result :=
            TAst.LambdaExpr(
                Node.Identity,
                newParams,
                newBody,
                L.Layout,
                L.Descriptor,
                L.Upvalues,
                L.HasNestedLambdas,
                L.IsPure,
                L.StaticType
            );
    end;
end;

function TAstTCO.VisitRecurNode(const Node: IAstNode): IAstNode;
var
    R: IRecurNode;
    newArgs: ITupleNode;
begin
    R := Node.AsRecur;
    if not FIsTailStack.Peek then
        raise EOptimizerException.Create('''recur'' can only be used in a tail position.');

    // Arguments are not in tail position
    FNextIsTail := False;
    newArgs := Accept(R.Arguments).AsTuple;

    if newArgs = R.Arguments then
        Result := Node
    else
        Result := TAst.Recur(Node.Identity, newArgs, R.StaticType);
end;

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

    if newBody = M.ExpandedBody then
        Result := Node
    else
        Result := TAst.MacroExpansionNode(Node.Identity, M.CallNode, newBody);
end;

function TAstTCO.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    isTailCall: Boolean;
    newCallee: IAstNode;
    newArgs: ITupleNode;
begin
    C := Node.AsFunctionCall;
    isTailCall := FIsTailStack.Peek;

    // Callee/Arguments are not in tail position
    FNextIsTail := False;

    newCallee := Accept(C.Callee);
    newArgs := Accept(C.Arguments).AsTuple;

    if (newCallee = C.Callee) and (newArgs = C.Arguments) and (isTailCall = C.IsTailCall) then
    begin
        Result := Node;
        exit;
    end;

    // Use factory to create new node with TCO status
    Result := TAst.FunctionCall(Node.Identity, newCallee, newArgs, C.StaticType, isTailCall, C.StaticTarget, C.IsTargetPure);
end;

function TAstTCO.VisitTuple(const Node: IAstNode): IAstNode;
var
    T: ITupleNode;
    newElements: TArray<IAstNode>;
    i: Integer;
    hasChanged: Boolean;
    savedNextIsTail: Boolean;
    elements: TArray<IAstNode>;
begin
    T := Node.AsTuple;
    savedNextIsTail := FNextIsTail; // Zustand sichern

    // Elemente in einem Tupel (Liste, Vektor, Argumente) sind NIEMALS in Tail-Position
    FNextIsTail := False;

    hasChanged := False;
    elements := T.Elements;
    var count := Length(elements);
    SetLength(newElements, count);

    try
        for i := 0 to count - 1 do
        begin
            newElements[i] := Accept(elements[i]);
            if newElements[i] <> elements[i] then
                hasChanged := True;
        end;
    finally
        FNextIsTail := savedNextIsTail;
    end;

    if hasChanged then
        Result := TAst.Tuple(Node.Identity, newElements, T.StaticType)
    else
        Result := Node;
end;

end.