Scf

class RehoistConstInLoops(RewritePattern):
    """
    Carry out const definitions from the loops.
    In the future this will probably be done by the pattern rewriter itself, like it's
    done in the MLIR's applyPatternsAndFoldGreedily.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: scf.ForOp, rewriter: PatternRewriter) -> None:
        for child_op in op.body.ops:
            if child_op.has_trait(ConstantLike):
                # we only rehoist consts that are not embeded in another region inside the loop
                rewriter.insert(new_const := child_op.clone())
                rewriter.replace(child_op, (), new_const.results)

class SimplifyTrivialLoops(RewritePattern):
    """
    Rewriting pattern that erases loops that are known not to iterate, replaces
    single-iteration loops with their bodies, and removes empty loops that iterate at
    least once and only return values defined outside of the loop.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: scf.ForOp, rewriter: PatternRewriter) -> None:
        # If the upper bound is the same as the lower bound, the loop does not iterate,
        # just remove it.
        if (lb := const_evaluate_operand(op.lb)) is None:
            return
        if (ub := const_evaluate_operand(op.ub)) is None:
            return

        if lb == ub:
            rewriter.replace(op, (), op.iter_args)
            return

        # If the loop is known to have 0 iterations, remove it.
        if (diff := ub - lb) <= 0:
            rewriter.replace(op, (), op.iter_args)
            return

        if (step := const_evaluate_operand(op.step)) is None:
            return

        # If the loop is known to have 1 iteration, inline its body and remove the loop.
        # TODO: handle signless values
        if step >= diff:
            block_args = (op.lb, *op.iter_args)
            replace_op_with_region(
                rewriter,
                op,
                op.body,
                block_args,
            )

class IfPropagateConstantCondition(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: scf.IfOp, rewriter: PatternRewriter) -> None:
        if (cond := const_evaluate_operand(op.cond)) is None:
            return
        if not cond and not op.false_region.blocks:
            # Cannot use helper below as false region is not single-block
            rewriter.erase(op)
            return
        region = op.true_region if cond else op.false_region
        replace_op_with_region(rewriter, op, region)

class SingleBlockExecuteInliner(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(
        self, op: scf.ExecuteRegionOp, rewriter: PatternRewriter
    ) -> None:
        assert op.region.first_block is not None
        if op.region.first_block is not op.region.last_block:
            return
        replace_op_with_region(rewriter, op, op.region)