Cf

class AssertTrue(RewritePattern):
    """Erase assertion if argument is constant true."""

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.AssertOp, rewriter: PatternRewriter):
        owner = op.arg.owner

        if not isinstance(owner, arith.ConstantOp):
            return

        value = owner.value

        if not isinstance(value, IntegerAttr):
            return

        if not value.value.data:
            return

        rewriter.replace(op, [])

class SimplifyBrToBlockWithSinglePred(RewritePattern):
    """
    Simplify a branch to a block that has a single predecessor. This effectively
    merges the two blocks.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.BranchOp, rewriter: PatternRewriter):
        succ = op.successor
        parent = op.parent_block()
        if parent is None:
            return

        # Check that the successor block has a single predecessor
        if succ == parent or len(succ.predecessors()) != 1:
            return

        br_operands = op.operands
        rewriter.erase(op)
        rewriter.inline_block(succ, InsertPoint.at_end(parent), br_operands)

class SimplifyPassThroughBr(RewritePattern):
    """
      br ^bb1
    ^bb1
      br ^bbN(...)

     -> br ^bbN(...)
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.BranchOp, rewriter: PatternRewriter):
        # Check the successor doesn't point back to the current block
        parent = op.parent_block()
        if parent is None or op.successor == parent:
            return

        ret = collapse_branch(op.successor, op.arguments)
        if ret is None:
            return
        (block, args) = ret

        rewriter.replace(op, cf.BranchOp(block, *args))

class SimplifyConstCondBranchPred(RewritePattern):
    """
    cf.cond_br true, ^bb1, ^bb2
     -> br ^bb1
    cf.cond_br false, ^bb1, ^bb2
     -> br ^bb2
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.ConditionalBranchOp, rewriter: PatternRewriter):
        # Check if cond operand is constant
        cond = const_evaluate_operand(op.cond)

        if cond is None:
            return

        if cond:
            rewriter.replace(op, cf.BranchOp(op.then_block, *op.then_arguments))
        else:
            rewriter.replace(op, cf.BranchOp(op.else_block, *op.else_arguments))

class SimplifyPassThroughCondBranch(RewritePattern):
    """
      cf.cond_br %cond, ^bb1, ^bb2
    ^bb1
      br ^bbN(...)
    ^bb2
      br ^bbK(...)

     -> cf.cond_br %cond, ^bbN(...), ^bbK(...)
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.ConditionalBranchOp, rewriter: PatternRewriter):
        # Try to collapse both branches
        collapsed_then = collapse_branch(op.then_block, op.then_arguments)
        collapsed_else = collapse_branch(op.else_block, op.else_arguments)

        # If neither collapsed then we return
        if collapsed_then is None and collapsed_else is None:
            return

        (new_then, new_then_args) = collapsed_then or (op.then_block, op.then_arguments)

        (new_else, new_else_args) = collapsed_else or (op.else_block, op.else_arguments)

        rewriter.replace(
            op,
            cf.ConditionalBranchOp(
                op.cond, new_then, new_then_args, new_else, new_else_args
            ),
        )

class SimplifyCondBranchIdenticalSuccessors(RewritePattern):
    """
    cf.cond_br %cond, ^bb1(A, ..., N), ^bb1(A, ..., N)
     -> br ^bb1(A, ..., N)

    cf.cond_br %cond, ^bb1(A), ^bb1(B)
     -> %select = arith.select %cond, A, B
        br ^bb1(%select)
    """

    @staticmethod
    def _merge_operand(
        op1: SSAValue,
        op2: SSAValue,
        rewriter: PatternRewriter,
        cond_br: cf.ConditionalBranchOp,
    ) -> SSAValue:
        if op1 == op2:
            return op1
        select = arith.SelectOp(cond_br.cond, op1, op2)
        rewriter.insert(select, InsertPoint.before(cond_br))
        return select.result

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.ConditionalBranchOp, rewriter: PatternRewriter):
        # Check that the true and false destinations are the same
        if op.then_block != op.else_block:
            return

        merged_operands = tuple(
            self._merge_operand(op1, op2, rewriter, op)
            for (op1, op2) in zip(op.then_arguments, op.else_arguments, strict=True)
        )

        rewriter.replace(op, cf.BranchOp(op.then_block, *merged_operands))

class CondBranchTruthPropagation(RewritePattern):
    """
      cf.cond_br %arg0, ^trueB, ^falseB

    ^trueB:
      "test.consumer1"(%arg0) : (i1) -> ()
       ...

    ^falseB:
      "test.consumer2"(%arg0) : (i1) -> ()
      ...

    ->

      cf.cond_br %arg0, ^trueB, ^falseB
    ^trueB:
      "test.consumer1"(%true) : (i1) -> ()
      ...

    ^falseB:
      "test.consumer2"(%false) : (i1) -> ()
      ...
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.ConditionalBranchOp, rewriter: PatternRewriter):
        if len(op.then_block.predecessors()) == 1:
            if any(
                use.operation.parent_block() is op.then_block for use in op.cond.uses
            ):
                const_true = arith.ConstantOp(BoolAttr.from_bool(True))
                rewriter.insert(const_true, InsertPoint.before(op))
                rewriter.replace_uses_with_if(
                    op.cond,
                    const_true.result,
                    lambda use: use.operation.parent_block() is op.then_block,
                )
        if len(op.else_block.predecessors()) == 1:
            if any(
                use.operation.parent_block() is op.else_block for use in op.cond.uses
            ):
                const_false = arith.ConstantOp(BoolAttr.from_bool(False))
                rewriter.insert(const_false, InsertPoint.before(op))
                rewriter.replace_uses_with_if(
                    op.cond,
                    const_false.result,
                    lambda use: use.operation.parent_block() is op.else_block,
                )

class SimplifySwitchWithOnlyDefault(RewritePattern):
    """
    switch %flag : i32, [
      default:  ^bb1
    ]
     -> br ^bb1
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.SwitchOp, rewriter: PatternRewriter):
        if not op.case_blocks:
            rewriter.replace(op, cf.BranchOp(op.default_block, *op.default_operands))

class DropSwitchCasesThatMatchDefault(RewritePattern):
    """
    switch %flag : i32, [
      default: ^bb1,
      42: ^bb1,
      43: ^bb2
    ]
    ->
    switch %flag : i32, [
      default: ^bb1,
      43: ^bb2
    ]
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.SwitchOp, rewriter: PatternRewriter):
        def predicate(
            switch_case: IntegerAttr,
            block: Block,
            operands: Sequence[Operation | SSAValue],
        ) -> bool:
            return block == op.default_block and operands == op.default_operands

        drop_case_helper(rewriter, op, predicate)

class SimplifyConstSwitchValue(RewritePattern):
    """
    switch %c_42 : i32, [
      default: ^bb1,
      42: ^bb2,
      43: ^bb3
    ]
    -> br ^bb2
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.SwitchOp, rewriter: PatternRewriter):
        if (flag := const_evaluate_operand(op.flag)) is not None:
            fold_switch(op, rewriter, flag)

class SimplifyPassThroughSwitch(RewritePattern):
    """
    switch %c_42 : i32, [
      default: ^bb1,
      42: ^bb2,
    ]
    ^bb2:
      br ^bb3
    ->
    switch %c_42 : i32, [
      default: ^bb1,
      42: ^bb3,
    ]
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.SwitchOp, rewriter: PatternRewriter):
        requires_change = False

        new_case_blocks: list[Block] = []
        new_case_operands: list[Sequence[Operation | SSAValue]] = []

        for block, operands in zip(op.case_blocks, op.case_operand, strict=True):
            collapsed = collapse_branch(block, operands)
            requires_change |= collapsed is not None
            (new_block, new_operands) = collapsed or (block, operands)
            new_case_blocks.append(new_block)
            new_case_operands.append(new_operands)

        collapsed = collapse_branch(op.default_block, op.default_operands)

        requires_change |= collapsed is not None

        (default_block, default_operands) = collapsed or (
            op.default_block,
            op.default_operands,
        )

        if requires_change:
            rewriter.replace(
                op,
                cf.SwitchOp(
                    op.flag,
                    default_block,
                    default_operands,
                    op.case_values,
                    new_case_blocks,
                    new_case_operands,
                ),
            )

class SimplifySwitchFromSwitchOnSameCondition(RewritePattern):
    """
    switch %flag : i32, [
      default: ^bb1,
      42: ^bb2,
    ]
    ^bb2:
      switch %flag : i32, [
        default: ^bb3,
        42: ^bb4
      ]
    ->
    switch %flag : i32, [
      default: ^bb1,
      42: ^bb2,
    ]
    ^bb2:
      br ^bb4

     and

    switch %flag : i32, [
      default: ^bb1,
      42: ^bb2,
    ]
    ^bb2:
      switch %flag : i32, [
        default: ^bb3,
        43: ^bb4
      ]
    ->
    switch %flag : i32, [
      default: ^bb1,
      42: ^bb2,
    ]
    ^bb2:
      br ^bb3

    and

    switch %flag : i32, [
      default: ^bb1,
      42: ^bb2
    ]
    ^bb1:
      switch %flag : i32, [
        default: ^bb3,
        42: ^bb4,
        43: ^bb5
      ]
    ->
    switch %flag : i32, [
      default: ^bb1,
      42: ^bb2,
    ]
    ^bb1:
      switch %flag : i32, [
        default: ^bb3,
        43: ^bb5
      ]
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: cf.SwitchOp, rewriter: PatternRewriter):
        block = op.parent_block()
        if block is None:
            return
        if (pred := block.get_unique_use()) is None:
            return
        switch = pred.operation
        if not isinstance(switch, cf.SwitchOp):
            return

        if switch.flag != op.flag:
            return

        case_values = switch.case_values
        if case_values is None:
            return

        if pred.index != 0:
            fold_switch(
                op,
                rewriter,
                case_values.get_values()[pred.index - 1],
            )
        else:

            def predicate(
                switch_case: IntegerAttr,
                block: Block,
                operands: Sequence[Operation | SSAValue],
            ) -> bool:
                return switch_case in case_values.get_attrs()

            drop_case_helper(rewriter, op, predicate)