# RUN: %PYTHON %s | FileCheck %s --enable-var-scope=false

import gc
from mlir.ir import *
from mlir.dialects import func


def run(f):
    print("\nTEST:", f.__name__)
    f()
    gc.collect()
    assert Context._get_live_count() == 0
    return f


# CHECK-LABEL: TEST: testCapsuleConversions
@run
def testCapsuleConversions():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        value = Operation.create("custom.op1", results=[i32]).result
        value_capsule = value._CAPIPtr
        assert '"mlir.ir.Value._CAPIPtr"' in repr(value_capsule)
        value2 = Value._CAPICreate(value_capsule)
        assert value2 == value


# CHECK-LABEL: TEST: testOpResultOwner
@run
def testOpResultOwner():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        op = Operation.create("custom.op1", results=[i32])
        assert op.result.owner == op


# CHECK-LABEL: TEST: testBlockArgOwner
@run
def testBlockArgOwner():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    module = Module.parse(
        r"""
    func.func @foo(%arg0: f32) {
      return
    }""",
        ctx,
    )
    func = module.body.operations[0]
    block = func.regions[0].blocks[0]
    assert block.arguments[0].owner == block


# CHECK-LABEL: TEST: testValueIsInstance
@run
def testValueIsInstance():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    module = Module.parse(
        r"""
    func.func @foo(%arg0: f32) {
      %0 = "some_dialect.some_op"() : () -> f64
      return
    }""",
        ctx,
    )
    func = module.body.operations[0]
    assert BlockArgument.isinstance(func.regions[0].blocks[0].arguments[0])
    assert not OpResult.isinstance(func.regions[0].blocks[0].arguments[0])

    op = func.regions[0].blocks[0].operations[0]
    assert not BlockArgument.isinstance(op.results[0])
    assert OpResult.isinstance(op.results[0])


# CHECK-LABEL: TEST: testValueHash
@run
def testValueHash():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    module = Module.parse(
        r"""
    func.func @foo(%arg0: f32) -> f32 {
      %0 = "some_dialect.some_op"(%arg0) : (f32) -> f32
      return %0 : f32
    }""",
        ctx,
    )

    [func] = module.body.operations
    block = func.entry_block
    op, ret = block.operations
    assert hash(block.arguments[0]) == hash(op.operands[0])
    assert hash(op.result) == hash(ret.operands[0])


# CHECK-LABEL: TEST: testValueUses
@run
def testValueUses():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            value = Operation.create("custom.op1", results=[i32]).results[0]
            op1 = Operation.create("custom.op2", operands=[value])
            op2 = Operation.create("custom.op2", operands=[value])

    # CHECK: Use owner: "custom.op2"
    # CHECK: Use operand_number: 0
    # CHECK: Use owner: "custom.op2"
    # CHECK: Use operand_number: 0
    for use in value.uses:
        assert use.owner in [op1, op2]
        print(f"Use owner: {use.owner}")
        print(f"Use operand_number: {use.operand_number}")


# CHECK-LABEL: TEST: testValueReplaceAllUsesWith
@run
def testValueReplaceAllUsesWith():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            value = Operation.create("custom.op1", results=[i32]).results[0]
            op1 = Operation.create("custom.op2", operands=[value])
            op2 = Operation.create("custom.op2", operands=[value])
            value2 = Operation.create("custom.op3", results=[i32]).results[0]
            value.replace_all_uses_with(value2)

    assert len(list(value.uses)) == 0

    # CHECK: Use owner: "custom.op2"
    # CHECK: Use operand_number: 0
    # CHECK: Use owner: "custom.op2"
    # CHECK: Use operand_number: 0
    for use in value2.uses:
        assert use.owner in [op1, op2]
        print(f"Use owner: {use.owner}")
        print(f"Use operand_number: {use.operand_number}")


# CHECK-LABEL: TEST: testValueReplaceAllUsesWithExcept
@run
def testValueReplaceAllUsesWithExcept():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            value = Operation.create("custom.op1", results=[i32]).results[0]
            op1 = Operation.create("custom.op1", operands=[value])
            op2 = Operation.create("custom.op2", operands=[value])
            value2 = Operation.create("custom.op3", results=[i32]).results[0]
            value.replace_all_uses_except(value2, op1)

    assert len(list(value.uses)) == 1

    # CHECK: Use owner: "custom.op2"
    # CHECK: Use operand_number: 0
    for use in value2.uses:
        assert use.owner in [op2]
        print(f"Use owner: {use.owner}")
        print(f"Use operand_number: {use.operand_number}")

    # CHECK: Use owner: "custom.op1"
    # CHECK: Use operand_number: 0
    for use in value.uses:
        assert use.owner in [op1]
        print(f"Use owner: {use.owner}")
        print(f"Use operand_number: {use.operand_number}")


# CHECK-LABEL: TEST: testValueReplaceAllUsesWithMultipleExceptions
@run
def testValueReplaceAllUsesWithMultipleExceptions():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            value = Operation.create("custom.op1", results=[i32]).results[0]
            op1 = Operation.create("custom.op1", operands=[value])
            op2 = Operation.create("custom.op2", operands=[value])
            op3 = Operation.create("custom.op3", operands=[value])
            value2 = Operation.create("custom.op4", results=[i32]).results[0]

            # Replace all uses of `value` with `value2`, except for `op1` and `op2`.
            value.replace_all_uses_except(value2, [op1, op2])

    # After replacement, only `op3` should use `value2`, while `op1` and `op2` should still use `value`.
    assert len(list(value.uses)) == 2
    assert len(list(value2.uses)) == 1

    # CHECK: Use owner: "custom.op3"
    # CHECK: Use operand_number: 0
    for use in value2.uses:
        assert use.owner in [op3]
        print(f"Use owner: {use.owner}")
        print(f"Use operand_number: {use.operand_number}")

    # CHECK: Use owner: "custom.op2"
    # CHECK: Use operand_number: 0
    # CHECK: Use owner: "custom.op1"
    # CHECK: Use operand_number: 0
    for use in value.uses:
        assert use.owner in [op1, op2]
        print(f"Use owner: {use.owner}")
        print(f"Use operand_number: {use.operand_number}")


# CHECK-LABEL: TEST: testValuePrintAsOperand
@run
def testValuePrintAsOperand():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            value = Operation.create("custom.op1", results=[i32]).results[0]
            # CHECK: Value(%[[VAL1:.*]] = "custom.op1"() : () -> i32)
            print(value)

            value2 = Operation.create("custom.op2", results=[i32]).results[0]
            # CHECK: Value(%[[VAL2:.*]] = "custom.op2"() : () -> i32)
            print(value2)

            topFn = func.FuncOp("test", ([i32, i32], []))
            entry_block = Block.create_at_start(topFn.operation.regions[0], [i32, i32])

            with InsertionPoint(entry_block):
                value3 = Operation.create("custom.op3", results=[i32]).results[0]
                # CHECK: Value(%[[VAL3:.*]] = "custom.op3"() : () -> i32)
                print(value3)
                value4 = Operation.create("custom.op4", results=[i32]).results[0]
                # CHECK: Value(%[[VAL4:.*]] = "custom.op4"() : () -> i32)
                print(value4)
                func.ReturnOp([])

        # CHECK: %[[VAL1]]
        print(value.get_name())
        # CHECK: %[[VAL2]]
        print(value2.get_name())
        # CHECK: %[[VAL3]]
        print(value3.get_name())
        # CHECK: %[[VAL4]]
        print(value4.get_name())

        print("With AsmState")
        # CHECK-LABEL: With AsmState
        state = AsmState(topFn.operation, use_local_scope=True)
        # CHECK: %0
        print(value3.get_name(state=state))
        # CHECK: %1
        print(value4.get_name(state=state))

        print("With use_local_scope")
        # CHECK-LABEL: With use_local_scope
        # CHECK: %0
        print(value3.get_name(use_local_scope=True))
        # CHECK: %1
        print(value4.get_name(use_local_scope=True))

        # CHECK: %[[ARG0:.*]]
        print(entry_block.arguments[0].get_name())
        # CHECK: %[[ARG1:.*]]
        print(entry_block.arguments[1].get_name())

        # CHECK: module {
        # CHECK:   %[[VAL1]] = "custom.op1"() : () -> i32
        # CHECK:   %[[VAL2]] = "custom.op2"() : () -> i32
        # CHECK:   func.func @test(%[[ARG0]]: i32, %[[ARG1]]: i32) {
        # CHECK:     %[[VAL3]] = "custom.op3"() : () -> i32
        # CHECK:     %[[VAL4]] = "custom.op4"() : () -> i32
        # CHECK:     return
        # CHECK:   }
        # CHECK: }
        print(module)

        value2.owner.detach_from_parent()
        # CHECK: %0
        print(value2.get_name())


# CHECK-LABEL: TEST: testValueSetType
@run
def testValueSetType():
    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        i64 = IntegerType.get_signless(64)
        module = Module.create()
        with InsertionPoint(module.body):
            value = Operation.create("custom.op1", results=[i32]).results[0]
            # CHECK: Value(%[[VAL1:.*]] = "custom.op1"() : () -> i32)
            print(value)

            value.set_type(i64)
            # CHECK: Value(%[[VAL1]] = "custom.op1"() : () -> i64)
            print(value)

            # CHECK: %[[VAL1]] = "custom.op1"() : () -> i64
            print(value.owner)


# CHECK-LABEL: TEST: testValueCasters
@run
def testValueCasters():
    class NOPResult(OpResult):
        def __init__(self, v):
            super().__init__(v)

        def __str__(self):
            return super().__str__().replace(Value.__name__, NOPResult.__name__)

    class NOPValue(Value):
        def __init__(self, v):
            super().__init__(v)

        def __str__(self):
            return super().__str__().replace(Value.__name__, NOPValue.__name__)

    class NOPBlockArg(BlockArgument):
        def __init__(self, v):
            super().__init__(v)

        def __str__(self):
            return super().__str__().replace(Value.__name__, NOPBlockArg.__name__)

    @register_value_caster(IntegerType.static_typeid)
    def cast_int(v) -> Value:
        print("in caster", v.__class__.__name__)
        if isinstance(v, OpResult):
            return NOPResult(v)
        if isinstance(v, BlockArgument):
            return NOPBlockArg(v)
        elif isinstance(v, Value):
            return NOPValue(v)

    ctx = Context()
    ctx.allow_unregistered_dialects = True
    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            values = Operation.create("custom.op1", results=[i32, i32]).results
            # CHECK: in caster OpResult
            # CHECK: result 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("result", values[0].result_number, values[0])
            # CHECK: in caster OpResult
            # CHECK: result 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("result", values[1].result_number, values[1])

            # CHECK: results slice 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("results slice", values[:1][0].result_number, values[:1][0])

            value0, value1 = values
            # CHECK: in caster OpResult
            # CHECK: result 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("result", value0.result_number, values[0])
            # CHECK: in caster OpResult
            # CHECK: result 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("result", value1.result_number, values[1])

            op1 = Operation.create("custom.op2", operands=[value0, value1])
            # CHECK: "custom.op2"(%0#0, %0#1) : (i32, i32) -> ()
            print(op1)

            # CHECK: in caster Value
            # CHECK: operand 0 NOPValue(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("operand 0", op1.operands[0])
            # CHECK: in caster Value
            # CHECK: operand 1 NOPValue(%0:2 = "custom.op1"() : () -> (i32, i32))
            print("operand 1", op1.operands[1])

            # CHECK: in caster BlockArgument
            # CHECK: in caster BlockArgument
            @func.FuncOp.from_py_func(i32, i32)
            def reduction(arg0, arg1):
                # CHECK: as func arg 0 NOPBlockArg
                print("as func arg", arg0.arg_number, arg0.__class__.__name__)
                # CHECK: as func arg 1 NOPBlockArg
                print("as func arg", arg1.arg_number, arg1.__class__.__name__)

            # CHECK: args slice 0 NOPBlockArg(<block argument> of type 'i32' at index: 0)
            print(
                "args slice",
                reduction.func_op.arguments[:1][0].arg_number,
                reduction.func_op.arguments[:1][0],
            )

    try:

        @register_value_caster(IntegerType.static_typeid)
        def dont_cast_int_shouldnt_register(v):
            ...

    except RuntimeError as e:
        # CHECK: Value caster is already registered: {{.*}}cast_int
        print(e)

    @register_value_caster(IntegerType.static_typeid, replace=True)
    def dont_cast_int(v) -> OpResult:
        assert isinstance(v, OpResult)
        print("don't cast", v.result_number, v)
        return v

    with Location.unknown(ctx):
        i32 = IntegerType.get_signless(32)
        module = Module.create()
        with InsertionPoint(module.body):
            # CHECK: don't cast 0 Value(%0 = "custom.op1"() : () -> i32)
            new_value = Operation.create("custom.op1", results=[i32]).result
            # CHECK: result 0 Value(%0 = "custom.op1"() : () -> i32)
            print("result", new_value.result_number, new_value)

            # CHECK: don't cast 0 Value(%1 = "custom.op2"() : () -> i32)
            new_value = Operation.create("custom.op2", results=[i32]).results[0]
            # CHECK: result 0 Value(%1 = "custom.op2"() : () -> i32)
            print("result", new_value.result_number, new_value)