# 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=missing-docstring
# ruff: noqa: E501, F841

from tvm_ffi.access_path import AccessPath

import tvm
import tvm.testing
from tvm import IRModule, relax, tirx
from tvm.runtime.script_printer import PrinterConfig, _script
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


def _assert_print(obj, expected):
    if not isinstance(obj, str):
        obj = obj.script(verbose_expr=True)
    obj = obj.strip()
    # compare line by line in case there is trailing whitespace in the _middle_
    for obj_line, expected_line in zip(obj.splitlines(), expected.strip().splitlines()):
        assert obj_line.strip() == expected_line.strip(), "\n" + obj


def test_function():
    @R.function
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):  # type: ignore
        R.func_attr({"some_attr": 1})
        return a

    _assert_print(
        func,
        """
# from tvm.script import relax as R

@R.function
def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
    R.func_attr({"some_attr": 1})
    return a""",
    )


def test_function_dependent_shape_escaped_source_spans():
    n = tirx.Var("n", "int64")
    cast = tirx.Cast("int64", n)
    x = relax.Var("x", relax.TensorType([cast], "float32"))
    ret_ty = relax.TensorType(dtype="float32", ndim=1)
    func = relax.Function([x], x, ret_ty=ret_ty).with_attr("global_symbol", "main")
    cast_path = (
        AccessPath.root()
        .attr("params")
        .array_item(0)
        .attr("ty")
        .attr("shape")
        .attr("values")
        .array_item(0)
    )

    def render(path):
        config = PrinterConfig(
            verbose_expr=True,
            num_context_lines=0,
            path_to_underline=[path],
            extra_config={"render_invisible_path_info": True},
        )
        first = _script(func, config)
        assert _script(func, config) == first
        assert first.count("Access path:") == 1
        lines = first.splitlines()
        definition_index = next(i for i, line in enumerate(lines) if "def main" in line)
        assert "Access path:" not in lines[definition_index]
        return lines[definition_index], lines[definition_index + 1]

    expression = r'"T.Cast(\"int64\", n)"'
    definition, underline = render(cast_path)
    expression_start = definition.index(expression)
    assert underline[expression_start : expression_start + len(expression)] == "^" * len(expression)
    assert underline.strip() == "^" * len(expression)

    escaped_dtype = r"\"int64\""
    definition, underline = render(cast_path.attr("dtype"))
    dtype_start = definition.index(escaped_dtype)
    assert underline[dtype_start : dtype_start + len(escaped_dtype)] == "^" * len(escaped_dtype)
    assert underline.strip() == "^" * len(escaped_dtype)

    definition, underline = render(cast_path.attr("value"))
    variable_start = definition.index(expression) + expression.rindex("n")
    assert underline[variable_start] == "^"
    assert underline.strip() == "^"


def test_lone_private_function():
    @R.function(private=True)
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):  # type: ignore
        R.func_attr({"some_attr": 1})
        return a

    # name prints as main because without a global symbol, the printer cannot assume a name
    _assert_print(
        func,
        """
# from tvm.script import relax as R

@R.function(private=True)
def main(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
    R.func_attr({"some_attr": 1})
    return a""",
    )


def test_extern_func():
    @R.function
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):  # type: ignore
        return a

    obj = IRModule(
        {
            "func": func,
            "my_ext": relax.ExternFunc("my_ext"),
        }
    )
    _assert_print(
        obj,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    my_ext = R.ExternFunc("my_ext")
    @R.function
    def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
        return a
""",
    )


def test_extern_func_with_ty():
    obj = IRModule(
        {
            "my_ext": relax.ExternFunc(
                "my_ext",
                relax.FuncType([], relax.TensorType(dtype="float32", ndim=2), purity=True),
            ),
        }
    )
    _assert_print(
        obj,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    my_ext = R.ExternFunc("my_ext", R.Callable((), R.Tensor(dtype="float32", ndim=2), True))
""",
    )


def test_extern_func_with_ty_roundtrip():
    mod = IRModule(
        {
            "my_ext": relax.ExternFunc(
                "my_ext",
                relax.FuncType([], relax.TensorType(dtype="float32", ndim=2), purity=True),
            ),
        }
    )
    roundtrip = tvm.script.from_source(mod.script(verbose_expr=True))
    tvm.ir.assert_structural_equal(mod, roundtrip)


def test_nested_function():
    @I.ir_module(s_tir=True)
    class NestedFunction:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            @R.function
            def nested(y: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
                return y

            z = nested(x)
            return z

    _assert_print(
        NestedFunction,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        # from tvm.script import relax as R

        @R.function
        def nested(y: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
            return y

        z: R.Tensor((), dtype="int32") = nested(x)
        return z
""",
    )


def test_object_ty():
    obj = relax.AnyType()
    _assert_print(
        obj,
        "R.Any",
    )


def test_prim_ty():
    obj = tvm.ir.PrimType("float32")
    _assert_print(obj, "T.float32")


def test_shape_ty_0():
    obj = relax.ShapeType(ndim=-1)
    _assert_print(obj, "R.Shape(ndim=-1)")


def test_shape_ty_1():
    obj = relax.ShapeType([1, 2, 3])
    _assert_print(obj, "R.Shape([1, 2, 3])")


def test_shape_ty_2():
    obj = relax.ShapeType([1, tirx.Var("a", "int64"), 3])
    _assert_print(
        obj,
        """
a = T.int64()
R.Shape([1, a, 3])""",
    )


def test_tensor_ty():
    obj = relax.TensorType(
        shape=relax.ShapeExpr([1, tirx.Var("a", "int64"), 3]),
        dtype="float32",
    )
    _assert_print(
        obj,
        """
a = T.int64()
R.Tensor((1, a, 3), dtype="float32")
""",
    )


def test_tuple_ty_empty():
    obj = relax.TupleType([])
    _assert_print(obj._relax_script(), "R.Tuple")  # pylint: disable=protected-access


def test_tuple_ty():
    obj = relax.TupleType(
        [
            tvm.ir.PrimType("float32"),
            relax.AnyType(),
            relax.ShapeType([1, tirx.Var("a", "int64"), 3]),
        ]
    )
    _assert_print(
        obj._relax_script(),  # pylint: disable=protected-access
        """
R.Tuple(T.float32, R.Any, R.Shape([1, a, 3]))
""",
    )


def test_func_ty():
    obj = relax.FuncType(
        params=[
            tvm.ir.PrimType("float32"),
            relax.AnyType(),
            relax.ShapeType([1, tirx.Var("a", "int64"), 3]),
            tvm.ir.PrimType("int64"),
        ],
        ret=relax.TensorType(
            shape=relax.ShapeExpr([1, 2, 3]),
            dtype="float32",
        ),
    )
    _assert_print(
        obj,
        "a = T.int64()\n"
        "R.Callable((T.float32, R.Any, R.Shape([1, a, 3]), T.int64), "
        'R.Tensor((1, 2, 3), dtype="float32"), True)',
    )


def test_shape_type():
    obj = relax.ShapeType(ndim=3)
    _assert_print(obj, "R.Shape(ndim=3)")


def test_object_type():
    obj = relax.AnyType()
    _assert_print(obj, "R.Any")


def test_dyn_tensor_type():
    obj = relax.TensorType()
    _assert_print(obj, 'R.Tensor(dtype="float32")')


def test_packed_func_type():
    obj = relax.PackedFuncType()
    _assert_print(obj, "R.PackedFunc")


def test_tuple_type():
    obj = relax.TupleType([relax.ShapeType(ndim=3), relax.AnyType()])
    _assert_print(
        obj._relax_script(),  # pylint: disable=protected-access
        "R.Tuple(R.Shape(ndim=3), R.Any)",
    )


def test_func_type():
    obj = relax.FuncType(
        params=[
            relax.AnyType(),
            relax.ShapeType(ndim=3),
        ],
        ret=relax.TensorType(
            ndim=3,
            dtype="float32",
        ),
    )
    _assert_print(
        obj._relax_script(),  # pylint: disable=protected-access
        'R.Callable((R.Any, R.Shape(ndim=3)), R.Tensor(dtype="float32", ndim=3), True)',
    )


def test_prim_value():
    obj = tirx.IntImm("int64", 1)
    _assert_print(obj, "T.int64(1)")

    @R.function
    def func() -> R.Prim("int64"):
        return R.prim_value(1)

    _assert_print(
        func,
        """
# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis
# from tvm.script import relax as R

@R.function
def func() -> T.int64:
    return 1""",
    )

    @R.function
    def float_func() -> R.Prim("float32"):
        return R.prim_value(T.float32(1.0))

    float_script = float_func.script(verbose_expr=True)
    assert "R.prim_value" not in float_script
    assert "return T.float32(" in float_script
    tvm.ir.assert_structural_equal(tvm.script.from_source(float_script), float_func)


def test_string_imm():
    obj = relax.StringImm("hello")
    _assert_print(obj, 'R.str("hello")')


def test_data_type_imm():
    obj = relax.DataTypeImm("float32")
    _assert_print(obj, 'R.dtype("float32")')


def test_var():
    obj = relax.Var("a", relax.TensorType([1, tirx.Var("x", "int64"), 3], "float32"))
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
a""",
    )


def test_dataflow_var():
    obj = relax.DataflowVar("a", relax.TensorType([1, tirx.Var("x", "int64"), 3], "float32"))
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
a""",
    )


def test_tuple():
    obj = relax.Tuple(
        [
            relax.Var("a", relax.TensorType([1, tirx.Var("x", "int64"), 3], "float32")),
            relax.Var("b", relax.TensorType([1, tirx.Var("y", "int64"), 3], "float32")),
            relax.Var("c", relax.TensorType([1, tirx.Var("z", "int64"), 3], "float32")),
        ]
    )
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
y = T.int64()
b: R.Tensor((1, y, 3), dtype="float32")
z = T.int64()
c: R.Tensor((1, z, 3), dtype="float32")
(a, b, c)
""",
    )


def test_tuple_get_item():
    obj = relax.TupleGetItem(
        relax.Tuple(
            [
                relax.Var("a", relax.TensorType([1, tirx.Var("x", "int64"), 3], "float32")),
                relax.Var("b", relax.TensorType([1, tirx.Var("y", "int64"), 3], "float32")),
                relax.Var("c", relax.TensorType([1, tirx.Var("z", "int64"), 3], "float32")),
            ]
        ),
        0,
    )
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
y = T.int64()
b: R.Tensor((1, y, 3), dtype="float32")
z = T.int64()
c: R.Tensor((1, z, 3), dtype="float32")
(a, b, c)[0]
""",
    )


def test_shape_expr():
    obj = relax.ShapeExpr([1, 2, 3])
    _assert_print(obj, "R.shape([1, 2, 3])")


def test_call():
    x = tirx.Var("x", "int64")
    a = relax.Var("a", relax.TensorType([1, x, 3], "float32"))
    o0 = relax.call_tir(relax.GlobalVar("tir_func"), args=a, out_ty=a.ty, tir_vars=[x])
    o1 = relax.call_dps_packed("my_dps_func", args=a, out_ty=a.ty)
    _assert_print(
        o0,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
R.call_tir(tir_func, (a,), out_ty=R.Tensor((1, x, 3), dtype="float32"), tir_vars=R.shape([x]))
""",
    )
    _assert_print(
        o1,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
R.call_dps_packed("my_dps_func", (a,), out_ty=R.Tensor((1, x, 3), dtype="float32"))
""",
    )


def test_call_tir_with_grad():
    x = tirx.Var("x", "int64")
    v0 = relax.Var("v0", R.Tensor([54, 96], "float32"))
    v1 = relax.call_tir_with_grad(
        relax.GlobalVar("tir_func"),
        (v0,),
        R.Tensor((54, 96), "float32"),
        te_grad_name="grad_func",
        te_grad_kwargs={"k": 1.0, "x": x},
    )

    _assert_print(
        v1,
        """
v0: R.Tensor((54, 96), dtype="float32")
x = T.int64()
R.call_tir_with_grad(tir_func, (v0,), out_ty=R.Tensor((54, 96), dtype="float32"), te_grad_name="grad_func", te_grad_kwargs={"k": 1.0, "x": x})
""",
    )


def test_call_tir_inplace():
    x = relax.Var("x", R.Tensor((32, 32), dtype="int32"))
    y = relax.Var("y", R.Tensor((32, 32), dtype="int32"))
    t = tirx.Var("t", dtype="int64")
    call = relax.call_tir_inplace(
        relax.GlobalVar("tir_func"),
        (
            x,
            y,
        ),
        inplace_indices=[-1, 0],
        out_ty=[R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32), dtype="int32")],
        tir_vars=[t],
    )
    _assert_print(
        call,
        """
x: R.Tensor((32, 32), dtype="int32")
y: R.Tensor((32, 32), dtype="int32")
t = T.int64()
R.call_tir_inplace(tir_func, (x, y), out_ty=[R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32), dtype="int32")], inplace_indices=[-1, 0], tir_vars=R.shape([t]))
        """,
    )


def test_seq_expr():
    x = tirx.Var("x", "int64")
    a = relax.Var("a", relax.TensorType([1, x, 3], "float32"))
    b = relax.DataflowVar("b", relax.TensorType([1, x, 3], "float32"))
    c = relax.Var("c", relax.TensorType([1, x, 3], "float32"))

    obj = relax.SeqExpr(
        blocks=[
            relax.DataflowBlock(
                bindings=[
                    relax.VarBinding(b, relax.op.sin(a)),
                    relax.VarBinding(c, relax.op.sin(b)),
                ]
            ),
        ],
        body=c,
    )
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
with R.dataflow():
    b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a)
    c: R.Tensor((1, x, 3), dtype="float32") = R.sin(b)
    R.output(c)
c
""",
    )


def test_binding_block():
    x = tirx.Var("x", "int64")
    a = relax.Var("a", relax.TensorType([1, x, 3], "float32"))
    b = relax.Var("b", relax.TensorType([1, x, 3], "float32"))
    c = relax.Var("c", relax.TensorType([1, x, 3], "float32"))
    obj = relax.BindingBlock(
        bindings=[
            relax.VarBinding(b, relax.op.sin(a)),
            relax.VarBinding(c, relax.op.sin(b)),
        ]
    )
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a)
c: R.Tensor((1, x, 3), dtype="float32") = R.sin(b)
""",
    )


def test_dataflow_block():
    x = tirx.Var("x", "int64")
    a = relax.Var("a", relax.TensorType([1, x, 3], "float32"))
    b = relax.DataflowVar("b", relax.TensorType([1, x, 3], "float32"))
    c = relax.Var("c", relax.TensorType([1, x, 3], "float32"))
    obj = relax.DataflowBlock(
        bindings=[
            relax.VarBinding(b, relax.op.sin(a)),
            relax.VarBinding(c, relax.op.sin(b)),
        ]
    )
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
with R.dataflow():
    b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a)
    c: R.Tensor((1, x, 3), dtype="float32") = R.sin(b)
    R.output(c)
""",
    )


def test_match_cast():
    x = tirx.Var("x", "int64")
    a = relax.Var("a", relax.TensorType([1, x, 3]))
    b = relax.Var("b", relax.TensorType([1, 5, 3]))
    obj = relax.MatchCast(
        var=b,
        value=a,
        ty=b.ty,
    )
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
b: R.Tensor((1, 5, 3), dtype="float32") = R.match_cast(a, R.Tensor((1, 5, 3), dtype="float32"))
""",
    )


def test_var_binding():
    x = tirx.Var("x", "int64")
    a = relax.Var("a", relax.TensorType([1, x, 3], "float32"))
    b = relax.Var("b", relax.TensorType([1, x, 3], "float32"))
    obj = relax.VarBinding(b, relax.op.sin(a))
    _assert_print(
        obj,
        """
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a)
""",
    )


def test_if():
    a = relax.Var("a", relax.TensorType([], "bool"))
    b = relax.Var("b", relax.TensorType([1, 2, 3], "float32"))
    c = relax.Var("c", relax.TensorType([1, 2, 3], "float32"))
    obj = relax.If(
        a,
        relax.SeqExpr([], b),
        relax.SeqExpr([], c),
    )
    _assert_print(
        obj,
        """
a: R.Tensor((), dtype="bool")
if a:
    b: R.Tensor((1, 2, 3), dtype="float32")
    b
else:
    c: R.Tensor((1, 2, 3), dtype="float32")
    c
""",
    )


def test_builtin_keywords():
    x = tirx.Var("x", "int64")
    a = relax.Var("R", relax.TensorType([1, x, 3], "float32"))
    b = relax.Var("T", relax.TensorType([1, x, 3], "float32"))
    obj = relax.VarBinding(b, relax.op.sin(a))
    _assert_print(
        obj,
        """
x = T.int64()
R_1: R.Tensor((1, x, 3), dtype="float32")
T_1: R.Tensor((1, x, 3), dtype="float32") = R.sin(R_1)
""",
    )


def test_module_cross_func_call():
    @I.ir_module(s_tir=True)
    class TestModule:
        @T.prim_func(s_tir=True)
        def tir_func(
            x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
        ):
            T.evaluate(0)

        @R.function
        def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
            cls = TestModule
            gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128,), dtype="float32"))
            return gv0

    # default behavior
    _assert_print(
        TestModule,
        """
# from tvm.script import ir as I
# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis
# from tvm.script import relax as R

@I.ir_module
class Module:
    @T.prim_func(s_tir=True)
    def tir_func(x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")):
        T.evaluate(0)

    @R.function
    def foo(x: R.Tensor((128,), dtype="float32")) -> R.Tensor((128,), dtype="float32"):
        cls = Module
        gv0 = R.call_tir(cls.tir_func, (x,), out_ty=R.Tensor((128,), dtype="float32"))
        return gv0
""",
    )

    # empty module alias
    module_str = TestModule.script(module_alias="")
    _assert_print(
        module_str,
        """
# from tvm.script import ir as I
# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis
# from tvm.script import relax as R

@I.ir_module
class Module:
    @T.prim_func(s_tir=True)
    def tir_func(x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")):
        T.evaluate(0)

    @R.function
    def foo(x: R.Tensor((128,), dtype="float32")) -> R.Tensor((128,), dtype="float32"):
        gv0 = R.call_tir(Module.tir_func, (x,), out_ty=R.Tensor((128,), dtype="float32"))
        return gv0
""",
    )


def test_assert_op():
    @I.ir_module(s_tir=True)
    class AssertOpMod:
        @R.function(pure=False)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
            return x

    _assert_print(
        AssertOpMod,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function(pure=False)
    def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        R.assert_op(R.const(False, "bool"), x, format=R.str("x: {}"))
        return x
""",
    )


def test_print():
    @I.ir_module(s_tir=True)
    class PrintMod:
        @R.function(pure=False)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.print(x, format="x: {}")
            return x

    _assert_print(
        PrintMod,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function(pure=False)
    def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        R.print(x, format=R.str("x: {}"))
        return x
""",
    )


def test_private_function():
    @I.ir_module(s_tir=True)
    class AddMod:
        @R.function(private=True)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y: R.Tensor((), dtype="int32") = R.add(x, x)
            return y

    _assert_print(
        AddMod,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function(private=True)
    def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        y: R.Tensor((), dtype="int32") = R.add(x, x)
        return y
""",
    )


def test_directly_construct_private_funcs():
    # public
    @R.function
    def foo(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
        y: R.Tensor((), dtype="int32") = R.add(x, x)
        return y

    # private
    @R.function(private=True)
    def bar(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
        y: R.Tensor((), dtype="int32") = R.multiply(x, x)
        return y

    # public but there's another attribute
    @R.function
    def baz(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
        R.func_attr({"relax.force_pure": True})
        y: R.Tuple = R.print(format="Hi there!")
        z: R.Tensor((), dtype="int32") = R.add(x, x)
        return z

    # private with an attribute
    @R.function(private=True)
    def quux(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
        R.func_attr({"relax.force_pure": True})
        y: R.Tuple = R.print(format="Lol")
        z: R.Tensor((), dtype="int32") = R.multiply(x, x)
        return z

    obj = IRModule(
        {
            "foo": foo,
            "bar": bar,
            "baz": baz,
            "quux": quux,
        }
    )
    _assert_print(
        obj,
        """
# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function(private=True)
    def bar(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        y: R.Tensor((), dtype="int32") = R.multiply(x, x)
        return y

    @R.function
    def baz(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        R.func_attr({"relax.force_pure": True})
        R.print(format=R.str("Hi there!"))
        z: R.Tensor((), dtype="int32") = R.add(x, x)
        return z

    @R.function
    def foo(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        y: R.Tensor((), dtype="int32") = R.add(x, x)
        return y

    @R.function(private=True)
    def quux(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        R.func_attr({"relax.force_pure": True})
        R.print(format=R.str("Lol"))
        z: R.Tensor((), dtype="int32") = R.multiply(x, x)
        return z
""",
    )


def test_reused_extern_func():
    """An ExternFunc used in a variable binding should be explicit"""

    @R.function
    def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
        extern_func = R.ExternFunc("extern_func")
        y = R.call_dps_packed(extern_func, (x,), out_ty=R.Tensor((128, 128), dtype="float32"))
        z = R.call_dps_packed(extern_func, (y,), out_ty=R.Tensor((128, 128), dtype="float32"))
        return z

    _assert_print(
        func,
        """
# from tvm.script import relax as R

@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
    extern_func: R.Callable = R.ExternFunc("extern_func")
    y = R.call_dps_packed(extern_func, (x,), out_ty=R.Tensor((128, 128), dtype="float32"))
    z = R.call_dps_packed(extern_func, (y,), out_ty=R.Tensor((128, 128), dtype="float32"))
    return z
                  """,
    )


def test_inline_extern_func():
    """An ExternFunc used in-line may be printed as a string"""

    @R.function
    def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
        y = R.call_dps_packed(
            R.ExternFunc("extern_func"), (x,), out_ty=R.Tensor((128, 128), dtype="float32")
        )
        z = R.call_dps_packed(
            R.ExternFunc("extern_func"), (y,), out_ty=R.Tensor((128, 128), dtype="float32")
        )
        return z

    _assert_print(
        func,
        """
# from tvm.script import relax as R

@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
    y = R.call_dps_packed("extern_func", (x,), out_ty=R.Tensor((128, 128), dtype="float32"))
    z = R.call_dps_packed("extern_func", (y,), out_ty=R.Tensor((128, 128), dtype="float32"))
    return z
                  """,
    )


def test_hide_inferable_ty():
    """Redundant type annotations can be omitted

    When `show_all_ty=False`, TVMScript type annotations that
    provide redundant type can be omitted.
    """

    @R.function
    def func(A: R.Tensor([10, 20], "float32"), B: R.Tensor(ndim=2, dtype="float32")):
        # R.match_cast has the type as an argument, so it can
        # be omitted from the variable annotation.
        B2 = R.match_cast(B, R.Tensor([10, 20], "float32"))

        # Call nodes may have inferable shapes from their arguments.
        C = R.add(A, B2)

        # Trivial bindings can be inferred to have the same struct
        # info as the RHS.
        D = C

        # Here, the type cannot be omitted.  `R.add(D,B)` has
        # type `R.Tensor(ndim=2)`, but the variable has a shape
        # `R.Tensor([10,20])`.  This is compatible, so it is not an
        # error to have this annotation, but it is not inferrable from
        # the RHS.  Therefore, it must still be printed.
        E: R.Tensor([10, 20], "float32") = R.add(D, B)

        # The return type can be inferred from function body, but is
        # still always printed in the TVMScript.  When parsing an
        # IRModule with functions calling each other, the return type
        # of each callee must be available for use in the caller's
        # shape inference.
        return E

    _assert_print(
        func.script(show_all_ty=False),
        """
# from tvm.script import relax as R

@R.function
def func(A: R.Tensor((10, 20), dtype="float32"), B: R.Tensor(dtype="float32", ndim=2)) -> R.Tensor((10, 20), dtype="float32"):
    B2 = R.match_cast(B, R.Tensor((10, 20), dtype="float32"))
    C = R.add(A, B2)
    D = C
    E: R.Tensor((10, 20), dtype="float32") = R.add(D, B)
    return E""",
    )


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