MycLib/Src/AST/Myc.Ast.Lowering.pas

unit Myc.Ast.Lowering;

interface

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

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

    // This transformer runs *last* (after TypeChecker).
    // It "lowers" the AST by rewriting complex nodes into simpler ones
    // that the evaluator can understand (e.g., operator folding).
    // It also performs Tail Call Optimization (TCO) identification.
    TAstLowerer = class(TAstTransformer, IAstLowerer)
    private
        FBinaryOperators: TDictionary<string, TScalar.TBinaryOp>;
        FUnaryOperators: TDictionary<string, TScalar.TUnaryOp>;
        FIsTailStack: TStack<Boolean>;
        FNextIsTail: Boolean;
        function SetType(const Node: IAstNode; const AType: IStaticType): IAstNode;
    protected
        function Accept(const Node: IAstNode): TDataValue; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; override;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; override;
        function VisitRecurNode(const Node: IRecurNode): TDataValue; override;
        function VisitMacroExpansionNode(const Node: IMacroExpansionNode): TDataValue; override;

    public
        constructor Create;
        destructor Destroy; override;
        function Execute(const RootNode: IAstNode): IAstNode;

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

implementation

uses
    System.Generics.Defaults,
    Myc.Data.Keyword;

{ TAstLowerer }

constructor TAstLowerer.Create;
var
    op: TScalar.TBinaryOp;
begin
    inherited Create;

    // TCO state
    FIsTailStack := TStack<Boolean>.Create;
    FNextIsTail := True; // The root expression is in tail position

    // Operator folding maps
    FBinaryOperators := TDictionary<string, TScalar.TBinaryOp>.Create;
    for op := Low(TScalar.TBinaryOp) to High(TScalar.TBinaryOp) do
        FBinaryOperators.Add(op.ToString, op);

    FUnaryOperators := TDictionary<string, TScalar.TUnaryOp>.Create;
    FUnaryOperators.Add('not', TScalar.TUnaryOp.Not);
end;

destructor TAstLowerer.Destroy;
begin
    FIsTailStack.Free;
    FUnaryOperators.Free;
    FBinaryOperators.Free;
    inherited;
end;

class function TAstLowerer.Lower(const RootNode: IAstNode): IAstNode;
begin
    var lowerer := TAstLowerer.Create as IAstLowerer;
    Result := lowerer.Execute(RootNode);
end;

function TAstLowerer.Execute(const RootNode: IAstNode): IAstNode;
begin
    var transformedValue := Accept(RootNode);
    if transformedValue.IsVoid then
        Result := TAst.Block([])
    else
        Result := transformedValue.AsIntf<IAstNode>;
    (Result as TAstNode).StaticType := (Result as TAstNode).StaticType;
end;

function TAstLowerer.SetType(const Node: IAstNode; const AType: IStaticType): IAstNode;
begin
    (Node as TAstNode).StaticType := AType;
    Result := Node;
end;

function TAstLowerer.Accept(const Node: IAstNode): TDataValue;
begin
    // Added for TCO
    if (not Assigned(Node)) or Done then
        exit;

    FIsTailStack.Push(FNextIsTail);
    try
        Result := inherited Accept(Node);
    finally
        FNextIsTail := FIsTailStack.Pop;
    end;
end;

function TAstLowerer.VisitMacroExpansionNode(const Node: IMacroExpansionNode): TDataValue;
begin
    // This node is just a wrapper. We must propagate the TCO context
    // to the expanded body.
    // We don't need to rebuild the node, as the evaluator just executes the body.
    Accept(Node.Callee);
    AcceptNodes<IAstNode>(Node.Arguments);

    // Propagate tail call status to the expanded body
    FNextIsTail := FIsTailStack.Peek;
    Result := Accept(Node.ExpandedBody);
end;

function TAstLowerer.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 TAstLowerer.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 TAstLowerer.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 TAstLowerer.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 TAstLowerer.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);

    // We don't modify the node, just traverse for TCO
    Result := TDataValue.FromIntf<ILambdaExpressionNode>(Node);
end;

function TAstLowerer.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
var
    calleeIdentifier: TIdentifierNode;
    binaryOp: TScalar.TBinaryOp;
    unaryOp: TScalar.TUnaryOp;
    left, right: IAstNode;
    callee: IAstNode;
    args: TArray<IAstNode>;
    isTailCall: Boolean;
begin
    // --- Transformation: Keyword-as-Function ---
    if (Node.Callee is TKeywordNode) then
    begin
        // This node should have been transformed by the TAstBinder.
        // If we see it here, it's an error, but we just traverse it.
        exit(inherited VisitFunctionCall(Node));
    end;

    // --- Optimization: Operator Folding ---
    if (Node.Callee is TIdentifierNode) then
    begin
        calleeIdentifier := Node.Callee as TIdentifierNode;

        if (Length(Node.Arguments) = 2) then
        begin
            if FBinaryOperators.TryGetValue(calleeIdentifier.Name, binaryOp) then
            begin
                FNextIsTail := False; // Ops are not tail calls
                left := Accept(Node.Arguments[0]).AsIntf<IAstNode>;
                right := Accept(Node.Arguments[1]).AsIntf<IAstNode>;
                var binExpr := TAst.BinaryExpr(left, binaryOp, right);
                Result := TDataValue.FromIntf<IAstNode>(SetType(binExpr, (Node as TAstNode).StaticType));
                exit;
            end;
        end;

        if (Length(Node.Arguments) = 1) then
        begin
            if FUnaryOperators.TryGetValue(calleeIdentifier.Name, unaryOp) then
            begin
                FNextIsTail := False; // Ops are not tail calls
                right := Accept(Node.Arguments[0]).AsIntf<IAstNode>;
                var unExpr := TAst.UnaryExpr(unaryOp, right);
                Result := TDataValue.FromIntf<IAstNode>(SetType(unExpr, (Node as TAstNode).StaticType));
                exit;
            end;

            if (calleeIdentifier.Name = '-') then
            begin
                FNextIsTail := False; // Ops are not tail calls
                right := Accept(Node.Arguments[0]).AsIntf<IAstNode>;
                var unExpr := TAst.UnaryExpr(TScalar.TUnaryOp.Negate, right);
                Result := TDataValue.FromIntf<IAstNode>(SetType(unExpr, (Node as TAstNode).StaticType));
                exit;
            end;
        end;
    end;

    // --- Standard Function Call (TCO Check) ---
    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);
        Result := TDataValue.FromIntf<IAstNode>(SetType(newBoundCall, (Node as TAstNode).StaticType));
    end
    else
        Result := TDataValue.FromIntf<IAstNode>(Node); // No change
end;

end.