Function constant pinning

class FunctionConstantPinning(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, func_op: func.FuncOp, rewriter: PatternRewriter, /):
        # can't rewrite nested functions yet
        if not isinstance(func_op.parent_op(), builtin.ModuleOp):
            return

        # check if the function contains a "pin_to_constants" annotated operation
        split_op = func_contains_pinning_annotation(func_op)
        if split_op is None:
            return
        # can't do splits on multi-value ops
        if len(split_op.results) != 1:
            return

        # get list of vals to pin:
        pinned_vals = get_pinned_vals_for_op(split_op)
        # drop malformed or empty pinning:
        if pinned_vals is None or len(pinned_vals) == 0:
            split_op.attributes.pop(PIN_CONSTANT_VALS)
            return

        return_types = tuple(func_op.function_type.outputs)

        function_remainder = split_op.next_op
        if function_remainder is None:
            # "Can't ping values for terminator"
            return

        # grab the first value to pin:
        val = pinned_vals.pop()
        # generate the function containing pinned value:
        new_func = generate_func_with_pinned_val(func_op, val, rewriter)
        # insert the specialized function after the generic function (the one we matched on)
        rewriter.insert(new_func, InsertPoint.after(func_op))
        # insert a compare to the value we specialize and, and branch on if we are equal
        rewriter.insert(
            [
                cst := arith.ConstantOp(val, split_op.results[0].type),
                is_eq := arith.CmpiOp(split_op.results[0], cst, "eq"),
                scf_if := scf.IfOp(
                    is_eq,
                    return_types,
                    [
                        # if we are equal to the specialized value, call the function:
                        call_op := func.CallOp(
                            new_func.sym_name,
                            func_op.body.block.args,
                            return_types,
                        ),
                        # yield call results
                        scf.YieldOp(*call_op.results),
                    ],
                    # empty region placeholder, will be filled in later
                    # grab a reference to it
                    Region(dest_block := Block()),
                ),
            ],
            InsertPoint.after(split_op),
        )

        # iterate over the remainder of the function:
        # grab a reference to the next operation in the remainder.
        # this is because we will modify the op and therefore loose the "old" next op.
        next_op = function_remainder.next_op

        # unless we already hit the block terminator
        # while we haven't reached the return statement:
        while (
            next_op is not None and function_remainder is not func_op.body.block.last_op
        ):
            # detatch the function
            function_remainder.detach()
            # re-insert it inside the else block of the if statement
            rewriter.insert(function_remainder, InsertPoint.at_end(dest_block))
            # go to next op
            function_remainder = next_op
            next_op = function_remainder.next_op

        # insert a yield that yields the return values
        rewriter.insert(
            scf.YieldOp(*function_remainder.operands), InsertPoint.at_end(dest_block)
        )
        # return the results of the scf.if
        rewriter.replace(function_remainder, func.ReturnOp(*scf_if.results))

        # remove pinning attribute
        if pinned_vals:
            split_op.attributes[PIN_CONSTANT_VALS] = ArrayAttr(pinned_vals)
        else:
            split_op.attributes.pop(PIN_CONSTANT_VALS)

class FunctionConstantPinningPass(ModulePass):
    """
    This pass consumes IR annotated with special hints to generate new functions that have certain SSA values pinned
    to a constant, usually to enable further optimization options on this pinned function.

    The original function is changed to dynamically dispatch to this pinned function when the ssa value matches the
    given constant.

    Any single-result operation annotated with a "pin_to_constants" attribute containing an array of values, that is
    located within a function body triggers this optimization. These annotations are usually inserted by previous
    passes that know that they would want to generate a more optimized version of their function for specific values
    of a run-time determined variable.

    An example might be a function that branches repeatedly on a specific variable:

    ```
    function test() {
        x = calc_condition()
        if (x) {
          specific_thing()
        }

        some_thing() // A

        if (x) {
          another_thing()
        }

        some_thing() // B
    }
    ```

    if we can pin `x` to `true`, we are suddenly able to generate two much simple function bodies (after constant folding)

    ```
    function test() {
        x = calc_condition()
        if (x) {
            test_pinned()
            return
        }

        some_thing() // A
        some_thing() // B
    }

    function test_pinned() {
        specific_thing()
        some_thing() // A
        another_thing()
        some_thing() // B
    }
    ```

    Note that the function `test_pinned` might be much easier to optimize for a compiler if there are state
    dependencies between `specific_thing`, `another_thing` and `some_thing`.
    """

    name = "function-constant-pinning"

    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
        PatternRewriteWalker(FunctionConstantPinning()).rewrite_module(op)