# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E501, F841
import pytest

import tvm
import tvm.testing
from tvm import relax, tirx
from tvm.ir import Call, Op
from tvm.ir.base import assert_structural_equal
from tvm.relax import PyExprMutator, PyExprVisitor
from tvm.relax.expr import (
    BindingBlock,
    Constant,
    DataflowBlock,
    DataflowVar,
    DataTypeImm,
    Expr,
    ExternFunc,
    Function,
    GlobalVar,
    If,
    MatchCast,
    SeqExpr,
    ShapeExpr,
    StringImm,
    Tuple,
    TupleGetItem,
    Var,
    VarBinding,
)
from tvm.script import relax as R

m, n = tirx.Var("m", "int64"), tirx.Var("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
y = relax.Var("y", R.Tensor([m, n], "float32"))
bb = relax.BlockBuilder()


@relax.expr_functor.visitor
class BasicVisitor(PyExprVisitor):
    """Default ExprVisitor"""


class ASTLog:
    """Helper class to log AST"""

    def __init__(self) -> None:
        self.log = []
        self.indent = "\t"
        self.level = 0

    def push_scope(self):
        self.level += 1

    def pop_scope(self):
        self.level -= 1

    def add(self, s: str):
        self.log.append(self.indent * self.level + s)

    def __str__(self) -> str:
        return "\n".join(self.log)


@relax.expr_functor.visitor
class ASTPrinter(PyExprVisitor):
    """Print relax AST in structured format. The shape of Node is ignored."""

    def __init__(self) -> None:
        super().__init__()
        self.log = ASTLog()

    def visit_constant_(self, op: Constant) -> None:
        self.log.add("Constant")

    def visit_global_var_(self, op: GlobalVar) -> None:
        self.log.add("GlobalVar")

    def visit_tuple_(self, op: Tuple) -> None:
        self.log.add("Tuple")
        self.log.push_scope()
        for field in op.fields:
            self.visit_expr(field)
        self.log.pop_scope()

    def visit_var_(self, op: Var) -> None:
        self.log.add("Var")

    def visit_dataflow_var_(self, op: DataflowVar) -> None:
        self.log.add("DataflowVar")

    def visit_function_(self, op: Function) -> None:
        self.log.add("Function")
        self.log.push_scope()
        for param in op.params:
            self.visit_var_def(param)

        self.visit_expr(op.body)
        self.log.pop_scope()

    def visit_call_(self, op: Call) -> None:
        self.log.add("Call")
        self.log.push_scope()
        self.visit_expr(op.op)

        for arg in op.args:
            self.visit_expr(arg)
        self.log.pop_scope()

    def visit_if_(self, op: If) -> None:
        self.log.add("If")
        self.log.push_scope()
        self.visit_expr(op.cond)
        self.visit_expr(op.true_branch)
        self.visit_expr(op.false_branch)
        self.log.pop_scope()

    def visit_op_(self, op: Op) -> None:
        self.log.add("Op")

    def visit_tuple_getitem_(self, op: TupleGetItem) -> None:
        self.log.add("TupleGetItem")
        self.log.push_scope()
        self.visit_expr(op.tuple_value)
        self.log.pop_scope()

    def visit_expr_fallback_(self, op: Expr) -> None:
        self.log.add("ExprFallback")

    def visit_string_imm_(self, op: StringImm) -> None:
        self.log.add("StringImm")

    def visit_data_type_imm_(self, op: DataTypeImm) -> None:
        self.log.add("DataTypeImm")

    def visit_shape_expr_(self, op: ShapeExpr) -> None:
        self.log.add("ShapeExpr")

    def visit_extern_func_(self, op: ExternFunc) -> None:
        self.log.add("ExternFunc")

    def visit_seq_expr_(self, op: SeqExpr) -> None:
        self.log.add("SeqExpr")
        self.log.push_scope()
        for block in op.blocks:
            self.visit_binding_block(block)
        self.visit_expr(op.body)
        self.log.pop_scope()

    def visit_var_binding_(self, binding: VarBinding) -> None:
        self.log.add("VarBinding")
        self.log.push_scope()
        self.visit_expr(binding.value)
        self.visit_var_def(binding.var)
        self.log.pop_scope()

    def visit_match_cast_(self, binding: MatchCast) -> None:
        self.log.add("MatchCast")
        self.log.push_scope()
        self.visit_var_def(binding.var)
        self.visit_expr(binding.value)
        self.log.pop_scope()

    def visit_binding_block_(self, block: BindingBlock) -> None:
        self.log.add("BindingBlock")
        self.log.push_scope()
        for binding in block.bindings:
            self.visit_binding(binding)
        self.log.pop_scope()

    def visit_dataflow_block_(self, block: DataflowBlock) -> None:
        self.log.add("DataflowBlock")
        self.log.push_scope()
        for binding in block.bindings:
            self.visit_binding(binding)
        self.log.pop_scope()

    def visit_var_def_(self, var: Var) -> None:
        self.log.add("VarDef")

    def visit_dataflow_var_def_(self, var: DataflowVar) -> None:
        self.log.add("DataflowVarDef")


@relax.expr_functor.mutator
class BasicMutator(PyExprMutator):
    """Default ExprMutator"""


@relax.expr_functor.mutator
class ASTPostPrinterMutator(PyExprMutator):
    """Print relax AST in the post order format."""

    def __init__(self) -> None:
        super().__init__()
        self.log = ASTLog()

    def visit_constant_(self, op: Constant) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("Constant")
        return op

    def visit_global_var_(self, op: GlobalVar) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("GlobalVar")
        return op

    def visit_tuple_(self, op: Tuple) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("Tuple")
        return op

    def visit_var_(self, op: Var) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("Var")
        return op

    def visit_dataflow_var_(self, op: DataflowVar) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("DataflowVar")
        return op

    def visit_function_(self, op: Function) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("Function")
        return op

    def visit_call_(self, op: Call) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("Call")
        return op

    def visit_if_(self, op: If) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("If")
        return op

    def visit_op_(self, op: Op) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("Op")
        return op

    def visit_tuple_getitem_(self, op: TupleGetItem) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("TupleGetItem")
        return op

    def visit_expr_fallback_(self, op: Expr) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("ExprFallback")
        return op

    def visit_string_imm_(self, op: StringImm) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("StringImm")
        return op

    def visit_data_type_imm_(self, op: DataTypeImm) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("DataTypeImm")
        return op

    def visit_shape_expr_(self, op: ShapeExpr) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("ShapeExpr")
        return op

    def visit_extern_func_(self, op: ExternFunc) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("ExternFunc")
        return op

    def visit_seq_expr_(self, op: SeqExpr) -> Expr:
        op = self.visit_expr_post_order(op)
        self.log.add("SeqExpr")
        return op

    def visit_var_binding_(self, binding: VarBinding) -> None:
        """Identical with ExprMutator::VisitBinding_(const VarBindingNode* binding) on the C++ side."""
        new_value = self.visit_expr(binding.value)
        new_var = self.visit_var_def(binding.var)

        self.log.add("VarBinding")
        if binding.var.same_as(new_var) and binding.value.same_as(new_value):
            self.builder_.emit_normalized(binding)
            return

        temp = self.with_type(new_var, new_value.ty)
        if not temp.same_as(new_var):
            new_var = temp
            self.set_var_remap(binding.var, new_var)

        self.builder_.emit_normalized(VarBinding(new_var, new_value))

    def visit_match_cast_(self, binding: MatchCast) -> None:
        """Identical with ExprMutator::VisitBinding_(const MatchCastNode* binding) on the C++ side."""
        new_var = self.visit_var_def(binding.var)
        new_value = self.visit_expr(binding.value)

        temp = self.with_type(new_var, binding.ty)
        if not temp.same_as(new_var):
            new_var = temp
            self.set_var_remap(binding.var, new_var)

        self.log.add("MatchCast")
        self.builder_.emit_normalized(MatchCast(new_var, new_value, binding.ty))

    def visit_binding_block_(self, block: BindingBlock) -> BindingBlock:
        """Identical with ExprMutator::VisitBindingBlock_(const BindingBlockNode* block) on the C++ side."""
        self.builder_._begin_binding_block()
        for binding in block.bindings:
            self.visit_binding(binding)
        self.log.add("BindingBlock")
        return self.builder_._end_block()

    def visit_dataflow_block_(self, block: DataflowBlock) -> None:
        """Identical with ExprMutator::VisitBindingBlock_(const DataflowBlockNode* block) on the C++ side."""
        self.builder_._begin_dataflow_block()
        for binding in block.bindings:
            self.visit_binding(binding)
        self.log.add("DataflowBlock")
        return self.builder_._end_block()

    def visit_var_def_(self, var: Var) -> None:
        """Identical with ExprMutator::VisitVarDef_(const VarNode* var) on the C++ side."""
        self.log.add("VarDef")
        return var

    def visit_dataflow_var_def_(self, var: DataflowVar) -> None:
        """Identical with ExprMutator::VisitVarDef_(const DataflowVarNode* var) on the C++ side."""
        self.log.add("DataflowVarDef")
        return var


def basic_check(expr, visitor_str, mutator_str):
    def visit(f, expr):
        if isinstance(expr, relax.Expr):
            return f.visit_expr(expr)
        elif isinstance(expr, relax.BindingBlock):
            return f.visit_binding_block(expr)

    # check no overloading case
    basic_visitor = BasicVisitor()
    visit(basic_visitor, expr)

    # check the output log
    log_visitor = ASTPrinter()
    visit(log_visitor, expr)
    assert str(log_visitor.log) == visitor_str

    # check no overloading case
    basic_mutator = BasicMutator()
    # skip normalize GlobalVar since it requires context IRModule to get ty
    if isinstance(expr, relax.Expr) and not isinstance(expr, relax.GlobalVar):
        expr = bb.normalize(expr)
        assert_structural_equal(visit(basic_mutator, expr), expr)

    # check the output log and return value
    post_log_mutator = ASTPostPrinterMutator()
    if isinstance(expr, relax.Expr) and not isinstance(expr, relax.GlobalVar):
        expr = bb.normalize(expr)
        assert_structural_equal(visit(post_log_mutator, expr), expr)
        assert str(post_log_mutator.log) == mutator_str


def test_constant():
    basic_check(relax.const(1.0), "Constant", "Constant")


def test_var():
    basic_check(x, "Var", "Var")


def test_dataflow_var():
    lv = relax.DataflowVar("lv", R.Tensor([n], "float32"))
    basic_check(lv, "DataflowVar", "DataflowVar")


def test_tuple():
    t = relax.Tuple([x, y])
    basic_check(t, "\n".join(["Tuple", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Tuple"]))


def test_global_var():
    gv = relax.GlobalVar("gv")
    basic_check(gv, "GlobalVar", "GlobalVar")


def test_seq_expr():
    bindings = [relax.VarBinding(x, relax.const(1))]
    blocks = [relax.BindingBlock(bindings)]
    seq_expr = relax.SeqExpr(blocks, x)
    basic_check(
        seq_expr,
        "\n".join(
            [
                "SeqExpr",
                "\tBindingBlock",
                "\t\tVarBinding",
                "\t\t\tConstant",
                "\t\t\tVarDef",
                "\tVar",
            ]
        ),
        "\n".join(["Constant", "VarDef", "VarBinding", "BindingBlock", "Var", "SeqExpr"]),
    )


def test_shape_expr():
    x = relax.ShapeExpr([m, n])
    basic_check(x, "ShapeExpr", "\n".join(["ExprFallback", "ExprFallback", "ShapeExpr"]))


def test_prim_expr():
    basic_check(tvm.tirx.IntImm("int64", 1), "ExprFallback", "ExprFallback")
    basic_check(tvm.tirx.FloatImm("float32", 1.0), "ExprFallback", "ExprFallback")
    basic_check(m + n, "ExprFallback", "ExprFallback")
    basic_check(
        tvm.tirx.call_extern("int32", "test"),
        "\n".join(["Call", "\tOp", "\tExprFallback"]),
        "\n".join(["Op", "ExprFallback", "Call"]),
    )


def test_call():
    call_node = relax.op.add(x, y)
    basic_check(
        call_node,
        "\n".join(["Call", "\tOp", "\tVar", "\tVar"]),
        "\n".join(["Op", "Var", "Var", "ExprFallback", "ExprFallback", "ShapeExpr", "Call"]),
    )


def test_if():
    if_node = relax.If(x, x, x)
    basic_check(
        if_node,
        "\n".join(["If", "\tVar", "\tSeqExpr", "\t\tVar", "\tSeqExpr", "\t\tVar"]),
        "\n".join(["Var", "Var", "SeqExpr", "Var", "SeqExpr", "If"]),
    )


def test_tuple_getitem():
    tuple_getitem_node = relax.TupleGetItem(relax.Tuple([x, y]), 0)
    basic_check(
        tuple_getitem_node,
        "\n".join(["TupleGetItem", "\tTuple", "\t\tVar", "\t\tVar"]),
        "\n".join(["Var", "Var", "Tuple", "TupleGetItem"]),
    )


def test_binding_block():
    bb._begin_binding_block()
    gv0 = bb.emit(relax.op.add(x, y))
    gv1 = bb.match_cast(y, R.Tensor([m, n], "float32"))
    b0 = bb._end_block()
    basic_check(
        b0,
        "\n".join(
            [
                "BindingBlock",
                "\tVarBinding",
                "\t\tCall",
                "\t\t\tOp",
                "\t\t\tVar",
                "\t\t\tVar",
                "\t\tVarDef",
                "\tMatchCast",
                "\t\tVarDef",
                "\t\tVar",
            ]
        ),
        "\n".join(
            [
                "Op",
                "Var",
                "Var",
                "Call",
                "ShapeExpr",
                "VarDef",
                "VarBinding",
                "Var",
                "ShapeExpr",
                "ShapeExpr",
                "VarDef",
                "MatchCast",
                "BindingBlock",
            ]
        ),
    )


def test_dataflow_block():
    bb._begin_dataflow_block()
    lv0 = bb.emit(relax.op.add(x, y))
    gv1 = bb.match_cast(y, R.Tensor([m, n], "float32"))
    b0 = bb._end_block()
    basic_check(
        b0,
        "\n".join(
            [
                "DataflowBlock",
                "\tVarBinding",
                "\t\tCall",
                "\t\t\tOp",
                "\t\t\tVar",
                "\t\t\tVar",
                "\t\tDataflowVarDef",
                "\tMatchCast",
                "\t\tDataflowVarDef",
                "\t\tVar",
            ]
        ),
        "\n".join(
            [
                "Op",
                "Var",
                "Var",
                "Call",
                "ShapeExpr",
                "DataflowVarDef",
                "VarBinding",
                "Var",
                "ShapeExpr",
                "ShapeExpr",
                "DataflowVarDef",
                "MatchCast",
                "DataflowBlock",
            ]
        ),
    )


def test_function():
    bindings = [relax.VarBinding(x, relax.const(1))]
    blocks = [relax.BindingBlock(bindings)]
    seq_expr = relax.SeqExpr(blocks, x)
    func = relax.Function([x], seq_expr, R.Tensor([n], "float32"))
    basic_check(
        func,
        "\n".join(
            [
                "Function",
                "\tVarDef",
                "\tSeqExpr",
                "\t\tBindingBlock",
                "\t\t\tVarBinding",
                "\t\t\t\tConstant",
                "\t\t\t\tVarDef",
                "\t\tVar",
            ]
        ),
        "\n".join(
            [
                "VarDef",
                "Constant",
                "VarDef",
                "VarBinding",
                "BindingBlock",
                "Var",
                "SeqExpr",
                "Function",
            ]
        ),
    )


def test_extern_func():
    func = relax.ExternFunc("f")
    basic_check(func, "ExternFunc", "ExternFunc")


def test_inherit():
    # The internal class is not instantiated.
    class InternalVisitor(PyExprVisitor):
        def __init__(self) -> None:
            super().__init__()
            self.log = ASTLog()

        def visit_call_(self, op: Call) -> None:
            self.log.add("InternalCall")
            self.log.push_scope()
            self.visit_expr(op.op)

            for arg in op.args:
                self.visit_expr(arg)
            self.log.pop_scope()

        def visit_var_(self, op: Var) -> None:
            self.log.add("Var")

        def visit_op_(self, op: Op) -> None:
            self.log.add("Op")

    @relax.expr_functor.visitor
    class LeafVisitor(InternalVisitor):
        def visit_call_(self, op: Call) -> None:
            self.log.add("LeafCall")
            self.log.push_scope()
            self.visit_expr(op.op)

            for arg in op.args:
                self.visit_expr(arg)
            self.log.pop_scope()

    call_node = relax.op.add(x, y)
    lv = LeafVisitor()
    lv.visit_expr(call_node)
    assert str(lv.log) == "\n".join(["LeafCall", "\tOp", "\tVar", "\tVar"])


def test_inherit_with_cls():
    # The decorator converts `InternalVisitor` to a wrapper class.
    @relax.expr_functor.visitor
    class InternalVisitor(PyExprVisitor):
        def __init__(self) -> None:
            super().__init__()
            self.log = ASTLog()

        def visit_call_(self, op: Call) -> None:
            self.log.add("InternalCall")
            self.log.push_scope()
            self.visit_expr(op.op)

            for arg in op.args:
                self.visit_expr(arg)
            self.log.pop_scope()

        def visit_var_(self, op: Var) -> None:
            self.log.add("Var")

        def visit_op_(self, op: Op) -> None:
            self.log.add("Op")

    # `InternalVisitor._cls` refers to the original `InternalVisitor` users defined.
    @relax.expr_functor.visitor
    class LeafVisitor(InternalVisitor._cls):
        def visit_call_(self, op: Call) -> None:
            self.log.add("LeafCall")
            self.log.push_scope()
            self.visit_expr(op.op)

            for arg in op.args:
                self.visit_expr(arg)
            self.log.pop_scope()

    call_node = relax.op.add(x, y)
    iv = InternalVisitor()
    iv.visit_expr(call_node)
    assert str(iv.log) == "\n".join(["InternalCall", "\tOp", "\tVar", "\tVar"])

    lv = LeafVisitor()
    lv.visit_expr(call_node)
    assert str(lv.log) == "\n".join(["LeafCall", "\tOp", "\tVar", "\tVar"])


def test_wrong_inherit():
    @relax.expr_functor.visitor
    class InternalVisitor(PyExprVisitor):
        def visit_call_(self, op: Call) -> None:
            pass

    with pytest.raises(
        TypeError,
        match="Inheritance from a decorated object `LeafVisitor` is not allowed. Please inherit from `LeafVisitor._cls`.",
    ):

        @relax.expr_functor.visitor
        class LeafVisitor(InternalVisitor):
            def visit_call_(self, op: Call) -> None:
                pass


@R.function
def dummy(x: R.Tensor((10, 10))):
    lv = R.add(x, R.const(1))
    with R.dataflow():
        gv = lv
        R.output(gv)
    return gv


def test_call_visitor_super():
    @relax.expr_functor.visitor
    class InternalVisitor(PyExprVisitor):
        def __init__(self) -> None:
            super().__init__()
            self.log = ASTLog()

        def visit_binding_block_(self, block: relax.BindingBlock) -> None:
            self.log.add("BindingBlock")
            super().visit_binding_block_(block)

        def visit_dataflow_block_(self, block: DataflowBlock) -> None:
            self.log.add("DataflowBlock")
            super().visit_dataflow_block_(block)

        def visit_var_binding_(self, binding: relax.VarBinding) -> None:
            self.log.add("VarBinding")
            super().visit_var_binding_(binding)

        def visit_call_(self, op: Call) -> None:
            self.log.add("InternalCall")
            super().visit_call_(op)  # call PyExprVisitor.visit_call_

        def visit_var_def_(self, var: Var) -> None:
            self.log.add("VarDef")
            super().visit_var_def_(var)

        def visit_dataflow_var_def_(self, var: Var) -> None:
            self.log.add("DataflowVarDef")
            super().visit_dataflow_var_def_(var)

        def visit_var_(self, op: Var) -> None:
            self.log.add("Var")

        def visit_op_(self, op: Op) -> None:
            self.log.add("Op")

    @relax.expr_functor.visitor
    class LeafVisitor(InternalVisitor._cls):
        def visit_call_(self, op: Call) -> None:
            self.log.add("LeafCall")
            super().visit_call_(op)  # call InternalVisit.visit_call_

    call_node = relax.op.add(x, y)
    iv = InternalVisitor()
    iv.visit_expr(call_node)
    assert str(iv.log) == "\n".join(["InternalCall", "Op", "Var", "Var"])

    lv = LeafVisitor()
    lv.visit_expr(call_node)
    assert str(lv.log) == "\n".join(["LeafCall", "InternalCall", "Op", "Var", "Var"])

    lv = LeafVisitor()
    lv.visit_expr(dummy)
    assert str(lv.log) == "\n".join(
        [
            "VarDef",
            "BindingBlock",
            "VarBinding",
            "LeafCall",
            "InternalCall",
            "Op",
            "Var",
            "VarDef",
            "DataflowBlock",
            "VarBinding",
            "Var",
            "VarDef",
            "Var",
        ]
    )


def test_call_mutator_super():
    @relax.expr_functor.mutator
    class InternalMutator(PyExprMutator):
        def __init__(self) -> None:
            super().__init__()
            self.log = ASTLog()

        def visit_binding_block_(self, block: relax.BindingBlock) -> None:
            self.log.add("BindingBlock")
            return super().visit_binding_block_(block)

        def visit_dataflow_block_(self, block: DataflowBlock) -> None:
            self.log.add("DataflowBlock")
            return super().visit_dataflow_block_(block)

        def visit_var_binding_(self, binding: relax.VarBinding) -> None:
            self.log.add("VarBinding")
            return super().visit_var_binding_(binding)

        def visit_call_(self, op: Call) -> None:
            self.log.add("InternalCall")
            return super().visit_call_(op)  # call PyExprMutator.visit_call_

        def visit_var_def_(self, var: Var) -> None:
            self.log.add("VarDef")
            return super().visit_var_def_(var)

        def visit_dataflow_var_def_(self, var: Var) -> None:
            self.log.add("DataflowVarDef")
            return super().visit_dataflow_var_def_(var)

        def visit_var_(self, op: Var) -> None:
            self.log.add("Var")
            return super().visit_var_(op)  # call PyExprMutator.visit_var_

        def visit_op_(self, op: Op) -> None:
            self.log.add("Op")
            return super().visit_op_(op)  # call PyExprMutator.visit_op_

    @relax.expr_functor.mutator
    class LeafMutator(InternalMutator._cls):
        def visit_call_(self, op: Call) -> None:
            self.log.add("LeafCall")
            return super().visit_call_(op)  # call InternalMutator.visit_call_

    call_node = relax.op.add(x, y)
    im = InternalMutator()
    im.visit_expr(call_node)
    assert str(im.log) == "\n".join(["InternalCall", "Op", "Var", "Var"])

    lm = LeafMutator()
    lm.visit_expr(call_node)
    assert str(lm.log) == "\n".join(["LeafCall", "InternalCall", "Op", "Var", "Var"])

    lm = LeafMutator()
    lm.visit_expr(dummy)
    assert str(lm.log) == "\n".join(
        [
            "VarDef",
            "BindingBlock",
            "VarBinding",
            "LeafCall",
            "InternalCall",
            "Op",
            "Var",
            "VarDef",
            "DataflowBlock",
            "VarBinding",
            "Var",
            "VarDef",
            "Var",
        ]
    )


def test_function_parameter_mutation():
    @relax.expr_functor.mutator
    class ParamMutator(PyExprMutator):
        def __init__(self, shape_replacements):
            super().__init__()
            self.shape_replacements = shape_replacements

        def visit_var_def_(self, var):
            if var.name_hint in self.shape_replacements:
                new_shape = self.shape_replacements[var.name_hint]
                new_ty = relax.TensorType(new_shape, dtype=var.ty.dtype)
                return relax.Var(f"{var.name_hint}_with_new_shape", new_ty)
            else:
                return var

    @R.function(private=True)
    def before(A: R.Tensor((16, 32), "float32"), B: R.Tensor((32, 64), "float32")) -> R.Tensor(
        (16, 64), "float32"
    ):
        return R.matmul(A, B)

    @R.function(private=True)
    def expected(A: R.Tensor((1, 32), "float32"), B: R.Tensor((32, 64), "float32")) -> R.Tensor(
        (1, 64), "float32"
    ):
        return R.matmul(A, B)

    after = ParamMutator({"A": (1, 32)}).visit_expr(before)
    tvm.ir.assert_structural_equal(expected, after)


if __name__ == "__main__":
    tvm.testing.main()