unit Myc.Ast.Visitor;

interface

uses
    System.SysUtils,
    System.Generics.Collections,
    Myc.Data.Value,
    Myc.Ast,
    Myc.Ast.Nodes;

type
    // A fully abstract base class for any visitor of the AST.
    TAstVisitor = class abstract(TInterfacedObject, IAstVisitor)
    public
        function VisitConstant(const Node: IConstantNode): TDataValue; virtual; abstract;
        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; virtual; abstract;
        function VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue; virtual; abstract;
        function VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue; virtual; abstract;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; virtual; abstract;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; virtual; abstract;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; virtual; abstract;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; virtual; abstract;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; virtual; abstract;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue; virtual; abstract;
        function VisitAssignment(const Node: IAssignmentNode): TDataValue; virtual; abstract;
        function VisitIndexer(const Node: IIndexerNode): TDataValue; virtual; abstract;
        function VisitMemberAccess(const Node: IMemberAccessNode): TDataValue; virtual; abstract;
        function VisitCreateSeries(const Node: ICreateSeriesNode): TDataValue; virtual; abstract;
        function VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue; virtual; abstract;
        function VisitSeriesLength(const Node: ISeriesLengthNode): TDataValue; virtual; abstract;
        function VisitRecurNode(const Node: IRecurNode): TDataValue; virtual; abstract;
    end;

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

    // A base visitor that walks the AST and rebuilds it node by node.
    // This class uses the Template Method pattern:
    // - The Visit... methods are the template and should not be overridden.
    // - The virtual Transform... methods are the hooks for derived classes to customize behavior.
    TAstTransformer = class abstract(TAstVisitor, IAstTransformer)
    private
        FDone: Boolean;

    protected
        // Dispatch methods for each node type. Derived classes override these to change the transformation.
        function TransformConstant(const Node: IConstantNode): IConstantNode; virtual;
        function TransformIdentifier(const Node: IIdentifierNode): IIdentifierNode; virtual;
        function TransformBinaryExpression(const Node: IBinaryExpressionNode): IBinaryExpressionNode; virtual;
        function TransformUnaryExpression(const Node: IUnaryExpressionNode): IUnaryExpressionNode; virtual;
        function TransformIfExpression(const Node: IIfExpressionNode): IIfExpressionNode; virtual;
        function TransformTernaryExpression(const Node: ITernaryExpressionNode): ITernaryExpressionNode; virtual;
        function TransformLambdaExpression(const Node: ILambdaExpressionNode): ILambdaExpressionNode; virtual;
        function TransformFunctionCall(const Node: IFunctionCallNode): IFunctionCallNode; virtual;
        function TransformBlockExpression(const Node: IBlockExpressionNode): IBlockExpressionNode; virtual;
        function TransformVariableDeclaration(const Node: IVariableDeclarationNode): IVariableDeclarationNode; virtual;
        function TransformAssignment(const Node: IAssignmentNode): IAssignmentNode; virtual;
        function TransformIndexer(const Node: IIndexerNode): IIndexerNode; virtual;
        function TransformMemberAccess(const Node: IMemberAccessNode): IMemberAccessNode; virtual;
        function TransformCreateSeries(const Node: ICreateSeriesNode): ICreateSeriesNode; virtual;
        function TransformAddSeriesItem(const Node: IAddSeriesItemNode): IAddSeriesItemNode; virtual;
        function TransformSeriesLength(const Node: ISeriesLengthNode): ISeriesLengthNode; virtual;
        function TransformRecur(const Node: IRecurNode): IRecurNode; virtual;

        // Helper to transform an array of nodes.
        function TransformNodes<T: IAstNode>(const Nodes: TArray<T>): TArray<T>;

        function Accept(const Node: IAstNode): TDataValue; virtual;

        property Done: Boolean read FDone write FDone;
    public
        // This is the main entry point for the transformer, returning a transformed node.
        function Execute(const RootNode: IAstNode): IAstNode;

        // Final Visit methods - these should not be overridden. They call the virtual Transform methods.
        function VisitConstant(const Node: IConstantNode): TDataValue; override; final;
        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; override; final;
        function VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue; override; final;
        function VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue; override; final;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; override; final;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; override; final;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; override; final;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override; final;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override; final;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue; override; final;
        function VisitAssignment(const Node: IAssignmentNode): TDataValue; override; final;
        function VisitIndexer(const Node: IIndexerNode): TDataValue; override; final;
        function VisitMemberAccess(const Node: IMemberAccessNode): TDataValue; override; final;
        function VisitCreateSeries(const Node: ICreateSeriesNode): TDataValue; override; final;
        function VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue; override; final;
        function VisitSeriesLength(const Node: ISeriesLengthNode): TDataValue; override; final;
        function VisitRecurNode(const Node: IRecurNode): TDataValue; override; final;
    end;

    // TAstTraverser simply inherits the identity-transform (cloning) behavior from TAstTransformer.
    TAstTraverser = class(TAstTransformer);

    // Generic traverser for managing state during AST walks.
    TAstTraverser<T> = class(TAstTraverser)
    private
        FData: TStack<T>;
    protected
        function Accept(const Node: IAstNode): TDataValue; override;
        // Called before a node is visited. The returned value is pushed onto the state stack.
        function EnterNode(const Node: IAstNode): T; virtual; abstract;
        // Called after a node has been visited.
        procedure ExitNode(const Node: IAstNode; const Data: T); virtual;
        property Data: TStack<T> read FData;
    public
        constructor Create;
        destructor Destroy; override;
    end;

implementation

{ TAstTransformer }

function TAstTransformer.Accept(const Node: IAstNode): TDataValue;
begin
    if (not Assigned(Node)) or FDone then
        exit;
    Result := Node.Accept(Self);
end;

function TAstTransformer.Execute(const RootNode: IAstNode): IAstNode;
begin
    Result := Accept(RootNode).AsIntf<IAstNode>;
end;

// --- Default implementations for Transform... methods (Identity Transformation / Cloning) ---

function TAstTransformer.TransformConstant(const Node: IConstantNode): IConstantNode;
begin
    // The node itself is immutable, return it directly.
    Result := Node;
end;

function TAstTransformer.TransformIdentifier(const Node: IIdentifierNode): IIdentifierNode;
begin
    Result := Node;
end;

function TAstTransformer.TransformBinaryExpression(const Node: IBinaryExpressionNode): IBinaryExpressionNode;
begin
    var left := Accept(Node.Left).AsIntf<IAstNode>;
    var right := Accept(Node.Right).AsIntf<IAstNode>;
    if (left = Node.Left) and (right = Node.Right) then
        Result := Node
    else
        Result := TAst.BinaryExpr(left, Node.Operator, right);
end;

function TAstTransformer.TransformUnaryExpression(const Node: IUnaryExpressionNode): IUnaryExpressionNode;
begin
    var right := Accept(Node.Right).AsIntf<IAstNode>;
    if right = Node.Right then
        Result := Node
    else
        Result := TAst.UnaryExpr(Node.Operator, right);
end;

function TAstTransformer.TransformIfExpression(const Node: IIfExpressionNode): IIfExpressionNode;
begin
    var condition := Accept(Node.Condition).AsIntf<IAstNode>;
    var thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    var elseBranch := Accept(Node.ElseBranch).AsIntf<IAstNode>;
    if (condition = Node.Condition) and (thenBranch = Node.ThenBranch) and (elseBranch = Node.ElseBranch) then
        Result := Node
    else
        Result := TAst.IfExpr(condition, thenBranch, elseBranch);
end;

function TAstTransformer.TransformTernaryExpression(const Node: ITernaryExpressionNode): ITernaryExpressionNode;
begin
    var condition := Accept(Node.Condition).AsIntf<IAstNode>;
    var thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    var elseBranch := Accept(Node.ElseBranch).AsIntf<IAstNode>;
    if (condition = Node.Condition) and (thenBranch = Node.ThenBranch) and (elseBranch = Node.ElseBranch) then
        Result := Node
    else
        Result := TAst.TernaryExpr(condition, thenBranch, elseBranch);
end;

function TAstTransformer.TransformLambdaExpression(const Node: ILambdaExpressionNode): ILambdaExpressionNode;
begin
    var parameters := TransformNodes<IIdentifierNode>(Node.Parameters);
    var body := Accept(Node.Body).AsIntf<IAstNode>;
    if (parameters = Node.Parameters) and (body = Node.Body) then
        Result := Node
    else
        Result := TAst.LambdaExpr(parameters, body);
end;

function TAstTransformer.TransformFunctionCall(const Node: IFunctionCallNode): IFunctionCallNode;
begin
    var callee := Accept(Node.Callee).AsIntf<IAstNode>;
    var args := TransformNodes<IAstNode>(Node.Arguments);
    if (callee = Node.Callee) and (args = Node.Arguments) then
        Result := Node
    else
        Result := TAst.FunctionCall(callee, args);
end;

function TAstTransformer.TransformBlockExpression(const Node: IBlockExpressionNode): IBlockExpressionNode;
begin
    // TList is mutable, so we always transform.
    var exprs: TArray<IAstNode>;
    var i: Integer;
    SetLength(exprs, Node.Expressions.Count);
    for i := 0 to Node.Expressions.Count - 1 do
        exprs[i] := Accept(Node.Expressions[i]).AsIntf<IAstNode>;
    Result := TAst.Block(exprs);
end;

function TAstTransformer.TransformVariableDeclaration(const Node: IVariableDeclarationNode): IVariableDeclarationNode;
begin
    var identifier := Accept(Node.Identifier).AsIntf<IIdentifierNode>;
    var initializer := Accept(Node.Initializer).AsIntf<IAstNode>;
    if (identifier = Node.Identifier) and (initializer = Node.Initializer) then
        Result := Node
    else
        Result := TAst.VarDecl(identifier, initializer);
end;

function TAstTransformer.TransformAssignment(const Node: IAssignmentNode): IAssignmentNode;
begin
    var identifier := Accept(Node.Identifier).AsIntf<IIdentifierNode>;
    var value := Accept(Node.Value).AsIntf<IAstNode>;
    if (identifier = Node.Identifier) and (value = Node.Value) then
        Result := Node
    else
        Result := TAst.Assign(identifier, value);
end;

function TAstTransformer.TransformIndexer(const Node: IIndexerNode): IIndexerNode;
begin
    var base := Accept(Node.Base).AsIntf<IAstNode>;
    var index := Accept(Node.Index).AsIntf<IAstNode>;
    if (base = Node.Base) and (index = Node.Index) then
        Result := Node
    else
        Result := TAst.Indexer(base, index);
end;

function TAstTransformer.TransformMemberAccess(const Node: IMemberAccessNode): IMemberAccessNode;
begin
    var base := Accept(Node.Base).AsIntf<IAstNode>;
    var member := Node.Member;
    if (base = Node.Base) and (member = Node.Member) then
        Result := Node
    else
        Result := TAst.MemberAccess(base, member);
end;

function TAstTransformer.TransformCreateSeries(const Node: ICreateSeriesNode): ICreateSeriesNode;
begin
    Result := Node;
end;

function TAstTransformer.TransformAddSeriesItem(const Node: IAddSeriesItemNode): IAddSeriesItemNode;
begin
    var series := Accept(Node.Series).AsIntf<IIdentifierNode>;
    var value := Accept(Node.Value).AsIntf<IAstNode>;
    var lookback := Accept(Node.Lookback).AsIntf<IAstNode>;
    if (series = Node.Series) and (value = Node.Value) and (lookback = Node.Lookback) then
        Result := Node
    else
        Result := TAst.AddSeriesItem(series, value, lookback);
end;

function TAstTransformer.TransformSeriesLength(const Node: ISeriesLengthNode): ISeriesLengthNode;
begin
    var series := Accept(Node.Series).AsIntf<IIdentifierNode>;
    if series = Node.Series then
        Result := Node
    else
        Result := TAst.SeriesLength(series);
end;

function TAstTransformer.TransformRecur(const Node: IRecurNode): IRecurNode;
begin
    var args := TransformNodes<IAstNode>(Node.Arguments);
    if args = Node.Arguments then
        Result := Node
    else
        Result := TAst.Recur(args);
end;

function TAstTransformer.TransformNodes<T>(const Nodes: TArray<T>): TArray<T>;
var
    i: Integer;
    hasChanged: boolean;
begin
    hasChanged := False;
    SetLength(Result, Length(Nodes));
    for i := 0 to High(Nodes) do
    begin
        Result[i] := Accept(Nodes[i]).AsIntf<T>;
        if Result[i] <> Nodes[i] then
            hasChanged := True;
    end;
    if not hasChanged then
        Result := Nodes;
end;

function TAstTransformer.VisitConstant(const Node: IConstantNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IConstantNode>(TransformConstant(Node));
end;

function TAstTransformer.VisitIdentifier(const Node: IIdentifierNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IIdentifierNode>(TransformIdentifier(Node));
end;

function TAstTransformer.VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IBinaryExpressionNode>(TransformBinaryExpression(Node));
end;

function TAstTransformer.VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IUnaryExpressionNode>(TransformUnaryExpression(Node));
end;

function TAstTransformer.VisitIfExpression(const Node: IIfExpressionNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IIfExpressionNode>(TransformIfExpression(Node));
end;

function TAstTransformer.VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue;
begin
    Result := TDataValue.FromIntf<ITernaryExpressionNode>(TransformTernaryExpression(Node));
end;

function TAstTransformer.VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue;
begin
    Result := TDataValue.FromIntf<ILambdaExpressionNode>(TransformLambdaExpression(Node));
end;

function TAstTransformer.VisitFunctionCall(const Node: IFunctionCallNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IFunctionCallNode>(TransformFunctionCall(Node));
end;

function TAstTransformer.VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IBlockExpressionNode>(TransformBlockExpression(Node));
end;

function TAstTransformer.VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IVariableDeclarationNode>(TransformVariableDeclaration(Node));
end;

function TAstTransformer.VisitAssignment(const Node: IAssignmentNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IAssignmentNode>(TransformAssignment(Node));
end;

function TAstTransformer.VisitIndexer(const Node: IIndexerNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IIndexerNode>(TransformIndexer(Node));
end;

function TAstTransformer.VisitMemberAccess(const Node: IMemberAccessNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IMemberAccessNode>(TransformMemberAccess(Node));
end;

function TAstTransformer.VisitCreateSeries(const Node: ICreateSeriesNode): TDataValue;
begin
    Result := TDataValue.FromIntf<ICreateSeriesNode>(TransformCreateSeries(Node));
end;

function TAstTransformer.VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IAddSeriesItemNode>(TransformAddSeriesItem(Node));
end;

function TAstTransformer.VisitSeriesLength(const Node: ISeriesLengthNode): TDataValue;
begin
    Result := TDataValue.FromIntf<ISeriesLengthNode>(TransformSeriesLength(Node));
end;

function TAstTransformer.VisitRecurNode(const Node: IRecurNode): TDataValue;
begin
    Result := TDataValue.FromIntf<IRecurNode>(TransformRecur(Node));
end;

{ TAstTraverser<T> }

constructor TAstTraverser<T>.Create;
begin
    inherited Create;
    FData := TStack<T>.Create;
end;

destructor TAstTraverser<T>.Destroy;
begin
    FData.Free;
    inherited;
end;

function TAstTraverser<T>.Accept(const Node: IAstNode): TDataValue;
begin
    if not Assigned(Node) then
        exit;

    FData.Push(EnterNode(Node));
    try
        // Call the inherited Accept, which will dispatch to the virtual Visit... methods
        // in the base TAstTransformer, performing the transformation/cloning.
        Result := inherited Accept(Node);
    finally
        var data := FData.Pop;
        ExitNode(Node, data);
    end;
end;

procedure TAstTraverser<T>.ExitNode(const Node: IAstNode; const Data: T);
begin
end;

end.