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)
    protected
        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 VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue; virtual; abstract;
        function VisitQuasiquote(const Node: IQuasiquoteNode): TDataValue; virtual; abstract;
        function VisitUnquote(const Node: IUnquoteNode): TDataValue; virtual; abstract;
        function VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): 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.
    // Derived classes can override the virtual Visit... methods to customize behavior.
    TAstTransformer = class abstract(TAstVisitor, IAstTransformer)
    private
        FDone: Boolean;

    protected
        // These virtual methods implement the default identity-transformation (cloning) behavior.
        function VisitConstant(const Node: IConstantNode): TDataValue; override;
        function VisitIdentifier(const Node: IIdentifierNode): TDataValue; override;
        function VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue; override;
        function VisitUnaryExpression(const Node: IUnaryExpressionNode): TDataValue; override;
        function VisitIfExpression(const Node: IIfExpressionNode): TDataValue; override;
        function VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue; override;
        function VisitLambdaExpression(const Node: ILambdaExpressionNode): TDataValue; override;
        function VisitFunctionCall(const Node: IFunctionCallNode): TDataValue; override;
        function VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue; override;
        function VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue; override;
        function VisitAssignment(const Node: IAssignmentNode): TDataValue; override;
        function VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue; override;
        function VisitQuasiquote(const Node: IQuasiquoteNode): TDataValue; override;
        function VisitUnquote(const Node: IUnquoteNode): TDataValue; override;
        function VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): TDataValue; override;
        function VisitIndexer(const Node: IIndexerNode): TDataValue; override;
        function VisitMemberAccess(const Node: IMemberAccessNode): TDataValue; override;
        function VisitCreateSeries(const Node: ICreateSeriesNode): TDataValue; override;
        function VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue; override;
        function VisitSeriesLength(const Node: ISeriesLengthNode): TDataValue; override;
        function VisitRecurNode(const Node: IRecurNode): TDataValue; override;

        // 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;
    end;

    // Generic traverser for managing state during AST walks.
    TAstTraverser<T> = class(TAstTransformer)
    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;

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

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

function TAstTransformer.VisitBinaryExpression(const Node: IBinaryExpressionNode): TDataValue;
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 := TDataValue.FromIntf<IBinaryExpressionNode>(Node)
    else
        Result := TDataValue.FromIntf<IBinaryExpressionNode>(TAst.BinaryExpr(left, Node.Operator, right));
end;

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

function TAstTransformer.VisitIfExpression(const Node: IIfExpressionNode): TDataValue;
begin
    var condition := Accept(Node.Condition).AsIntf<IAstNode>;
    var thenBranch := Accept(Node.ThenBranch).AsIntf<IAstNode>;
    var elseBranchValue := Accept(Node.ElseBranch);
    var elseBranch :=
        if elseBranchValue.IsVoid then nil
        else elseBranchValue.AsIntf<IAstNode>;

    if (condition = Node.Condition) and (thenBranch = Node.ThenBranch) and (elseBranch = Node.ElseBranch) then
        Result := TDataValue.FromIntf<IIfExpressionNode>(Node)
    else
        Result := TDataValue.FromIntf<IIfExpressionNode>(TAst.IfExpr(condition, thenBranch, elseBranch));
end;

function TAstTransformer.VisitTernaryExpression(const Node: ITernaryExpressionNode): TDataValue;
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 := TDataValue.FromIntf<ITernaryExpressionNode>(Node)
    else
        Result := TDataValue.FromIntf<ITernaryExpressionNode>(TAst.TernaryExpr(condition, thenBranch, elseBranch));
end;

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

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

function TAstTransformer.VisitBlockExpression(const Node: IBlockExpressionNode): TDataValue;
var
    expressions: TArray<IAstNode>;
    i: Integer;
    hasChanged: Boolean;
begin
    // Check if any child expression has changed before creating a new node.
    hasChanged := false;
    SetLength(expressions, Length(Node.Expressions));
    for i := 0 to High(Node.Expressions) do
    begin
        expressions[i] := Accept(Node.Expressions[i]).AsIntf<IAstNode>;
        if expressions[i] <> Node.Expressions[i] then
            hasChanged := true;
    end;

    if hasChanged then
        Result := TDataValue.FromIntf<IBlockExpressionNode>(TAst.Block(expressions))
    else
        Result := TDataValue.FromIntf<IBlockExpressionNode>(Node);
end;

function TAstTransformer.VisitVariableDeclaration(const Node: IVariableDeclarationNode): TDataValue;
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 := TDataValue.FromIntf<IVariableDeclarationNode>(Node)
    else
        Result := TDataValue.FromIntf<IVariableDeclarationNode>(TAst.VarDecl(identifier, initializer));
end;

function TAstTransformer.VisitAssignment(const Node: IAssignmentNode): TDataValue;
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 := TDataValue.FromIntf<IAssignmentNode>(Node)
    else
        Result := TDataValue.FromIntf<IAssignmentNode>(TAst.Assign(identifier, value));
end;

function TAstTransformer.VisitMacroDefinition(const Node: IMacroDefinitionNode): TDataValue;
begin
    var name := Accept(Node.Name).AsIntf<IIdentifierNode>;
    var parameters := TransformNodes<IIdentifierNode>(Node.Parameters);
    var body := Accept(Node.Body).AsIntf<IAstNode>;
    if (name = Node.Name) and (parameters = Node.Parameters) and (body = Node.Body) then
        Result := TDataValue.FromIntf<IMacroDefinitionNode>(Node)
    else
        Result := TDataValue.FromIntf<IMacroDefinitionNode>(TAst.MacroDef(name, parameters, body));
end;

function TAstTransformer.VisitQuasiquote(const Node: IQuasiquoteNode): TDataValue;
begin
    var expr := Accept(Node.Expression).AsIntf<IAstNode>;
    if expr = Node.Expression then
        Result := TDataValue.FromIntf<IQuasiquoteNode>(Node)
    else
        Result := TDataValue.FromIntf<IQuasiquoteNode>(TAst.Quasiquote(expr));
end;

function TAstTransformer.VisitUnquote(const Node: IUnquoteNode): TDataValue;
begin
    var expr := Accept(Node.Expression).AsIntf<IAstNode>;
    if expr = Node.Expression then
        Result := TDataValue.FromIntf<IUnquoteNode>(Node)
    else
        Result := TDataValue.FromIntf<IUnquoteNode>(TAst.Unquote(expr));
end;

function TAstTransformer.VisitUnquoteSplicing(const Node: IUnquoteSplicingNode): TDataValue;
begin
    var expr := Accept(Node.Expression).AsIntf<IAstNode>;
    if expr = Node.Expression then
        Result := TDataValue.FromIntf<IUnquoteSplicingNode>(Node)
    else
        Result := TDataValue.FromIntf<IUnquoteSplicingNode>(TAst.UnquoteSplicing(expr));
end;

function TAstTransformer.VisitIndexer(const Node: IIndexerNode): TDataValue;
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 := TDataValue.FromIntf<IIndexerNode>(Node)
    else
        Result := TDataValue.FromIntf<IIndexerNode>(TAst.Indexer(base, index));
end;

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

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

function TAstTransformer.VisitAddSeriesItem(const Node: IAddSeriesItemNode): TDataValue;
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 := TDataValue.FromIntf<IAddSeriesItemNode>(Node)
    else
        Result := TDataValue.FromIntf<IAddSeriesItemNode>(TAst.AddSeriesItem(series, value, lookback));
end;

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

function TAstTransformer.VisitRecurNode(const Node: IRecurNode): TDataValue;
begin
    var args := TransformNodes<IAstNode>(Node.Arguments);
    if args = Node.Arguments then
        Result := TDataValue.FromIntf<IRecurNode>(Node)
    else
        Result := TDataValue.FromIntf<IRecurNode>(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;

{ 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.