# 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.
# pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name, arguments-differ
# ruff: noqa: RUF012
"""The expression and statement functor of TIR."""

from collections.abc import Callable

import tvm_ffi

from tvm.ir import Call, Expr
from tvm.ir.utils import derived_object

from . import _ffi_api
from .expr import (
    EQ,
    GE,
    GT,
    LE,
    LT,
    NE,
    Add,
    And,
    Broadcast,
    BufferLoad,
    Cast,
    Div,
    FloatImm,
    FloorDiv,
    FloorMod,
    IntImm,
    Let,
    Max,
    Min,
    Mod,
    Mul,
    Not,
    Or,
    ProducerLoad,
    Ramp,
    Reduce,
    Select,
    Shuffle,
    StringImm,
    Sub,
    Var,
)
from .stmt import (
    AllocBuffer,
    AssertStmt,
    AttrStmt,
    Bind,
    BufferStore,
    DeclBuffer,
    Evaluate,
    For,
    IfThenElse,
    SBlock,
    SBlockRealize,
    SeqStmt,
    Stmt,
    While,
)

# visitor and mutator are aliases for derived_object
visitor = derived_object
mutator = derived_object

"""
A decorator to wrap user-customized PyStmtExprMutator as TVM object _PyStmtExprMutator.

Parameters
----------
mutator_cls : PyStmtExprMutator
    The user-customized PyStmtExprMutator.

Returns
-------
cls : _PyStmtExprMutator
    The decorated TVM object _PyStmtExprMutator(StmtExprMutator on the C++ side).

Example
-------
.. code-block:: python

    @tirx.functor.stmt_expr_mutator
    class MyStmtExprMutator(PyStmtExprMutator):
        # customize rewrite function
        def visit_add_(self, op: Add) -> Expr:
            # just for demo purposes
            ...

    # mymutator is now a special mutator that rewrite every Add with
    # user-customized visit_add_
    mymutator = MyStmtExprMutator()
    # apply mymutator to Expr and Stmt
    mymutator.visit_expr(expr)
    mymutator.visit_stmt(stmt)
"""


@tvm_ffi.register_object("tirx.PyStmtExprVisitor")
class _PyStmtExprVisitor(tvm_ffi.core.Object):
    """
    An internal wrapper to interface between C++ and Python StmtExprVisitor.
    This is the TVM object that wraps PyStmtExprVisitor.

    Do not use this class directly. Use PyStmtExprVisitor instead.

    See also: PyStmtExprVisitor, stmt_expr_visitor
    """

    def __init__(
        self,
        f_visit_stmt: Callable | None = None,
        f_visit_expr: Callable | None = None,
        # Stmt
        f_visit_bind: Callable | None = None,
        f_visit_attr_stmt: Callable | None = None,
        f_visit_if_then_else: Callable | None = None,
        f_visit_for: Callable | None = None,
        f_visit_while: Callable | None = None,
        f_visit_alloc_buffer: Callable | None = None,
        f_visit_decl_buffer: Callable | None = None,
        f_visit_buffer_store: Callable | None = None,
        f_visit_assert_stmt: Callable | None = None,
        f_visit_seq_stmt: Callable | None = None,
        f_visit_evaluate: Callable | None = None,
        f_visit_block: Callable | None = None,
        f_visit_sblock_realize: Callable | None = None,
        # Expr
        f_visit_var: Callable | None = None,
        f_visit_buffer_load: Callable | None = None,
        f_visit_producer_load: Callable | None = None,
        f_visit_let: Callable | None = None,
        f_visit_call: Callable | None = None,
        f_visit_add: Callable | None = None,
        f_visit_sub: Callable | None = None,
        f_visit_mul: Callable | None = None,
        f_visit_div: Callable | None = None,
        f_visit_mod: Callable | None = None,
        f_visit_floor_div: Callable | None = None,
        f_visit_floor_mod: Callable | None = None,
        f_visit_min: Callable | None = None,
        f_visit_max: Callable | None = None,
        f_visit_eq: Callable | None = None,
        f_visit_ne: Callable | None = None,
        f_visit_lt: Callable | None = None,
        f_visit_le: Callable | None = None,
        f_visit_gt: Callable | None = None,
        f_visit_ge: Callable | None = None,
        f_visit_and: Callable | None = None,
        f_visit_or: Callable | None = None,
        f_visit_reduce: Callable | None = None,
        f_visit_cast: Callable | None = None,
        f_visit_not: Callable | None = None,
        f_visit_select: Callable | None = None,
        f_visit_ramp: Callable | None = None,
        f_visit_broadcast: Callable | None = None,
        f_visit_shuffle: Callable | None = None,
        f_visit_int_imm: Callable | None = None,
        f_visit_float_imm: Callable | None = None,
        f_visit_string_imm: Callable | None = None,
    ) -> None:
        """Constructor."""
        self.__init_handle_by_constructor__(
            _ffi_api.MakePyStmtExprVisitor,  # type: ignore
            f_visit_stmt,
            f_visit_expr,
            # Stmt
            f_visit_bind,
            f_visit_attr_stmt,
            f_visit_if_then_else,
            f_visit_for,
            f_visit_while,
            f_visit_alloc_buffer,
            f_visit_decl_buffer,
            f_visit_buffer_store,
            f_visit_assert_stmt,
            f_visit_seq_stmt,
            f_visit_evaluate,
            f_visit_block,
            f_visit_sblock_realize,
            # Expr
            f_visit_var,
            f_visit_buffer_load,
            f_visit_producer_load,
            f_visit_let,
            f_visit_call,
            f_visit_add,
            f_visit_sub,
            f_visit_mul,
            f_visit_div,
            f_visit_mod,
            f_visit_floor_div,
            f_visit_floor_mod,
            f_visit_min,
            f_visit_max,
            f_visit_eq,
            f_visit_ne,
            f_visit_lt,
            f_visit_le,
            f_visit_gt,
            f_visit_ge,
            f_visit_and,
            f_visit_or,
            f_visit_reduce,
            f_visit_cast,
            f_visit_not,
            f_visit_select,
            f_visit_ramp,
            f_visit_broadcast,
            f_visit_shuffle,
            f_visit_int_imm,
            f_visit_float_imm,
            f_visit_string_imm,
        )


class PyStmtExprVisitor:
    """
    A Python StmtExprVisitor to define custom visitor for both Stmt and Expr.

    Users can customize any of the visit function.
    """

    _tvm_metadata = {
        "cls": _PyStmtExprVisitor,
        "methods": [
            "visit_stmt",
            "visit_expr",
            # Stmt
            "visit_bind_",
            "visit_attr_stmt_",
            "visit_if_then_else_",
            "visit_for_",
            "visit_while_",
            "visit_alloc_buffer_",
            "visit_decl_buffer_",
            "visit_buffer_store_",
            "visit_assert_stmt_",
            "visit_seq_stmt_",
            "visit_evaluate_",
            "visit_sblock_",
            "visit_sblock_realize_",
            # Expr
            "visit_var_",
            "visit_buffer_load_",
            "visit_producer_load_",
            "visit_let_",
            "visit_call_",
            "visit_add_",
            "visit_sub_",
            "visit_mul_",
            "visit_div_",
            "visit_mod_",
            "visit_floor_div_",
            "visit_floor_mod_",
            "visit_min_",
            "visit_max_",
            "visit_eq_",
            "visit_ne_",
            "visit_lt_",
            "visit_le_",
            "visit_gt_",
            "visit_ge_",
            "visit_and_",
            "visit_or_",
            "visit_reduce_",
            "visit_cast_",
            "visit_not_",
            "visit_select_",
            "visit_ramp_",
            "visit_broadcast_",
            "visit_shuffle_",
            "visit_int_imm_",
            "visit_float_imm_",
            "visit_string_imm_",
        ],
    }

    def visit_stmt(self, stmt: Stmt) -> None:
        """Visit a Stmt.

        Parameters
        ----------
        stmt : Stmt
            The Stmt to be visited.
        """
        _ffi_api.PyStmtExprVisitorVisitStmt(self._outer(), stmt)  # type: ignore

    def visit_expr(self, expr: Expr) -> None:
        """Visit a Expr.

        Parameters
        ----------
        expr : Expr
            The Expr to be visited.
        """
        _ffi_api.PyStmtExprVisitorVisitExpr(self._outer(), expr)  # type: ignore

    def visit_attr_stmt_(self, op: AttrStmt) -> None:
        """Visit AttrStmt.
        Users can customize this function to overwrite VisitStmt_(const AttrStmtNode* op)
        on the C++ side.

        Parameters
        ----------
        op : AttrStmt
            The AttrStmt to be visited.
        """
        print("visit_attr_stmt_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_if_then_else_(self, op: IfThenElse) -> None:
        """Visit IfThenElse.
        Users can customize this function to overwrite VisitStmt_(const IfThenElseNode* op)
        on the C++ side.

        Parameters
        ----------
        op : IfThenElse
            The IfThenElse to be visited.
        """
        print("visit_if_then_else_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_bind_(self, op: Bind) -> None:
        """Visit Bind.
        Users can customize this function to overwrite VisitStmt_(const BindNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Bind
            The Bind node to be visited.
        """
        print("visit_bind_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_for_(self, op: For) -> None:
        """Visit For.
        Users can customize this function to overwrite VisitStmt_(const ForNode* op)
        on the C++ side.

        Parameters
        ----------
        op : For
            The For to be visited.
        """
        print("visit_for_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_while_(self, op: While) -> None:
        """Visit While.
        Users can customize this function to overwrite VisitStmt_(const WhileNode* op)
        on the C++ side.

        Parameters
        ----------
        op : While
            The While to be visited.
        """
        print("visit_while_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_alloc_buffer_(self, op: AllocBuffer) -> None:
        """Visit AllocBuffer.
        Users can customize this function to overwrite VisitStmt_(const AllocBufferNode* op)
        on the C++ side.

        Parameters
        ----------
        op : AllocBuffer
            The AllocBuffer to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_decl_buffer_(self, op: DeclBuffer) -> None:
        """Visit DeclBuffer.
        Users can customize this function to overwrite VisitStmt_(const DeclBufferNode* op)
        on the C++ side.

        Parameters
        ----------
        op : DeclBuffer
            The DeclBuffer to be visited.
        """
        print("visit_decl_buffer_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_buffer_store_(self, op: BufferStore) -> None:
        """Visit BufferStore.
        Users can customize this function to overwrite VisitStmt_(const BufferStoreNode* op)
        on the C++ side.

        Parameters
        ----------
        op : BufferStore
            The BufferStore to be visited.
        """
        print("visit_buffer_store_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_assert_stmt_(self, op: AssertStmt) -> None:
        """Visit AssertStmt.
        Users can customize this function to overwrite VisitStmt_(const AssertStmtNode* op)
        on the C++ side.

        Parameters
        ----------
        op : AssertStmt
            The AssertStmt to be visited.
        """
        print("visit_assert_stmt_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_seq_stmt_(self, op: SeqStmt) -> None:
        """Visit SeqStmt.
        Users can customize this function to overwrite VisitStmt_(const SeqStmtNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SeqStmt
            The SeqStmt to be visited.
        """
        print("visit_seq_stmt_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_evaluate_(self, op: Evaluate) -> None:
        """Visit Evaluate.
        Users can customize this function to overwrite VisitStmt_(const EvaluateNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Evaluate
            The Evaluate to be visited.
        """
        print("visit_evaluate_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_sblock_(self, op: SBlock) -> None:
        """Visit SBlock.
        Users can customize this function to overwrite VisitStmt_(const SBlockNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SBlock
            The SBlock to be visited.
        """
        print("visit_sblock_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_sblock_realize_(self, op: SBlockRealize) -> None:
        """Visit BlockRealize.
        Users can customize this function to overwrite VisitStmt_(const SBlockRealizeNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SBlockRealize
            The BlockRealize to be visited.
        """
        print("visit_sblock_realize_", op)
        _ffi_api.PyStmtExprVisitorDefaultVisitStmt(self._outer(), op)  # type: ignore

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

        Users can customize this function to overwrite VisitVar_(const VarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Var
            The Var to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_buffer_load_(self, op: BufferLoad) -> None:
        """Visit BufferLoad.

        Users can customize this function to overwrite VisitBufferLoad_(const BufferLoadNode* op)
        on the C++ side.

        Parameters
        ----------
        op : BufferLoad
            The BufferLoad to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_producer_load_(self, op: ProducerLoad) -> None:
        """Visit ProducerLoad.

        Users can customize this function to overwrite
        VisitProducerLoad_(const ProducerLoadNode* op) on the C++ side.

        Parameters
        ----------
        op : ProducerLoad
            The ProducerLoad to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_let_(self, op: Let) -> None:
        """Visit Let.

        Users can customize this function to overwrite VisitLet_(const LetNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Let
            The Let to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_call_(self, op: Call) -> None:
        """Visit Call.

        Users can customize this function to overwrite VisitCall_(const CallNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Call
            The Call to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_add_(self, op: Add) -> None:
        """Visit Add.

        Users can customize this function to overwrite VisitAdd_(const AddNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Add
            The Add to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_sub_(self, op: Sub) -> None:
        """Visit Sub.

        Users can customize this function to overwrite VisitSub_(const SubNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Sub
            The Sub to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_mul_(self, op: Mul) -> None:
        """Visit Mul.

        Users can customize this function to overwrite VisitMul_(const MulNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Mul
            The Mul to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_div_(self, op: Div) -> None:
        """Visit Div.

        Users can customize this function to overwrite VisitDiv_(const DivNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Div
            The Div to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_mod_(self, op: Mod) -> None:
        """Visit Mod.

        Users can customize this function to overwrite VisitMod_(const ModNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Mod
            The Mod to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_floor_div_(self, op: FloorDiv) -> None:
        """Visit FloorDiv.

        Users can customize this function to overwrite VisitFloorDiv_(const FloorDivNode* op)
        on the C++ side.

        Parameters
        ----------
        op : FloorDiv
            The FloorDiv to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_floor_mod_(self, op: FloorMod) -> None:
        """Visit FloorMod.

        Users can customize this function to overwrite VisitFloorMod_(const FloorModNode* op)
        on the C++ side.

        Parameters
        ----------
        op : FloorMod
            The FloorMod to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_min_(self, op: Min) -> None:
        """Visit Min.

        Users can customize this function to overwrite VisitMin_(const MinNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Min
            The Min to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_max_(self, op: Max) -> None:
        """Visit Max.

        Users can customize this function to overwrite VisitMax_(const MaxNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Max
            The Max to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_eq_(self, op: EQ) -> None:
        """Visit EQ.

        Users can customize this function to overwrite VisitEQ_(const EQNode* op)
        on the C++ side.

        Parameters
        ----------
        op : EQ
            The EQ to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_ne_(self, op: NE) -> None:
        """Visit NE.

        Users can customize this function to overwrite VisitNE_(const NENode* op)
        on the C++ side.

        Parameters
        ----------
        op : NE
            The NE to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_lt_(self, op: LT) -> None:
        """Visit LT.

        Users can customize this function to overwrite VisitLT_(const LTNode* op)
        on the C++ side.

        Parameters
        ----------
        op : LT
            The LT to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_le_(self, op: LE) -> None:
        """Visit LE.

        Users can customize this function to overwrite VisitLE_(const LENode* op)
        on the C++ side.

        Parameters
        ----------
        op : LE
            The LE to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_gt_(self, op: GT) -> None:
        """Visit GT.

        Users can customize this function to overwrite VisitGT_(const GTNode* op)
        on the C++ side.

        Parameters
        ----------
        op : GT
            The GT to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_ge_(self, op: GE) -> None:
        """Visit GE.

        Users can customize this function to overwrite VisitGE_(const GENode* op)
        on the C++ side.

        Parameters
        ----------
        op : GE
            The GE to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_and_(self, op: And) -> None:
        """Visit And.

        Users can customize this function to overwrite VisitAnd_(const AndNode* op)
        on the C++ side.

        Parameters
        ----------
        op : And
            The And to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_or_(self, op: Or) -> None:
        """Visit Or.

        Users can customize this function to overwrite VisitOr_(const OrNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Or
            The Or to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_reduce_(self, op: Reduce) -> None:
        """Visit Reduce.

        Users can customize this function to overwrite VisitReduce_(const ReduceNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Reduce
            The Reduce to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_cast_(self, op: Cast) -> None:
        """Visit Cast.

        Users can customize this function to overwrite VisitCast_(const CastNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Cast
            The Cast to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_not_(self, op: Not) -> None:
        """Visit Not.

        Users can customize this function to overwrite VisitNot_(const NotNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Not
            The Not to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_select_(self, op: Select) -> None:
        """Visit Select.

        Users can customize this function to overwrite VisitSelect_(const SelectNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Select
            The Select to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_ramp_(self, op: Ramp) -> None:
        """Visit Ramp.

        Users can customize this function to overwrite VisitRamp_(const RampNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Ramp
            The Ramp to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_broadcast_(self, op: Broadcast) -> None:
        """Visit Broadcast.

        Users can customize this function to overwrite VisitBroadcast_(const BroadcastNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Broadcast
            The Broadcast to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_shuffle_(self, op: Shuffle) -> None:
        """Visit Shuffle.

        Users can customize this function to overwrite VisitShuffle_(const ShuffleNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Shuffle
            The Shuffle to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_int_imm_(self, op: IntImm) -> None:
        """Visit IntImm.

        Users can customize this function to overwrite VisitIntImm_(const IntImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : IntImm
            The IntImm to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_float_imm_(self, op: FloatImm) -> None:
        """Visit FloatImm.

        Users can customize this function to overwrite VisitFloatImm_(const FloatImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : FloatImm
            The FloatImm to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore

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

        Users can customize this function to overwrite VisitStringImm_(const StringImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : StringImm
            The StringImm to be visited.
        """
        _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op)  # type: ignore


@tvm_ffi.register_object("tirx.PyStmtExprMutator")
class _PyStmtExprMutator(tvm_ffi.core.Object):
    """
    A TVM object to support customization of StmtExprMutator on the python side.
    This is the decorated result returned from stmt_expr_mutator decorator.

    WARNING: This is NOT the user facing class for method overwriting inheritance.

    See also: stmt_expr_mutator, PyStmtExprMutator
    """

    def __init__(
        self,
        f_visit_stmt: Callable | None = None,
        f_visit_expr: Callable | None = None,
        # Stmt
        f_visit_bind: Callable | None = None,
        f_visit_attr_stmt: Callable | None = None,
        f_visit_if_then_else: Callable | None = None,
        f_visit_for: Callable | None = None,
        f_visit_while: Callable | None = None,
        f_visit_alloc_buffer: Callable | None = None,
        f_visit_decl_buffer: Callable | None = None,
        f_visit_buffer_store: Callable | None = None,
        f_visit_assert_stmt: Callable | None = None,
        f_visit_seq_stmt: Callable | None = None,
        f_visit_evaluate: Callable | None = None,
        f_visit_block: Callable | None = None,
        f_visit_sblock_realize: Callable | None = None,
        # Expr
        f_visit_var: Callable | None = None,
        f_visit_buffer_load: Callable | None = None,
        f_visit_producer_load: Callable | None = None,
        f_visit_let: Callable | None = None,
        f_visit_call: Callable | None = None,
        f_visit_add: Callable | None = None,
        f_visit_sub: Callable | None = None,
        f_visit_mul: Callable | None = None,
        f_visit_div: Callable | None = None,
        f_visit_mod: Callable | None = None,
        f_visit_floor_div: Callable | None = None,
        f_visit_floor_mod: Callable | None = None,
        f_visit_min: Callable | None = None,
        f_visit_max: Callable | None = None,
        f_visit_eq: Callable | None = None,
        f_visit_ne: Callable | None = None,
        f_visit_lt: Callable | None = None,
        f_visit_le: Callable | None = None,
        f_visit_gt: Callable | None = None,
        f_visit_ge: Callable | None = None,
        f_visit_and: Callable | None = None,
        f_visit_or: Callable | None = None,
        f_visit_reduce: Callable | None = None,
        f_visit_cast: Callable | None = None,
        f_visit_not: Callable | None = None,
        f_visit_select: Callable | None = None,
        f_visit_ramp: Callable | None = None,
        f_visit_broadcast: Callable | None = None,
        f_visit_shuffle: Callable | None = None,
        f_visit_int_imm: Callable | None = None,
        f_visit_float_imm: Callable | None = None,
        f_visit_string_imm: Callable | None = None,
    ) -> None:
        """Constructor."""
        self.__init_handle_by_constructor__(
            _ffi_api.MakePyStmtExprMutator,  # type: ignore
            f_visit_stmt,
            f_visit_expr,
            # Stmt
            f_visit_bind,
            f_visit_attr_stmt,
            f_visit_if_then_else,
            f_visit_for,
            f_visit_while,
            f_visit_alloc_buffer,
            f_visit_decl_buffer,
            f_visit_buffer_store,
            f_visit_assert_stmt,
            f_visit_seq_stmt,
            f_visit_evaluate,
            f_visit_block,
            f_visit_sblock_realize,
            # Expr
            f_visit_var,
            f_visit_buffer_load,
            f_visit_producer_load,
            f_visit_let,
            f_visit_call,
            f_visit_add,
            f_visit_sub,
            f_visit_mul,
            f_visit_div,
            f_visit_mod,
            f_visit_floor_div,
            f_visit_floor_mod,
            f_visit_min,
            f_visit_max,
            f_visit_eq,
            f_visit_ne,
            f_visit_lt,
            f_visit_le,
            f_visit_gt,
            f_visit_ge,
            f_visit_and,
            f_visit_or,
            f_visit_reduce,
            f_visit_cast,
            f_visit_not,
            f_visit_select,
            f_visit_ramp,
            f_visit_broadcast,
            f_visit_shuffle,
            f_visit_int_imm,
            f_visit_float_imm,
            f_visit_string_imm,
        )


class PyStmtExprMutator:
    """
    A Python StmtExprMutator to define custom mutator for both Stmt and Expr.

    Users can customize any of the visit function.
    """

    _tvm_metadata = {
        "cls": _PyStmtExprMutator,
        "methods": [
            "visit_stmt",
            "visit_expr",
            # Stmt
            "visit_bind_",
            "visit_attr_stmt_",
            "visit_if_then_else_",
            "visit_for_",
            "visit_while_",
            "visit_alloc_buffer_",
            "visit_decl_buffer_",
            "visit_buffer_store_",
            "visit_assert_stmt_",
            "visit_seq_stmt_",
            "visit_evaluate_",
            "visit_sblock_",
            "visit_sblock_realize_",
            # Expr
            "visit_var_",
            "visit_buffer_load_",
            "visit_producer_load_",
            "visit_let_",
            "visit_call_",
            "visit_add_",
            "visit_sub_",
            "visit_mul_",
            "visit_div_",
            "visit_mod_",
            "visit_floor_div_",
            "visit_floor_mod_",
            "visit_min_",
            "visit_max_",
            "visit_eq_",
            "visit_ne_",
            "visit_lt_",
            "visit_le_",
            "visit_gt_",
            "visit_ge_",
            "visit_and_",
            "visit_or_",
            "visit_reduce_",
            "visit_cast_",
            "visit_not_",
            "visit_select_",
            "visit_ramp_",
            "visit_broadcast_",
            "visit_shuffle_",
            "visit_int_imm_",
            "visit_float_imm_",
            "visit_string_imm_",
        ],
    }

    def visit_expr(self, expr: Expr) -> Expr:
        """Visit Expr.
        Users can customize this function to overwrite VisitExpr(const Expr& expr)
        on the C++ side.

        Parameters
        ----------
        expr : Expr
            The Expr to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorVisitExpr(self._outer(), expr)  # type: ignore

    def visit_stmt(self, stmt: Stmt) -> Stmt:
        """Visit Stmt.
        Users can customize this function to overwrite VisitStmt(const Stmt& stmt)
        on the C++ side.

        Parameters
        ----------
        stmt : Stmt
            The Stmt to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorVisitStmt(self._outer(), stmt)  # type: ignore

    def visit_attr_stmt_(self, op: AttrStmt) -> Stmt:
        """Visit AttrStmt.
        Users can customize this function to overwrite VisitStmt_(const AttrStmtNode* op)
        on the C++ side.

        Parameters
        ----------
        op : AttrStmt
            The AttrStmt to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_if_then_else_(self, op: IfThenElse) -> Stmt:
        """Visit IfThenElse.
        Users can customize this function to overwrite VisitStmt_(const IfThenElseNode* op)
        on the C++ side.

        Parameters
        ----------
        op : IfThenElse
            The IfThenElse to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_bind_(self, op: Bind) -> Stmt:
        """Visit Bind.
        Users can customize this function to overwrite VisitStmt_(const BindNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Bind
            The Bind node to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_for_(self, op: For) -> Stmt:
        """Visit For.
        Users can customize this function to overwrite VisitStmt_(const ForNode* op)
        on the C++ side.

        Parameters
        ----------
        op : For
            The For to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_while_(self, op: While) -> Stmt:
        """Visit While.
        Users can customize this function to overwrite VisitStmt_(const WhileNode* op)
        on the C++ side.

        Parameters
        ----------
        op : While
            The While to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_alloc_buffer_(self, op: AllocBuffer) -> Stmt:
        """Visit AllocBuffer.
        Users can customize this function to overwrite VisitStmt_(const AllocBufferNode* op)
        on the C++ side.

        Parameters
        ----------
        op : AllocBuffer
            The AllocBuffer to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_decl_buffer_(self, op: DeclBuffer) -> Stmt:
        """Visit DeclBuffer.
        Users can customize this function to overwrite VisitStmt_(const DeclBufferNode* op)
        on the C++ side.

        Parameters
        ----------
        op : DeclBuffer
            The DeclBuffer to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_buffer_store_(self, op: BufferStore) -> Stmt:
        """Visit BufferStore.
        Users can customize this function to overwrite VisitStmt_(const BufferStoreNode* op)
        on the C++ side.

        Parameters
        ----------
        op : BufferStore
            The BufferStore to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_assert_stmt_(self, op: AssertStmt) -> Stmt:
        """Visit AssertStmt.
        Users can customize this function to overwrite VisitStmt_(const AssertStmtNode* op)
        on the C++ side.

        Parameters
        ----------
        op : AssertStmt
            The AssertStmt to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_seq_stmt_(self, op: SeqStmt) -> Stmt:
        """Visit SeqStmt.
        Users can customize this function to overwrite VisitStmt_(const SeqStmtNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SeqStmt
            The SeqStmt to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_evaluate_(self, op: Evaluate) -> Stmt:
        """Visit Evaluate.
        Users can customize this function to overwrite VisitStmt_(const EvaluateNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Evaluate
            The Evaluate to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_sblock_(self, op: SBlock) -> Stmt:
        """Visit SBlock.
        Users can customize this function to overwrite VisitStmt_(const SBlockNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SBlock
            The SBlock to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_sblock_realize_(self, op: SBlockRealize) -> Stmt:
        """Visit BlockRealize.
        Users can customize this function to overwrite VisitStmt_(const SBlockRealizeNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SBlockRealize
            The SBlockRealize to be visited.

        Returns
        -------
        result : Stmt
            The mutated Stmt.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitStmt(self._outer(), op)  # type: ignore

    def visit_var_(self, op: Var) -> Expr:
        """Visit Var.

        Users can customize this function to overwrite VisitVar_(const VarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Var
            The Var to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_buffer_load_(self, op: BufferLoad) -> Expr:
        """Visit BufferLoad.

        Users can customize this function to overwrite VisitBufferLoad_(const BufferLoadNode* op)
        on the C++ side.

        Parameters
        ----------
        op : BufferLoad
            The BufferLoad to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_producer_load_(self, op: ProducerLoad) -> Expr:
        """Visit ProducerLoad.

        Users can customize this function to overwrite
        VisitProducerLoad_(const ProducerLoadNode* op) on the C++ side.

        Parameters
        ----------
        op : ProducerLoad
            The ProducerLoad to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_let_(self, op: Let) -> Expr:
        """Visit Let.

        Users can customize this function to overwrite VisitLet_(const LetNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Let
            The Let to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_call_(self, op: Call) -> Expr:
        """Visit Call.

        Users can customize this function to overwrite VisitCall_(const CallNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Call
            The Call to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_add_(self, op: Add) -> Expr:
        """Visit Add.

        Users can customize this function to overwrite VisitAdd_(const AddNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Add
            The Add to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_sub_(self, op: Sub) -> Expr:
        """Visit Sub.

        Users can customize this function to overwrite VisitSub_(const SubNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Sub
            The Sub to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_mul_(self, op: Mul) -> Expr:
        """Visit Mul.

        Users can customize this function to overwrite VisitMul_(const MulNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Mul
            The Mul to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_div_(self, op: Div) -> Expr:
        """Visit Div.

        Users can customize this function to overwrite VisitDiv_(const DivNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Div
            The Div to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_mod_(self, op: Mod) -> Expr:
        """Visit Mod.

        Users can customize this function to overwrite VisitMod_(const ModNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Mod
            The Mod to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_floor_div_(self, op: FloorDiv) -> Expr:
        """Visit FloorDiv.

        Users can customize this function to overwrite VisitFloorDiv_(const FloorDivNode* op)
        on the C++ side.

        Parameters
        ----------
        op : FloorDiv
            The FloorDiv to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_floor_mod_(self, op: FloorMod) -> Expr:
        """Visit FloorMod.

        Users can customize this function to overwrite VisitFloorMod_(const FloorModNode* op)
        on the C++ side.

        Parameters
        ----------
        op : FloorMod
            The FloorMod to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_min_(self, op: Min) -> Expr:
        """Visit Min.

        Users can customize this function to overwrite VisitMin_(const MinNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Min
            The Min to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_max_(self, op: Max) -> Expr:
        """Visit Max.

        Users can customize this function to overwrite VisitMax_(const MaxNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Max
            The Max to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_eq_(self, op: EQ) -> Expr:
        """Visit EQ.

        Users can customize this function to overwrite VisitEQ_(const EQNode* op)
        on the C++ side.

        Parameters
        ----------
        op : EQ
            The EQ to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_ne_(self, op: NE) -> Expr:
        """Visit NE.

        Users can customize this function to overwrite VisitNE_(const NENode* op)
        on the C++ side.

        Parameters
        ----------
        op : NE
            The NE to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_lt_(self, op: LT) -> Expr:
        """Visit LT.

        Users can customize this function to overwrite VisitLT_(const LTNode* op)
        on the C++ side.

        Parameters
        ----------
        op : LT
            The LT to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_le_(self, op: LE) -> Expr:
        """Visit LE.

        Users can customize this function to overwrite VisitLE_(const LENode* op)
        on the C++ side.

        Parameters
        ----------
        op : LE
            The LE to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_gt_(self, op: GT) -> Expr:
        """Visit GT.

        Users can customize this function to overwrite VisitGT_(const GTNode* op)
        on the C++ side.

        Parameters
        ----------
        op : GT
            The GT to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_ge_(self, op: GE) -> Expr:
        """Visit GE.

        Users can customize this function to overwrite VisitGE_(const GENode* op)
        on the C++ side.

        Parameters
        ----------
        op : GE
            The GE to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_and_(self, op: And) -> Expr:
        """Visit And.

        Users can customize this function to overwrite VisitAnd_(const AndNode* op)
        on the C++ side.

        Parameters
        ----------
        op : And
            The And to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_or_(self, op: Or) -> Expr:
        """Visit Or.

        Users can customize this function to overwrite VisitOr_(const OrNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Or
            The Or to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_reduce_(self, op: Reduce) -> Expr:
        """Visit Reduce.

        Users can customize this function to overwrite VisitReduce_(const ReduceNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Reduce
            The Reduce to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_cast_(self, op: Cast) -> Expr:
        """Visit Cast.

        Users can customize this function to overwrite VisitCast_(const CastNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Cast
            The Cast to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_not_(self, op: Not) -> Expr:
        """Visit Not.

        Users can customize this function to overwrite VisitNot_(const NotNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Not
            The Not to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_select_(self, op: Select) -> Expr:
        """Visit Select.

        Users can customize this function to overwrite VisitSelect_(const SelectNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Select
            The Select to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_ramp_(self, op: Ramp) -> Expr:
        """Visit Ramp.

        Users can customize this function to overwrite VisitRamp_(const RampNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Ramp
            The Ramp to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_broadcast_(self, op: Broadcast) -> Expr:
        """Visit Broadcast.

        Users can customize this function to overwrite VisitBroadcast_(const BroadcastNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Broadcast
            The Broadcast to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_shuffle_(self, op: Shuffle) -> Expr:
        """Visit Shuffle.

        Users can customize this function to overwrite VisitShuffle_(const ShuffleNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Shuffle
            The Shuffle to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_int_imm_(self, op: IntImm) -> Expr:
        """Visit IntImm.

        Users can customize this function to overwrite VisitIntImm_(const IntImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : IntImm
            The IntImm to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_float_imm_(self, op: FloatImm) -> Expr:
        """Visit FloatImm.

        Users can customize this function to overwrite VisitFloatImm_(const FloatImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : FloatImm
            The FloatImm to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore

    def visit_string_imm_(self, op: StringImm) -> Expr:
        """Visit StringImm.

        Users can customize this function to overwrite VisitStringImm_(const StringImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : StringImm
            The StringImm to be visited.

        Returns
        -------
        result : Expr
            The mutated Expr.
        """
        return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)  # type: ignore
