unit Myc.Ast.Refactoring.Remove;

interface

uses
    System.SysUtils,
    System.Classes,
    System.Generics.Collections,
    Myc.Ast,
    Myc.Ast.Nodes,
    Myc.Ast.Visitor,
    Myc.Ast.Types;

type
    // Removes a specific node from the AST.
    TAstNodeRemover = class(TAstTransformer)
    private
        FTarget: IAstNode;
        FRemoved: Boolean;

        function FilterList<T: IAstNode>(const Source: INodeList<T>): TArray<T>;
        function EnsureNode(const Node: IAstNode; const ContextNode: IAstNode): IAstNode;

    strict private
        // List Containers (Filter Logic)
        function VisitBlockExpression(const Node: IAstNode): IAstNode;
        function VisitExpressionList(const Node: IAstNode): IAstNode;
        function VisitArgumentList(const Node: IAstNode): IAstNode;
        function VisitParameterList(const Node: IAstNode): IAstNode;
        function VisitRecordFieldList(const Node: IAstNode): IAstNode;

        // Structural Nodes (Replace Logic)
        function VisitIfExpression(const Node: IAstNode): IAstNode;
        function VisitCondExpression(const Node: IAstNode): IAstNode;
        function VisitAssignment(const Node: IAstNode): IAstNode;
        function VisitVariableDeclaration(const Node: IAstNode): IAstNode;
        function VisitLambdaExpression(const Node: IAstNode): IAstNode;
        function VisitFunctionCall(const Node: IAstNode): IAstNode;
        function VisitIndexer(const Node: IAstNode): IAstNode;
        function VisitMemberAccess(const Node: IAstNode): IAstNode;
        function VisitAddSeriesItem(const Node: IAstNode): IAstNode;
        function VisitTuple(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;
        function Accept(const Node: IAstNode): IAstNode; override;

    public
        constructor Create(const ATarget: IAstNode);
        class function TryRemove(const RootNode, TargetNode: IAstNode; out NewRootNode: IAstNode): Boolean; static;
    end;

implementation

{ TAstNodeRemover }

constructor TAstNodeRemover.Create(const ATarget: IAstNode);
begin
    inherited Create;
    FTarget := ATarget;
    FRemoved := False;
end;

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

    // Filter Logic
    Register(akBlockExpression, VisitBlockExpression);
    Register(akExpressionList, VisitExpressionList);
    Register(akArgumentList, VisitArgumentList);
    Register(akParameterList, VisitParameterList);
    Register(akRecordFieldList, VisitRecordFieldList);

    // Replace Logic
    Register(akIfExpression, VisitIfExpression);
    Register(akCondExpression, VisitCondExpression);
    Register(akAssignment, VisitAssignment);
    Register(akVariableDeclaration, VisitVariableDeclaration);
    Register(akLambdaExpression, VisitLambdaExpression);
    Register(akFunctionCall, VisitFunctionCall);
    Register(akIndexer, VisitIndexer);
    Register(akMemberAccess, VisitMemberAccess);
    Register(akAddSeriesItem, VisitAddSeriesItem);

    Register(akTuple, VisitTuple);
end;

class function TAstNodeRemover.TryRemove(const RootNode, TargetNode: IAstNode; out NewRootNode: IAstNode): Boolean;
var
    Remover: TAstNodeRemover;
begin
    if (RootNode = TargetNode) then
    begin
        NewRootNode := TAst.Nop(RootNode.Identity);
        Result := True;
        Exit;
    end;

    Remover := TAstNodeRemover.Create(TargetNode);
    try
        NewRootNode := Remover.Accept(RootNode);
        Result := Remover.FRemoved;

        if not Result then
            NewRootNode := RootNode
        else if (NewRootNode = nil) then
            NewRootNode := TAst.Nop(RootNode.Identity);
    finally
        Remover.Free;
    end;
end;

function TAstNodeRemover.Accept(const Node: IAstNode): IAstNode;
begin
    if Node = FTarget then
    begin
        FRemoved := True;
        Exit(nil);
    end;
    Result := inherited Accept(Node);
end;

function TAstNodeRemover.EnsureNode(const Node: IAstNode; const ContextNode: IAstNode): IAstNode;
begin
    if Assigned(Node) then
        Result := Node
    else
        Result := TAst.Nop(ContextNode.Identity);
end;

function TAstNodeRemover.FilterList<T>(const Source: INodeList<T>): TArray<T>;
var
    i: Integer;
    item: T;
    transformed: IAstNode;
    list: TList<T>;
begin
    list := TList<T>.Create;
    try
        for i := 0 to Source.Count - 1 do
        begin
            item := Source[i];
            transformed := Accept(item);
            if Assigned(transformed) then
                list.Add(T(transformed));
        end;
        Result := list.ToArray;
    finally
        list.Free;
    end;
end;

// -----------------------------------------------------------------------------
// List Visitors (Filtering)
// -----------------------------------------------------------------------------

function TAstNodeRemover.VisitBlockExpression(const Node: IAstNode): IAstNode;
var
    B: IBlockExpressionNode;
    newExprs: TArray<IAstNode>;
    newList: IExpressionList;
    changed: Boolean;
    i: Integer;
begin
    B := Node.AsBlockExpression;
    newExprs := FilterList<IAstNode>(B.Expressions);

    if Length(newExprs) = B.Expressions.Count then
    begin
        if not FRemoved then
            Exit(Node);
        changed := False;
        for i := 0 to High(newExprs) do
            if newExprs[i] <> B.Expressions[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;

    newList := TExpressionList.Create(newExprs, B.Expressions.Identity);
    Result := TAst.Block(Node.Identity, newList, B.StaticType);
end;

function TAstNodeRemover.VisitExpressionList(const Node: IAstNode): IAstNode;
var
    L: IExpressionList;
    newExprs: TArray<IAstNode>;
    changed: Boolean;
    i: Integer;
begin
    L := Node.AsExpressionList;
    newExprs := FilterList<IAstNode>(L);
    if Length(newExprs) = L.Count then
    begin
        changed := False;
        for i := 0 to High(newExprs) do
            if newExprs[i] <> L[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;
    Result := TExpressionList.Create(newExprs, Node.Identity);
end;

function TAstNodeRemover.VisitArgumentList(const Node: IAstNode): IAstNode;
var
    L: IArgumentList;
    newArgs: TArray<IAstNode>;
    changed: Boolean;
    i: Integer;
begin
    L := Node.AsArgumentList;
    newArgs := FilterList<IAstNode>(L);
    if Length(newArgs) = L.Count then
    begin
        changed := False;
        for i := 0 to High(newArgs) do
            if newArgs[i] <> L[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;
    Result := TArgumentList.Create(newArgs, Node.Identity);
end;

function TAstNodeRemover.VisitParameterList(const Node: IAstNode): IAstNode;
var
    L: IParameterList;
    newParams: TArray<IIdentifierNode>;
    changed: Boolean;
    i: Integer;
begin
    L := Node.AsParameterList;
    newParams := FilterList<IIdentifierNode>(L);
    if Length(newParams) = L.Count then
    begin
        changed := False;
        for i := 0 to High(newParams) do
            if newParams[i] <> L[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;
    Result := TParameterList.Create(newParams, Node.Identity);
end;

function TAstNodeRemover.VisitRecordFieldList(const Node: IAstNode): IAstNode;
var
    L: IRecordFieldList;
    newFields: TArray<IRecordFieldNode>;
    changed: Boolean;
    i: Integer;
begin
    L := Node.AsRecordFieldList;
    newFields := FilterList<IRecordFieldNode>(L);
    if Length(newFields) = L.Count then
    begin
        changed := False;
        for i := 0 to High(newFields) do
            if newFields[i] <> L[i] then
                changed := True;
        if not changed then
            Exit(Node);
    end;
    Result := TRecordFieldList.Create(newFields, Node.Identity);
end;

// -----------------------------------------------------------------------------
// Structural Visitors (Replacement with Nop/Void)
// -----------------------------------------------------------------------------

function TAstNodeRemover.VisitIfExpression(const Node: IAstNode): IAstNode;
var
    E: IIfExpressionNode;
    newCond, newThen, newElse: IAstNode;
begin
    E := Node.AsIfExpression;
    newCond := EnsureNode(Accept(E.Condition), Node);
    newThen := EnsureNode(Accept(E.ThenBranch), Node);
    newElse := Accept(E.ElseBranch); // Optional

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

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

    for i := 0 to High(E.Pairs) do
    begin
        newCond := EnsureNode(Accept(E.Pairs[i].Condition), Node);
        newBranch := EnsureNode(Accept(E.Pairs[i].Branch), Node);
        newPairs[i] := TCondPair.Create(newCond, newBranch);
        if (newCond <> E.Pairs[i].Condition) or (newBranch <> E.Pairs[i].Branch) then
            hasChanged := True;
    end;

    newElse := EnsureNode(Accept(E.ElseBranch), Node);
    if newElse <> E.ElseBranch then
        hasChanged := True;

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

function TAstNodeRemover.VisitAssignment(const Node: IAstNode): IAstNode;
var
    A: IAssignmentNode;
    newTarget, newValue: IAstNode;
begin
    A := Node.AsAssignment;
    newTarget := EnsureNode(Accept(A.Target), Node);
    newValue := EnsureNode(Accept(A.Value), Node);
    if (newTarget = A.Target) and (newValue = A.Value) then
        Exit(Node);
    Result := TAst.Assign(Node.Identity, newTarget, newValue, A.StaticType);
end;

function TAstNodeRemover.VisitVariableDeclaration(const Node: IAstNode): IAstNode;
var
    V: IVariableDeclarationNode;
    newTarget, newInit: IAstNode;
begin
    V := Node.AsVariableDeclaration;
    newTarget := EnsureNode(Accept(V.Target), Node);
    newInit := Accept(V.Initializer);
    if (newTarget = V.Target) and (newInit = V.Initializer) then
        Exit(Node);
    Result := TAst.VarDecl(Node.Identity, newTarget, newInit, V.StaticType, V.IsBoxed);
end;

function TAstNodeRemover.VisitLambdaExpression(const Node: IAstNode): IAstNode;
var
    L: ILambdaExpressionNode;
    newParams, newBody: IAstNode;
begin
    L := Node.AsLambdaExpression;
    newParams := Accept(L.Parameters);
    if newParams = nil then
        newParams := TParameterList.Create([], L.Parameters.Identity);
    newBody := EnsureNode(Accept(L.Body), Node);

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

function TAstNodeRemover.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    newCallee, newArgs: IAstNode;
begin
    C := Node.AsFunctionCall;
    newCallee := EnsureNode(Accept(C.Callee), Node);
    newArgs := Accept(C.Arguments);
    if newArgs = nil then
        newArgs := TArgumentList.Create([], C.Arguments.Identity);

    if (newCallee = C.Callee) and (newArgs = C.Arguments) then
        Exit(Node);
    Result :=
        TAst.FunctionCall(Node.Identity, newCallee, newArgs.AsArgumentList, C.StaticType, C.IsTailCall, C.StaticTarget, C.IsTargetPure);
end;

function TAstNodeRemover.VisitIndexer(const Node: IAstNode): IAstNode;
var
    I: IIndexerNode;
    newBase, newIndex: IAstNode;
begin
    I := Node.AsIndexer;
    newBase := EnsureNode(Accept(I.Base), Node);
    newIndex := EnsureNode(Accept(I.Index), Node);
    if (newBase = I.Base) and (newIndex = I.Index) then
        Exit(Node);
    Result := TAst.Indexer(Node.Identity, newBase, newIndex, I.StaticType);
end;

function TAstNodeRemover.VisitMemberAccess(const Node: IAstNode): IAstNode;
var
    M: IMemberAccessNode;
    newBase, newMember: IAstNode;
begin
    M := Node.AsMemberAccess;
    newBase := EnsureNode(Accept(M.Base), Node);
    newMember := Accept(M.Member);
    if newMember = nil then
        Exit(TAst.Nop(Node.Identity)); // Member missing -> invalid access

    if (newBase = M.Base) and (newMember = M.Member) then
        Exit(Node);
    Result := TAst.MemberAccess(Node.Identity, newBase, newMember.AsKeyword, M.StaticType);
end;

function TAstNodeRemover.VisitAddSeriesItem(const Node: IAstNode): IAstNode;
var
    A: IAddSeriesItemNode;
    newSeries, newValue, newLookback: IAstNode;
begin
    A := Node.AsAddSeriesItem;
    newSeries := Accept(A.Series);
    if newSeries = nil then
        Exit(TAst.Nop(Node.Identity));
    newValue := EnsureNode(Accept(A.Value), Node);
    newLookback := Accept(A.Lookback);

    if (newSeries = A.Series) and (newValue = A.Value) and (newLookback = A.Lookback) then
        Exit(Node);
    Result := TAst.AddSeriesItem(Node.Identity, newSeries.AsIdentifier, newValue, newLookback, A.StaticType);
end;

function TAstNodeRemover.VisitTuple(const Node: IAstNode): IAstNode;
var
    T: ITupleNode;
    newElements: TArray<IAstNode>;
    list: TList<IAstNode>;
    item, transformed: IAstNode;
    i: Integer;
    hasChanged: Boolean;
begin
    T := Node.AsTuple;
    list := TList<IAstNode>.Create;
    hasChanged := False;
    try
        for i := 0 to High(T.Elements) do
        begin
            item := T.Elements[i];
            transformed := Accept(item);

            if transformed <> item then
                hasChanged := True;

            if Assigned(transformed) then
                list.Add(transformed)
            else
                hasChanged := True; // Item removed
        end;

        if not hasChanged then
            Exit(Node);

        Result := TAst.Tuple(Node.Identity, list.ToArray, T.StaticType);
    finally
        list.Free;
    end;
end;

end.