Pattern rewriter

@dataclass(eq=False)
class PatternRewriterListener(BuilderListener):
    """A listener for pattern rewriter events."""

    operation_removal_handler: list[Callable[[Operation], None]] = field(
        default_factory=list[Callable[[Operation], None]], kw_only=True
    )
    """Callbacks that are called when an operation is removed."""

    operation_modification_handler: list[Callable[[Operation], None]] = field(
        default_factory=list[Callable[[Operation], None]], kw_only=True
    )
    """Callbacks that are called when an operation is modified."""

    operation_replacement_handler: list[
        Callable[[Operation, Sequence[SSAValue | None]], None]
    ] = field(
        default_factory=list[Callable[[Operation, Sequence[SSAValue | None]], None]],
        kw_only=True,
    )
    """Callbacks that are called when an operation is replaced."""

    def handle_operation_removal(self, op: Operation) -> None:
        """Pass the operation that will be removed to the registered callbacks."""
        for handler in self.operation_removal_handler:
            handler(op)

    def handle_operation_modification(self, op: Operation) -> None:
        """Pass the operation that was just modified to the registered callbacks."""
        for handler in self.operation_modification_handler:
            handler(op)

    def handle_operation_replacement(
        self, op: Operation, new_results: Sequence[SSAValue | None]
    ) -> None:
        """Pass the operation that will be replaced to the registered callbacks."""
        for handler in self.operation_replacement_handler:
            handler(op, new_results)

    def extend_from_listener(self, listener: BuilderListener | PatternRewriterListener):
        """Forward all callbacks from `listener` to this listener."""
        super().extend_from_listener(listener)
        if isinstance(listener, PatternRewriterListener):
            self.operation_removal_handler.extend(listener.operation_removal_handler)
            self.operation_modification_handler.extend(
                listener.operation_modification_handler
            )
            self.operation_replacement_handler.extend(
                listener.operation_replacement_handler
            )

@dataclass(eq=False, init=False)
class PatternRewriter(Builder, PatternRewriterListener):
    """
    A rewriter used during pattern matching.
    Once an operation is matched, this rewriter is used to apply
    modification to the operation and its children.
    """

    _current_operation: Operation | None
    """The matched operation."""

    has_done_action: bool = field(default=False, init=False)
    """Has the rewriter done any action during the current match."""

    def __init__(self, insert_point: Operation | InsertPoint):
        PatternRewriterListener.__init__(self)
        if isinstance(insert_point, Operation):
            self._current_operation = insert_point
            insert_point = InsertPoint.before(insert_point)
        else:
            self._current_operation = None
        Builder.__init__(self, insert_point)

    @property
    @deprecated("PatternRewriter.current_operation is deprecated")
    def current_operation(self) -> Operation:
        """The matched operation."""
        assert self._current_operation is not None
        return self._current_operation

    @current_operation.setter
    def current_operation(self, current_operation: Operation) -> None:
        self._current_operation = current_operation

    def insert(
        self,
        op: InsertOpInvT,
        insertion_point: InsertPoint | None = None,
    ) -> InsertOpInvT:
        """Insert operations at a certain location in a block."""
        self.has_done_action = True
        return super().insert(op, insertion_point)

    def erase(self, op: Operation, safe_erase: bool = True):
        """
        Erase an operation.
        If safe_erase is True, check that the operation has no uses.
        Otherwise, replace its uses with ErasedSSAValue.
        """
        self.has_done_action = True
        self.handle_operation_removal(op)
        Rewriter.erase_op(op, safe_erase=safe_erase)

    @deprecated("Use .erase(op, safe_erase) instead")
    def erase_op(self, op: Operation, safe_erase: bool = True):
        """Erase an operation."""
        self.erase(op, safe_erase=safe_erase)

    def replace_all_uses_with(
        self, from_value: SSAValue, to_value: SSAValue | None, safe_erase: bool = True
    ):
        """Replace all uses of `old` with `new`."""
        if from_value is to_value:
            return

        modified_ops = [use.operation for use in from_value.uses]
        if to_value is None:
            self.has_done_action = True
            from_value.erase(safe_erase=safe_erase)
        else:
            from_value.replace_all_uses_with(to_value)

        if modified_ops:
            self.has_done_action = True
        for op in modified_ops:
            self.handle_operation_modification(op)

    def replace_uses_with_if(
        self,
        from_value: SSAValue,
        to_value: SSAValue,
        predicate: Callable[[Use], bool],
    ):
        """Replace uses of `old` satisfying `predicate` with `new`."""
        if from_value is to_value:
            return

        tracking = _TrackingPredicate(predicate)
        from_value.replace_uses_with_if(to_value, tracking)

        if tracking.modified_ops:
            self.has_done_action = True
        for op in tracking.modified_ops:
            self.handle_operation_modification(op)

    @deprecated("Use .replace(op, new_op) instead")
    def replace_matched_op(
        self,
        new_ops: Operation | Sequence[Operation],
        new_results: Sequence[SSAValue | None] | None = None,
        safe_erase: bool = True,
    ):
        """
        Replace the matched operation with new operations.
        Also, optionally specify SSA values to replace the operation results.
        If safe_erase is True, check that the operation has no uses.
        Otherwise, replace its uses with ErasedSSAValue.
        """
        assert self._current_operation is not None
        self.replace(
            self._current_operation, new_ops, new_results, safe_erase=safe_erase
        )

    def replace(
        self,
        op: Operation,
        new_ops: Operation | Sequence[Operation],
        new_results: Sequence[SSAValue | None] | None = None,
        safe_erase: bool = True,
    ):
        """
        Replace an operation with new operations.
        Also, optionally specify SSA values to replace the operation results.
        If safe_erase is True, check that the operation has no uses.
        Otherwise, replace its uses with ErasedSSAValue.
        """
        self.has_done_action = True

        if isinstance(new_ops, Operation):
            new_ops = (new_ops,)

        # First, insert the new operations before the matched operation
        self.insert(new_ops, InsertPoint.before(op))

        if new_results is None:
            new_results = new_ops[-1].results if new_ops else []

        if len(op.results) != len(new_results):
            raise ValueError(
                f"Expected {len(op.results)} new results, but got {len(new_results)}"
            )

        # Then, replace the results with new ones
        self.handle_operation_replacement(op, new_results)
        for old_result, new_result in zip(op.results, new_results):
            self.replace_all_uses_with(old_result, new_result, safe_erase=safe_erase)

            # Preserve name hints for ops with multiple results
            if new_result is not None and not new_result.name_hint:
                new_result.name_hint = old_result.name_hint

        # Add name hints for existing ops, only if there is a single new result
        if (
            len(new_results) == 1
            and (only_result := new_results[0]) is not None
            and (name_hint := only_result.name_hint) is not None
        ):
            for new_op in new_ops:
                for res in new_op.results:
                    if not res.name_hint:
                        res.name_hint = name_hint

        # Then, erase the original operation
        self.erase(op, safe_erase=safe_erase)

    @deprecated("Use .replace(op, new_ops, new_results, safe_erase) instead")
    def replace_op(
        self,
        op: Operation,
        new_ops: Operation | Sequence[Operation],
        new_results: Sequence[SSAValue | None] | None = None,
        safe_erase: bool = True,
    ):
        """Replace an operation with new operations."""
        self.replace(op, new_ops, new_results, safe_erase=safe_erase)

    def replace_value_with_new_type(
        self, val: SSAValue, new_type: Attribute
    ) -> SSAValue:
        """
        Replace a value with a value of a new type, and return the new value.
        This will insert the new value in the operation or block, and remove the existing
        value.
        """
        self.has_done_action = True
        if isinstance(val, OpResult):
            self.handle_operation_modification(val.op)
        if isinstance(val, BlockArgument):
            if (op := val.block.parent_op()) is not None:
                self.handle_operation_modification(op)
        return Rewriter.replace_value_with_new_type(val, new_type)

    def insert_block_argument(
        self, block: Block, index: int, arg_type: Attribute
    ) -> BlockArgument:
        """Insert a new block argument."""
        self.has_done_action = True
        return block.insert_arg(arg_type, index)

    def erase_block_argument(self, arg: BlockArgument, safe_erase: bool = True) -> None:
        """
        Erase a new block argument.
        If safe_erase is true, then raise an exception if the block argument has still
        uses, otherwise, replace it with an ErasedSSAValue.
        """
        self.has_done_action = True
        self.replace_all_uses_with(arg, None, safe_erase=safe_erase)
        arg.block.erase_arg(arg, safe_erase)

    def inline_block(
        self,
        block: Block,
        insertion_point: InsertPoint,
        arg_values: Sequence[SSAValue] = (),
    ):
        """
        Move the block operations to the specified insertion point.
        """
        self.has_done_action = True
        Rewriter.inline_block(block, insertion_point, arg_values=arg_values)

    def move_region_contents_to_new_regions(self, region: Region) -> Region:
        """Move the region blocks to a new region."""
        self.has_done_action = True
        return Rewriter.move_region_contents_to_new_regions(region)

    def inline_region(self, region: Region, insertion_point: BlockInsertPoint) -> None:
        """Move the region blocks to the specified insertion point."""
        self.has_done_action = True
        Rewriter.inline_region(region, insertion_point)

    def notify_op_modified(self, op: Operation) -> None:
        """
        Notify the rewriter that an operation was modified in the pattern.
        This will correctly update the rewriter state.
        """
        self.has_done_action = True
        self.handle_operation_modification(op)

class RewritePattern(ABC):
    """
    A side-effect free rewrite pattern matching on a DAG.
    """

    # The / in the function signature makes the previous arguments positional, see
    # https://peps.python.org/pep-0570/
    # This is used by the op_type_rewrite_pattern
    @abstractmethod
    def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter, /):
        """
        Match an operation, and optionally perform a rewrite using the rewriter.
        """
        ...

@dataclass
class TypeConversionPattern(RewritePattern):
    """
    Base pattern for type conversion. It is supposed to be inherited from, then one can
    implement `convert_type` to define the conversion.

    It will convert an Operations' result types, dictionary attributes, and block arguments.

    One can use `@attr_type_rewrite_pattern` on this defined method to automatically filter
    on the Attribute type used.

    This base pattern defines two flags:

    - `recursive` (defaulting to False): recurse over structured attributes to convert
      parameters.
      e.g. a recusrive `i32` to `index` conversion will convert `vector<i32>` to
      `vector<index>`.
    - `ops` (defaulting to any Operation) is a tuple of Operation types on which to apply
      the defined attribute conversion.
    """

    recursive: bool = False
    """
    recurse over structured attributes to convert parameters.
    Defaults to False.
    """
    ops: tuple[type[Operation], ...] | None = None
    """
    A tuple of Operation types on which to apply the defined attribute conversion.
    Defaults to any operation type.
    """

    @abstractmethod
    def convert_type(self, typ: Attribute, /) -> Attribute | None:
        """
        The method to implement to define a TypeConversionPattern

        This defines how the input Attribute should be converted.
        It allows returning None, meaning "this attribute should not
        be converted".
        """
        raise NotImplementedError()

    @final
    def _convert_type_rec(self, typ: Attribute) -> Attribute | None:
        """
        Provided recursion over structed/parameterized Attributes.
        """
        inp = typ
        if self.recursive:
            if isinstance(typ, ParametrizedAttribute):
                parameters = list(
                    self._convert_type_rec(p) or p for p in typ.parameters
                )
                inp = type(typ).new(parameters)
            if isa(typ, ArrayAttr[Attribute]):
                parameters = tuple(self._convert_type_rec(p) or p for p in typ)
                inp = type(typ).new(parameters)
            if isa(typ, DictionaryAttr):
                parameters = {k: self._convert_type_rec(v) for k, v in typ.data.items()}
                inp = type(typ).new(parameters)
        converted = self.convert_type(inp)
        return converted if converted is not None else inp

    @final
    def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter):
        """
        Pattern application implementation
        """
        if self.ops and not isinstance(op, self.ops):
            return
        new_result_types: list[Attribute] = []
        new_attributes: dict[str, Attribute] = {}
        new_properties: dict[str, Attribute] = {}
        changed: bool = False
        for result in op.results:
            converted = self._convert_type_rec(result.type)
            new_result_types.append(converted or result.type)
            if converted is not None and converted != result.type:
                changed = True
        for name, attribute in op.attributes.items():
            converted = self._convert_type_rec(attribute)
            new_attributes[name] = converted or attribute
            if converted is not None and converted != attribute:
                changed = True
        for name, attribute in op.properties.items():
            converted = self._convert_type_rec(attribute)
            new_properties[name] = converted or attribute
            if converted is not None and converted != attribute:
                changed = True
        for region in op.regions:
            for block in region.blocks:
                for arg in block.args:
                    converted = self._convert_type_rec(arg.type)
                    if converted is not None and converted != arg.type:
                        rewriter.replace_value_with_new_type(arg, converted)
        if changed:
            regions = [op.detach_region(r) for r in op.regions]
            new_op = type(op).create(
                operands=op.operands,
                result_types=new_result_types,
                properties=new_properties,
                attributes=new_attributes,
                successors=op.successors,
                regions=regions,
            )
            rewriter.replace(op, new_op)
            for new, old in zip(new_op.results, op.results):
                new.name_hint = old.name_hint

@dataclass(eq=False, repr=False)
class GreedyRewritePatternApplier(RewritePattern):
    """
    Apply a list of patterns in order until one pattern matches,
    and then use this rewrite.
    By default, the applier attempts to fold the operation first.
    If the operation is trivially dead, it is erased.
    """

    rewrite_patterns: list[RewritePattern]
    """The list of rewrites to apply in order."""

    ctx: Context | None = field(default=None)
    """Used to materialize constant operations when folding."""

    folding_enabled: bool = field(default=False, kw_only=True)
    """
    Whether the folders should be invoked.
    If this is True, the GreedyRewritePatternApplier must also have a context.
    """

    dce_enabled: bool = field(default=True, kw_only=True)
    """
    Whether trivial dead code elimination should be run on all operations before
    attempting to rewrite them.
    """

    def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter) -> None:
        from xdsl.transforms.dead_code_elimination import is_trivially_dead

        if self.dce_enabled and is_trivially_dead(op):
            rewriter.erase(op)
            return

        # Do not fold constant ops. That would lead to an infinite folding loop,
        # as every constant op would be folded to an Attribute and then
        # immediately be rematerialized as a constant op, which is then put
        # back into the worklist.
        if (
            self.folding_enabled
            and op.has_trait(HasFolder, value_if_unregistered=False)
            and not op.has_trait(ConstantLike, value_if_unregistered=True)
        ):
            if self.ctx is None:
                raise ValueError("Context is required for folding")
            folded = Folder(self.ctx).try_fold(op)
            if folded is not None:
                folded_values, folded_ops = folded
                rewriter.replace(op, new_ops=folded_ops, new_results=folded_values)
                return

        for pattern in self.rewrite_patterns:
            pattern.match_and_rewrite(op, rewriter)
            if rewriter.has_done_action:
                return
        return