import re
from collections.abc import Sequence

import pytest
from conftest import assert_print_op

from xdsl.builder import ImplicitBuilder
from xdsl.context import MLContext
from xdsl.dialects import test
from xdsl.dialects.arith import AddiOp, Arith, ConstantOp, MuliOp
from xdsl.dialects.builtin import (
    Builtin,
    IndexType,
    IntegerAttr,
    IntegerType,
    ModuleOp,
    StringAttr,
    UnitAttr,
    i32,
    i64,
)
from xdsl.ir import Block, Operation, SSAValue
from xdsl.irdl import BaseAttr
from xdsl.parser import Parser
from xdsl.pattern_rewriter import (
    GreedyRewritePatternApplier,
    PatternRewriter,
    PatternRewriterListener,
    PatternRewriteWalker,
    RewritePattern,
    TypeConversionPattern,
    attr_constr_rewrite_pattern,
    attr_type_rewrite_pattern,
    op_type_rewrite_pattern,
)
from xdsl.rewriter import BlockInsertPoint, InsertPoint


def rewrite_and_compare(
    prog: str,
    expected_prog: str,
    walker: PatternRewriteWalker,
    *,
    op_inserted: int = 0,
    op_removed: int = 0,
    op_modified: int = 0,
    op_replaced: int = 0,
    block_created: int = 0,
    expect_rewrite: bool = True,
):
    ctx = MLContext(allow_unregistered=True)
    ctx.load_dialect(Builtin)
    ctx.load_dialect(Arith)
    ctx.load_dialect(test.Test)

    num_op_inserted = 0
    num_op_removed = 0
    num_op_modified = 0
    num_op_replaced = 0
    num_block_created = 0

    def op_inserted_handler(op: Operation):
        nonlocal num_op_inserted
        num_op_inserted += 1

    def op_removed_handler(op: Operation):
        nonlocal num_op_removed
        num_op_removed += 1

    def op_modified_handler(op: Operation):
        nonlocal num_op_modified
        num_op_modified += 1

    def op_replaced_handler(op: Operation, values: Sequence[SSAValue | None]):
        nonlocal num_op_replaced
        num_op_replaced += 1

    def block_created_handler(block: Block):
        nonlocal num_block_created
        num_block_created += 1

    listener = PatternRewriterListener()
    listener.operation_insertion_handler = [op_inserted_handler]
    listener.operation_removal_handler = [op_removed_handler]
    listener.operation_modification_handler = [op_modified_handler]
    listener.operation_replacement_handler = [op_replaced_handler]
    listener.block_creation_handler = [block_created_handler]

    parser = Parser(ctx, prog)
    module = parser.parse_module()

    walker.listener = listener
    did_rewrite = walker.rewrite_module(module)

    assert_print_op(module, expected_prog, None)

    assert num_op_inserted == op_inserted
    assert num_op_removed == op_removed
    assert num_op_modified == op_modified
    assert num_op_replaced == op_replaced
    assert num_block_created == block_created
    assert did_rewrite == expect_rewrite


def test_non_recursive_rewrite():
    """Test a simple non-recursive rewrite"""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 43 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class RewriteConst(RewritePattern):
        def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter):
            if isinstance(op, ConstantOp):
                new_constant = ConstantOp.from_int_and_width(43, i32)
                rewriter.replace_matched_op([new_constant])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(RewriteConst(), apply_recursively=False),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
        op_modified=2,
    )


def test_non_recursive_rewrite_reversed():
    """Test a simple non-recursive rewrite with reverse walk order."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 43 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class RewriteConst(RewritePattern):
        def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter):
            if isinstance(op, ConstantOp):
                new_constant = ConstantOp.from_int_and_width(43, i32)
                rewriter.replace_matched_op([new_constant])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(
            RewriteConst(), apply_recursively=False, walk_reverse=True
        ),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
        op_modified=2,
    )


def test_op_type_rewrite_pattern_method_decorator():
    """Test op_type_rewrite_pattern decorator on methods."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 43 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class RewriteConst(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __op__: ConstantOp, rewriter: PatternRewriter):
            rewriter.replace_matched_op(ConstantOp.from_int_and_width(43, i32))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(RewriteConst(), apply_recursively=False),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
        op_modified=2,
    )


def test_op_type_rewrite_pattern_union_type():
    """Test op_type_rewrite_pattern decorator on static functions."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %2 = "test"(%0, %1) : (i32, i32) -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %2 = "test"(%0, %1) : (i32, i32) -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(
            self, __op__: ConstantOp | AddiOp, rewriter: PatternRewriter
        ):
            rewriter.replace_matched_op(ConstantOp.from_int_and_width(42, i32))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=2,
        op_removed=2,
        op_replaced=2,
        op_modified=4,
    )


def test_recursive_rewriter():
    """Test recursive walks on operations created by rewrites."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %1 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %2 = "arith.addi"(%0, %1) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %3 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %4 = "arith.addi"(%2, %3) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %5 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %6 = "arith.addi"(%4, %5) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %7 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %8 = "arith.addi"(%6, %7) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: ConstantOp, rewriter: PatternRewriter):
            if not isinstance(op_val := op.value, IntegerAttr):
                return
            val = op_val.value.data
            if val == 0 or val == 1:
                return
            constant_op = ConstantOp(IntegerAttr.from_int_and_width(val - 1, 32), i32)
            constant_one = ConstantOp(IntegerAttr.from_int_and_width(1, 32), i32)
            add_op = AddiOp(constant_op, constant_one)
            rewriter.replace_matched_op([constant_op, constant_one, add_op])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=True),
        op_inserted=12,
        op_removed=4,
        op_replaced=4,
        op_modified=3,
    )


def test_recursive_rewriter_reversed():
    """Test recursive walks on operations created by rewrites, in reverse walk order."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %1 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %2 = "arith.addi"(%0, %1) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %3 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %4 = "arith.addi"(%2, %3) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %5 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %6 = "arith.addi"(%4, %5) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  %7 = "arith.constant"() <{value = 1 : i32}> : () -> i32
  %8 = "arith.addi"(%6, %7) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: ConstantOp, rewriter: PatternRewriter):
            if not isinstance(op_val := op.value, IntegerAttr):
                return
            val = op_val.value.data
            if val == 0 or val == 1:
                return
            constant_op = ConstantOp(IntegerAttr.from_int_and_width(val - 1, 32), i32)
            constant_one = ConstantOp(IntegerAttr.from_int_and_width(1, 32), i32)
            add_op = AddiOp(constant_op, constant_one)
            rewriter.replace_matched_op([constant_op, constant_one, add_op])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=True, walk_reverse=True),
        op_inserted=12,
        op_removed=4,
        op_replaced=4,
        op_modified=3,
    )


def test_greedy_rewrite_pattern_applier():
    """Test GreedyRewritePatternApplier."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 43 : i32}> : () -> i32
  %1 = "arith.muli"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class ConstantRewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __op__: ConstantOp, rewriter: PatternRewriter):
            rewriter.replace_matched_op([ConstantOp.from_int_and_width(43, i32)])

    class AddiRewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: AddiOp, rewriter: PatternRewriter):
            rewriter.replace_matched_op([MuliOp(op.lhs, op.rhs)])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(
            GreedyRewritePatternApplier([ConstantRewrite(), AddiRewrite()]),
            apply_recursively=False,
        ),
        op_inserted=2,
        op_removed=2,
        op_replaced=2,
        op_modified=2,
    )


def test_insert_op_before_matched_op():
    """Test rewrites where operations are inserted before the matched operation."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  %1 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __cst__: ConstantOp, rewriter: PatternRewriter):
            new_cst = ConstantOp.from_int_and_width(42, i32)

            rewriter.insert_op_before_matched_op(new_cst)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
    )


def test_insert_op_at_start():
    """Test rewrites where operations are inserted with a given position."""

    prog = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    expected = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
    %1 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            new_cst = ConstantOp.from_int_and_width(42, i32)

            rewriter.insert_op(new_cst, InsertPoint.at_start(op.regions[0].blocks[0]))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
    )


def test_insert_op_before():
    """Test rewrites where operations are inserted before a given operation."""

    prog = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    expected = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
    %1 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            new_cst = ConstantOp.from_int_and_width(42, i32)

            first_op = op.regions[0].block.ops.first
            assert first_op is not None
            rewriter.insert_op(new_cst, InsertPoint.before(first_op))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
    )


def test_insert_op_after():
    """Test rewrites where operations are inserted after a given operation."""

    prog = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    expected = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
    %1 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            new_cst = ConstantOp.from_int_and_width(42, i32)

            first_op = op.regions[0].block.ops.first
            assert first_op is not None
            rewriter.insert_op(new_cst, InsertPoint.after(first_op))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
    )


def test_insert_op_after_matched_op():
    """Test rewrites where operations are inserted after a given operation."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  %1 = "arith.constant"() <{value = 42 : i32}> : () -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __cst__: ConstantOp, rewriter: PatternRewriter):
            new_cst = ConstantOp.from_int_and_width(42, i32)

            rewriter.insert_op_after_matched_op(new_cst)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
    )


def test_insert_op_after_matched_op_reversed():
    """Test rewrites where operations are inserted after a given operation."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  %1 = "arith.constant"() <{value = 42 : i32}> : () -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __cst__: ConstantOp, rewriter: PatternRewriter):
            new_cst = ConstantOp.from_int_and_width(42, i32)

            rewriter.insert_op_after_matched_op(new_cst)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False, walk_reverse=True),
        op_inserted=1,
    )


def test_operation_deletion():
    """Test rewrites where SSA values are deleted."""

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
^0:
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __op__: ConstantOp, rewriter: PatternRewriter):
            rewriter.erase_matched_op()

    rewrite_and_compare(prog, expected, PatternRewriteWalker(Rewrite()), op_removed=1)


def test_operation_deletion_reversed():
    """
    Test rewrites where SSA values are deleted.
    They have to be deleted in order for the rewrite to not fail.
    """

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    expected = """"builtin.module"() ({
^0:
}) : () -> ()"""

    class EraseAll(RewritePattern):
        def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter):
            if not isinstance(op, ModuleOp):
                rewriter.erase_matched_op()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(EraseAll(), walk_reverse=True),
        op_removed=2,
    )


def test_operation_deletion_failure():
    """Test rewrites where SSA values are deleted with still uses."""

    ctx = MLContext()
    ctx.load_dialect(Builtin)
    ctx.load_dialect(Arith)

    prog = """"builtin.module"() ({
  %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  %1 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, __op__: ConstantOp, rewriter: PatternRewriter):
            rewriter.erase_matched_op()

    parser = Parser(ctx, prog)
    module = parser.parse_module()
    walker = PatternRewriteWalker(Rewrite())

    # Check that the rewrite fails
    with pytest.raises(
        Exception,
        match="Attempting to delete SSA value that still has uses of result of operation",
    ):
        walker.rewrite_module(module)


def test_delete_inner_op():
    """Test rewrites where an operation inside a region of the matched op is deleted."""

    prog = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    expected = """"builtin.module"() ({
  "test.op"() ({
  ^0:
  }) : () -> ()
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            first_op = op.regions[0].block.ops.first

            assert first_op is not None
            rewriter.erase_op(first_op)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_removed=1,
    )


def test_replace_inner_op():
    """Test rewrites where an operation inside a region of the matched op is deleted."""

    prog = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 5 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    expected = """"builtin.module"() ({
  "test.op"() ({
    %0 = "arith.constant"() <{value = 42 : i32}> : () -> i32
  }) : () -> ()
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            first_op = op.regions[0].block.ops.first

            assert first_op is not None
            rewriter.replace_op(first_op, [ConstantOp.from_int_and_width(42, i32)])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
    )


def test_block_argument_type_change():
    """Test the modification of a block argument type."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() ({
  ^0(%1 : !test.type<"int">):
    %2 = "test.op"() : () -> !test.type<"int">
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() ({
  ^0(%1 : i64):
    %2 = "test.op"() : () -> !test.type<"int">
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                rewriter.modify_value_type(matched_op.regs[0].blocks[0].args[0], i64)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_modified=1,
    )


def test_block_argument_erasure():
    """Test the erasure of a block argument."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() ({
  ^0(%1 : !test.type<"int">):
    %2 = "test.op"() : () -> !test.type<"int">
  }, {
  ^0:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() ({
    %1 = "test.op"() : () -> !test.type<"int">
  }, {
  ^0:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                rewriter.erase_block_argument(matched_op.regs[0].blocks[0].args[0])

    rewrite_and_compare(
        prog, expected, PatternRewriteWalker(Rewrite(), apply_recursively=False)
    )


def test_block_argument_insertion():
    """Test the insertion of a block argument."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() ({
    %1 = "test.op"() : () -> !test.type<"int">
  }, {
  ^0:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() ({
  ^0(%1 : !test.type<"int">):
    %2 = "test.op"() : () -> !test.type<"int">
  }, {
  ^1:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                rewriter.insert_block_argument(
                    matched_op.regs[0].blocks[0], 0, test.TestType("int")
                )

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_inline_block_before_matched_op():
    """Test the inlining of a block before the matched operation."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
  ^0:
    %2 = "test.op"() : () -> !test.type<"int">
  }, {
  ^1:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() : () -> !test.type<"int">
  %2 = "test.op"() ({
  }, {
  ^0:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                rewriter.inline_block_before_matched_op(matched_op.regs[0].blocks[0])

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_inline_block_before():
    """Test the inlining of a block before an operation."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
    %1 = "test.op"() ({
      %1 = "test.op"() : () -> !test.type<"int">
    }, {
    ^1:
    }) : () -> !test.type<"int">
  }, {
  ^2:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
    %2 = "test.op"() : () -> !test.type<"int">
    %3 = "test.op"() ({
    }, {
    ^0:
    }) : () -> !test.type<"int">
  }, {
  ^1:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                first_op = matched_op.regs[0].blocks[0].first_op

                if isinstance(first_op, test.TestOp):
                    inner_block = first_op.regs[0].blocks[0]
                    rewriter.inline_block(inner_block, InsertPoint.before(first_op))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_inline_block_at_before_when_op_is_matched_op():
    """Test the inlining of a block before an operation, being the matched one."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
  ^0:
    %2 = "test.op"() : () -> !test.type<"int">
  }, {
  ^1:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() : () -> !test.type<"int">
  %2 = "test.op"() ({
  }, {
  ^0:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                rewriter.inline_block(
                    matched_op.regs[0].blocks[0], InsertPoint.before(matched_op)
                )

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_inline_block_before_with_args():
    """Test the inlining of a block before an operation."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
  ^0(%arg0 : !test.type<"int">):
    %1 = "test.op"() ({
    ^1(%arg1 : !test.type<"int">):
      %1 = "test.op"(%arg1) : (!test.type<"int">) -> !test.type<"int">
    }, {
    ^2:
    }) : () -> !test.type<"int">
  }, {
  ^3:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
  ^0(%arg0 : !test.type<"int">):
    %2 = "test.op"(%arg0) : (!test.type<"int">) -> !test.type<"int">
    %3 = "test.op"() ({
    }, {
    ^1:
    }) : () -> !test.type<"int">
  }, {
  ^2:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                outer_block = matched_op.regs[0].blocks[0]
                first_op = outer_block.first_op

                if isinstance(first_op, test.TestOp):
                    inner_block = first_op.regs[0].blocks[0]
                    rewriter.inline_block(
                        inner_block, InsertPoint.before(first_op), outer_block.args
                    )

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_inline_block_after():
    """Test the inlining of a block after an operation."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
    %1 = "test.op"() ({
      %1 = "test.op"() : () -> !test.type<"int">
    }, {
      ^1:
    }) : () -> !test.type<"int">
  }, {
  ^2:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
    %2 = "test.op"() ({
    }, {
    ^0:
    }) : () -> !test.type<"int">
    %3 = "test.op"() : () -> !test.type<"int">
  }, {
  ^1:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                first_op = matched_op.regs[0].blocks[0].first_op

                if first_op is not None and isinstance(first_op, test.TestOp):
                    if first_op.regs and first_op.regs[0].blocks:
                        inner_block = first_op.regs[0].blocks[0]
                        rewriter.inline_block(inner_block, InsertPoint.after(first_op))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_inline_block_after_matched():
    """Test the inlining of a block after an operation."""

    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
    %1 = "test.op"() ({
      %1 = "test.op"() : () -> !test.type<"int">
    }, {
    ^1:
    }) : () -> !test.type<"int">
  }, {
  ^2:
  }) : () -> !test.type<"int">
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> !test.type<"int">
  %1 = "test.op"() ({
    %2 = "test.op"() ({
    }, {
    ^0:
    }) : () -> !test.type<"int">
  }, {
  ^1:
  }) : () -> !test.type<"int">
  %3 = "test.op"() : () -> !test.type<"int">
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.regs and matched_op.regs[0].blocks:
                first_op = matched_op.regs[0].block.first_op

                if first_op is not None and isinstance(first_op, test.TestOp):
                    if first_op.regs and first_op.regs[0].blocks:
                        inner_block = first_op.regs[0].blocks[0]
                        rewriter.inline_block(
                            inner_block, InsertPoint.after(matched_op)
                        )

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
    )


def test_move_region_contents_to_new_regions():
    """Test moving a region outside of a region."""

    prog = """\
"builtin.module"() ({
  "test.op"() ({
    %0 = "test.op"() : () -> !test.type<"int">
    %1 = "test.op"() ({
    ^0:
      %2 = "test.op"() : () -> !test.type<"int">
    }) : () -> !test.type<"int">
  }) : () -> ()
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  "test.op"() ({
    %0 = "test.op"() : () -> !test.type<"int">
    %1 = "test.op"() ({
    }) : () -> !test.type<"int">
    %2 = "test.op"() ({
      %3 = "test.op"() : () -> !test.type<"int">
    }) : () -> !test.type<"int">
  }) : () -> ()
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            # Match the toplevel test.op
            if not isinstance(op.parent_op(), ModuleOp):
                return
            ops_iter = iter(op.regions[0].block.ops)

            _ = next(ops_iter)  # skip first op
            old_op = next(ops_iter)
            assert isinstance(old_op, test.TestOp)
            new_region = rewriter.move_region_contents_to_new_regions(old_op.regions[0])
            res_types = old_op.result_types
            new_op = test.TestOp.create(result_types=res_types, regions=[new_region])
            rewriter.insert_op(new_op, InsertPoint.after(old_op))

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=1,
    )


def test_insert_same_block():
    """Test rewriter on ops without results"""
    prog = """\
"builtin.module"() ({
  %0 = "test.op"() {label = "a"} : () -> i32
  %1 = "test.op"() {label = "b"} : () -> i32
  %2 = "test.op"() {label = "c"} : () -> i32
  "func.return"() : () -> ()
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() {label = "alloc"} : () -> i32
  %1 = "test.op"() {label = "a"} : () -> i32
  "test.op"(%0) {label = "init"} : (i32) -> ()
  %2 = "test.op"() {label = "c"} : () -> i32
  "test.op"(%0) {label = "dealloc"} : (i32) -> ()
  "func.return"() : () -> ()
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if op.attributes["label"] != StringAttr("b"):
                return

            block = op.parent

            if block is None:
                return

            last_op = block.last_op
            assert last_op is not None

            alloc = test.TestOp.build(
                operands=((),),
                attributes={"label": StringAttr("alloc")},
                regions=((),),
                result_types=((i32,),),
            )
            init = test.TestOp.build(
                operands=(alloc.res,),
                attributes={"label": StringAttr("init")},
                regions=((),),
                result_types=((),),
            )
            dealloc = test.TestOp.build(
                operands=(alloc.res,),
                attributes={"label": StringAttr("dealloc")},
                regions=((),),
                result_types=((),),
            )

            # Allocate before first use
            rewriter.insert_op(alloc, InsertPoint.at_start(block))
            # Deallocate after last use
            rewriter.insert_op(dealloc, InsertPoint.before(last_op))
            # Init instead of creating, and replace result with allocated value
            rewriter.replace_matched_op(init, alloc.res)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=3,
        op_removed=1,
        op_replaced=1,
    )


def test_inline_region_before():
    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
^0:
  "test.op"() ({
    ^1:
    %1 = "test.op"() : () -> f32
    ^2:
    %2 = "test.op"() : () -> f64
  }) {label = "а"} : () -> ()
  %2 = "test.op"() : () -> i64
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
^0:
  %1 = "test.op"() : () -> f32
^1:
  %2 = "test.op"() : () -> f64
^2:
  %3 = "test.op"() : () -> i64
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if op.attributes.get("label") != StringAttr("а"):
                return
            if op.parent is None:
                return

            rewriter.inline_region(op.regions[0], BlockInsertPoint.before(op.parent))
            rewriter.erase_matched_op()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=0,
        op_removed=1,
    )


def test_inline_region_after():
    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
  "test.op"() ({
    ^1:
    %1 = "test.op"() : () -> f32
    ^2:
    %2 = "test.op"() : () -> f64
  }) {label = "а"} : () -> ()
^0:
  %2 = "test.op"() : () -> i64
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
^0:
  %1 = "test.op"() : () -> f32
^1:
  %2 = "test.op"() : () -> f64
^2:
  %3 = "test.op"() : () -> i64
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if op.attributes.get("label") != StringAttr("а"):
                return
            if op.parent is None:
                return

            rewriter.inline_region(op.regions[0], BlockInsertPoint.after(op.parent))
            rewriter.erase_matched_op()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=0,
        op_removed=1,
    )


def test_inline_region_at_start():
    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
^0:
  "test.op"() ({
    ^1:
    %1 = "test.op"() : () -> f32
    ^2:
    %2 = "test.op"() : () -> f64
  }) {label = "а"} : () -> ()
  %2 = "test.op"() : () -> i64
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> f32
^0:
  %1 = "test.op"() : () -> f64
^1:
  %2 = "test.op"() : () -> i32
^2:
  %3 = "test.op"() : () -> i64
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if op.attributes.get("label") != StringAttr("а"):
                return
            parent_region = op.parent_region()
            if parent_region is None:
                return

            rewriter.inline_region(
                op.regions[0], BlockInsertPoint.at_start(parent_region)
            )
            rewriter.erase_matched_op()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=0,
        op_removed=1,
    )


def test_inline_region_at_end():
    prog = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
^0:
  "test.op"() ({
    ^1:
    %1 = "test.op"() : () -> f32
    ^2:
    %2 = "test.op"() : () -> f64
  }) {label = "а"} : () -> ()
  %2 = "test.op"() : () -> i64
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  %0 = "test.op"() : () -> i32
^0:
  %1 = "test.op"() : () -> i64
^1:
  %2 = "test.op"() : () -> f32
^2:
  %3 = "test.op"() : () -> f64
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if op.attributes.get("label") != StringAttr("а"):
                return
            parent_region = op.parent_region()
            if parent_region is None:
                return

            rewriter.inline_region(
                op.regions[0], BlockInsertPoint.at_end(parent_region)
            )
            rewriter.erase_matched_op()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=0,
        op_removed=1,
    )


def test_erased_ssavalue():
    prog = """\
builtin.module {
  "test.op"() ({
    %0 = "test.op"() : () -> i32
    "test.op"(%0) : (i32) -> ()
  }) : () -> ()
}
  """

    expected = """\
"builtin.module"() ({
  "test.op"() ({
  ^0:
  }) : () -> ()
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if op.results or op.operands:
                rewriter.erase_matched_op(safe_erase=False)

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=True),
        op_removed=2,
    )


def test_pattern_rewriter_as_op_builder():
    """Test that the PatternRewriter works as an OpBuilder."""
    prog = """
"builtin.module"() ({
  "test.op"() : () -> ()
  "test.op"() {nomatch} : () -> ()
  "test.op"() : () -> ()
}) : () -> ()"""
    expected = """
"builtin.module"() ({
  "test.op"() {inserted} : () -> ()
  "test.op"() {replaced} : () -> ()
  "test.op"() {nomatch} : () -> ()
  "test.op"() {inserted} : () -> ()
  "test.op"() {replaced} : () -> ()
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if "nomatch" in op.attributes:
                return
            with ImplicitBuilder(rewriter):
                test.TestOp.create(attributes={"inserted": UnitAttr()})
            rewriter.replace_matched_op(
                test.TestOp.create(attributes={"replaced": UnitAttr()})
            )

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=4,
        op_removed=2,
        op_replaced=2,
    )


def test_type_conversion():
    """Test rewriter on ops without results"""
    prog = """\
"builtin.module"() ({
  "func.func"() ({
  ^0(%arg : i32):
  }) : () -> ()
  %0 = "test.op"() {nested = memref<*xi32>} : () -> i32
  %1 = "test.op"() {type = () -> memref<*xi32>} : () -> f32
  %2 = "test.op"() <{prop = memref<*xi32>}> : () -> f32
  %3 = "test.op"(%0, %1) : (i32, f32) -> memref<*xi32>
  %4 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  "func.return"() : () -> ()
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  "func.func"() ({
  ^0(%arg : index):
  }) : () -> ()
  %0 = "test.op"() {nested = memref<*xindex>} : () -> index
  %1 = "test.op"() {type = () -> memref<*xindex>} : () -> f32
  %2 = "test.op"() <{prop = memref<*xindex>}> : () -> f32
  %3 = "test.op"(%0, %1) : (index, f32) -> memref<*xindex>
  %4 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
  "func.return"() : () -> ()
}) : () -> ()
"""

    class Rewrite(TypeConversionPattern):
        @attr_type_rewrite_pattern
        def convert_type(self, typ: IntegerType) -> IndexType:
            return IndexType()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(recursive=True), apply_recursively=False),
        op_inserted=5,
        op_removed=5,
        op_replaced=5,
        op_modified=5,
    )
    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(recursive=True), apply_recursively=True),
        op_inserted=5,
        op_removed=5,
        op_replaced=5,
        op_modified=5,
    )
    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(
            Rewrite(recursive=True), apply_recursively=False, walk_reverse=True
        ),
        op_inserted=5,
        op_removed=5,
        op_replaced=5,
        op_modified=5,
    )

    non_rec_expected = """\
"builtin.module"() ({
  "func.func"() ({
  ^0(%arg : index):
  }) : () -> ()
  %0 = "test.op"() {nested = memref<*xi32>} : () -> index
  %1 = "test.op"() {type = () -> memref<*xi32>} : () -> f32
  %2 = "test.op"() <{prop = memref<*xi32>}> : () -> f32
  %3 = "test.op"(%0, %1) : (index, f32) -> memref<*xi32>
  %4 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
  "func.return"() : () -> ()
}) : () -> ()
"""

    rewrite_and_compare(
        prog,
        non_rec_expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_inserted=2,
        op_removed=2,
        op_replaced=2,
        op_modified=4,
    )
    rewrite_and_compare(
        prog,
        non_rec_expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=True),
        op_inserted=2,
        op_removed=2,
        op_replaced=2,
        op_modified=4,
    )
    rewrite_and_compare(
        prog,
        non_rec_expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False, walk_reverse=True),
        op_inserted=2,
        op_removed=2,
        op_replaced=2,
        op_modified=4,
    )

    expected = """\
"builtin.module"() ({
  "func.func"() ({
  ^0(%arg : i32):
  }) : () -> ()
  %0 = "test.op"() {nested = memref<*xindex>} : () -> index
  %1 = "test.op"() {type = () -> memref<*xindex>} : () -> f32
  %2 = "test.op"() <{prop = memref<*xindex>}> : () -> f32
  %3 = "test.op"(%0, %1) : (index, f32) -> memref<*xindex>
  %4 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> i32
  "func.return"() : () -> ()
}) : () -> ()
"""

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(
            Rewrite(ops=(test.TestOp,), recursive=True), apply_recursively=False
        ),
        op_inserted=4,
        op_removed=4,
        op_replaced=4,
        op_modified=4,
    )
    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(
            Rewrite(ops=(test.TestOp,), recursive=True), apply_recursively=True
        ),
        op_inserted=4,
        op_removed=4,
        op_replaced=4,
        op_modified=4,
    )
    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(
            Rewrite(ops=(test.TestOp,), recursive=True),
            apply_recursively=False,
            walk_reverse=True,
        ),
        op_inserted=4,
        op_removed=4,
        op_replaced=4,
        op_modified=4,
    )

    class RewriteMaybe(TypeConversionPattern):
        @attr_type_rewrite_pattern
        def convert_type(self, typ: IntegerType) -> IndexType | None:
            return IndexType() if typ.width.data >= 8 else None

    prog = """\
"builtin.module"() ({
  "func.func"() ({
  ^0(%arg : i6):
  }) : () -> ()
  %0 = "test.op"() {nested = memref<*xi4>} : () -> i6
  %1 = "test.op"() {type = () -> memref<*xi32>} : () -> f32
  %2 = "test.op"() <{prop = memref<*xi4>}> : () -> i32
  %3 = "test.op"(%0, %1) : (i6, f32) -> memref<*xi32>
  %4 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i6, i6) -> i6
  "func.return"() : () -> ()
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
  "func.func"() ({
  ^0(%arg : i6):
  }) : () -> ()
  %0 = "test.op"() {nested = memref<*xi4>} : () -> i6
  %1 = "test.op"() {type = () -> memref<*xindex>} : () -> f32
  %2 = "test.op"() <{prop = memref<*xi4>}> : () -> index
  %3 = "test.op"(%0, %1) : (i6, f32) -> memref<*xindex>
  %4 = "arith.addi"(%0, %0) <{overflowFlags = #arith.overflow<none>}> : (i6, i6) -> i6
  "func.return"() : () -> ()
}) : () -> ()
"""

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(RewriteMaybe(recursive=True), apply_recursively=True),
        op_inserted=3,
        op_removed=3,
        op_replaced=3,
        op_modified=1,
    )

    prog = """\
"builtin.module"() ({
  "test.op"() {dict_nest = {hello = i32}} : () -> ()
}) : () -> ()
"""

    expected_recursive = """
"builtin.module"() ({
  "test.op"() {dict_nest = {hello = index}} : () -> ()
}) : () -> ()
"""

    rewrite_and_compare(
        prog,
        prog,
        PatternRewriteWalker(Rewrite(recursive=False)),
        expect_rewrite=False,
    )

    rewrite_and_compare(
        prog,
        expected_recursive,
        PatternRewriteWalker(Rewrite(recursive=True)),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
    )


def test_recursive_type_conversion_in_regions():
    prog = """\
"builtin.module"() ({
  "func.func"() <{function_type = (memref<2x4xui16>) -> (), sym_name = "main", sym_visibility = "private"}> ({
  ^bb0(%arg0 : memref<2x4xui16>):
    "func.return"() : () -> ()
  }) : () -> ()
}) : () -> ()
"""
    expected_prog = """\
"builtin.module"() ({
  "func.func"() <{function_type = (memref<2x4xindex>) -> (), sym_name = "main", sym_visibility = "private"}> ({
  ^0(%arg0 : memref<2x4xindex>):
    "func.return"() : () -> ()
  }) : () -> ()
}) : () -> ()
"""

    class IndexConversion(TypeConversionPattern):
        @attr_type_rewrite_pattern
        def convert_type(self, typ: IntegerType) -> IndexType:
            return IndexType()

    rewrite_and_compare(
        prog,
        expected_prog,
        PatternRewriteWalker(IndexConversion(recursive=True)),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
        op_modified=1,
    )


def test_no_change():
    """Test that doing nothing successfully does not report doing something."""

    prog = """\
"builtin.module"() ({
}) : () -> ()
"""

    expected = """\
"builtin.module"() ({
}) : () -> ()
"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            return

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        expect_rewrite=False,
    )


def test_error():
    prog = """\
builtin.module {
  "test.op"() {erroneous = false} : () -> ()
  "test.op"() : () -> ()
  "test.op"() {erroneous = true} : () -> ()
  "test.op"() : () -> ()
}
"""
    expected = """\
Error while applying pattern: Expected operation to not be erroneous!

"builtin.module"() ({
  "test.op"() {erroneous = false} : () -> ()
  "test.op"() : () -> ()
  "test.op"() {erroneous = true} : () -> ()
  ^^^^^^^^^--------------------------------------------------------------
  | Error while applying pattern: Expected operation to not be erroneous!
  -----------------------------------------------------------------------
  "test.op"() : () -> ()
}) : () -> ()

"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, matched_op: test.TestOp, rewriter: PatternRewriter):
            if matched_op.attributes.get(
                "erroneous", IntegerAttr.from_int_and_width(0, 1)
            ) == IntegerAttr.from_int_and_width(1, 1):
                raise ValueError("Expected operation to not be erroneous!")
            return

    ctx = MLContext(allow_unregistered=True)
    ctx.load_dialect(Builtin)
    ctx.load_dialect(Arith)
    ctx.load_dialect(test.Test)

    parser = Parser(ctx, prog)
    module = parser.parse_module()

    walker = PatternRewriteWalker(Rewrite())
    with pytest.raises(ValueError, match=re.escape(expected)):
        walker.rewrite_module(module)


def test_attr_constr_rewrite_pattern():
    prog = """\
"builtin.module"() ({
  "func.func"() <{function_type = (memref<2x4xui16>) -> (), sym_name = "main", sym_visibility = "private"}> ({
  ^bb0(%arg0 : memref<2x4xui16>):
    "func.return"() : () -> ()
  }) : () -> ()
}) : () -> ()
"""
    expected_prog = """\
"builtin.module"() ({
  "func.func"() <{function_type = (memref<2x4xindex>) -> (), sym_name = "main", sym_visibility = "private"}> ({
  ^0(%arg0 : memref<2x4xindex>):
    "func.return"() : () -> ()
  }) : () -> ()
}) : () -> ()
"""

    class IndexConversion(TypeConversionPattern):
        @attr_constr_rewrite_pattern(BaseAttr(IntegerType))
        def convert_type(self, typ: IntegerType) -> IndexType:
            return IndexType()

    rewrite_and_compare(
        prog,
        expected_prog,
        PatternRewriteWalker(IndexConversion(recursive=True)),
        op_inserted=1,
        op_removed=1,
        op_replaced=1,
        op_modified=1,
    )


def test_pattern_rewriter_erase_op_with_region():
    """Test that erasing an operation with a region works correctly."""
    prog = """
"builtin.module"() ({
  "test.op"() ({
    "test.op"() {error_if_matching} : () -> ()
  }): () -> ()
}) : () -> ()"""
    expected = """
"builtin.module"() ({
^0:
}) : () -> ()"""

    class Rewrite(RewritePattern):
        @op_type_rewrite_pattern
        def match_and_rewrite(self, op: test.TestOp, rewriter: PatternRewriter):
            if "error_if_matching" in op.attributes:
                raise Exception("operation that is supposed to be deleted was matched")
            assert not op.attributes
            rewriter.erase_matched_op()

    rewrite_and_compare(
        prog,
        expected,
        PatternRewriteWalker(Rewrite(), apply_recursively=False),
        op_removed=1,
    )
