cubecl-core 0.11.0-pre.3

use alloc::{boxed::Box, vec, vec::Vec};
use cubecl_environment::collections::HashMap;
use cubecl_ir::{
    VectorSize,
    attributes::{IndexAttr, ZeroAttr},
    dialect::{
        base::OperationPtrExt,
        cmp::IEqualOp,
        general::{CopyOp, ReinterpretCastOp, SelectOp},
        math::{IAddOp, IMulOp, UDivOp, URemOp},
        matrix,
        memory::{DeclareVariableOp, IndexOp},
        vector::{
            CompositeConstructOp, CompositeExtractOp, CompositeInsertOp, VectorExtractDynamicOp,
            VectorInsertDynamicOp,
        },
    },
    interfaces::{MaybeVectorizedType, TriviallyUnrollable, TypedExt},
    prelude::*,
    try_cast_op,
    types::{ArrayType, AtomicType, PointerType, RuntimeArrayType, VectorType},
    verify_op_succ, verify_ty_succ,
};
use pliron::{
    builtin::{
        ops::{ConstantOp, FuncOp},
        types::FunctionType,
    },
    graph::walkers::{WALKCONFIG_PREORDER_FORWARD, uninterruptible::mutable::walk_op},
};

type Mappings = HashMap<Value, Vec<Value>>;

#[derive(Debug, new)]
pub struct UnrollPass {
    max_vector_size: VectorSize,
}

#[type_interface]
pub trait UnrollableType: MaybeVectorizedType {
    verify_ty_succ!();
    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle;
}

#[type_interface_impl]
impl UnrollableType for VectorType {
    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle {
        VectorType::get(ctx, self.inner, vectorization).into()
    }
}

#[type_interface_impl]
impl UnrollableType for AtomicType {
    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle {
        let inner = self.inner.deref(ctx);
        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
        let new_inner = unrollable.with_vector_size(ctx, vectorization);
        AtomicType::get(ctx, new_inner).into()
    }
}

#[type_interface_impl]
impl UnrollableType for PointerType {
    fn with_vector_size(&self, ctx: &Context, vectorization: usize) -> TypeHandle {
        let inner = self.inner.deref(ctx);
        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
        let new_inner = unrollable.with_vector_size(ctx, vectorization);
        PointerType::get(ctx, new_inner, self.address_space).into()
    }
}

#[type_interface_impl]
impl UnrollableType for ArrayType {
    fn with_vector_size(&self, ctx: &Context, new_vec: usize) -> TypeHandle {
        let current_vec = self.vector_size(ctx);
        let inner = self.inner.deref(ctx);
        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
        let new_inner = unrollable.with_vector_size(ctx, new_vec);
        ArrayType::get(ctx, new_inner, self.length * current_vec / new_vec).into()
    }
}

#[type_interface_impl]
impl UnrollableType for RuntimeArrayType {
    fn with_vector_size(&self, ctx: &Context, new_vec: usize) -> TypeHandle {
        let inner = self.inner.deref(ctx);
        let unrollable = type_cast::<dyn UnrollableType>(&*inner).expect("Should be implemented");
        let new_inner = unrollable.with_vector_size(ctx, new_vec);
        RuntimeArrayType::get(ctx, new_inner).into()
    }
}

#[op_interface]
pub trait CustomUnrollOp {
    verify_op_succ!();
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState);
}

#[op_interface_impl]
impl CustomUnrollOp for DeclareVariableOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let value_ty = self.value_ty(ctx).get_type(ctx);
        let current_vec = value_ty.try_get_vector_size(ctx).unwrap_or(1);
        if current_vec <= state.max_vector_size {
            return;
        }

        state.result.ir_changed |= IRStatus::Changed;
        let result = self.get_result(ctx);
        let addr_space = *self.addr_space(ctx);
        // Align isn't handled properly, but targets that unroll ignore this anyways
        let align = *self.alignment(ctx);
        let new_value_ty = unroll_ty(ctx, value_ty, state.max_vector_size);
        let new_ptr_ty = PointerType::get(ctx, new_value_ty, addr_space.0);

        // The initializer has a type too, and it must match the declaration.
        // Unroll it as well, or we end up with a broken
        // `array<vec4<f32>, 8> = array<vec8<f32>, 4>()`.
        let unrolled_init = |ctx: &Context| {
            self.initializer(ctx)
                .map(|init| match init.is::<ZeroAttr>() {
                    true => ZeroAttr::new(new_value_ty).into(),
                    false => init.clone(),
                })
        };

        // Array doesn't change size, so no need to duplicate the declaration
        if new_value_ty.size(ctx) == value_ty.size(ctx) {
            let init = unrolled_init(ctx);
            self.set_value_ty(ctx, new_value_ty);
            self.get_result(ctx).set_type(ctx, new_ptr_ty.into());
            if let Some(init) = init {
                self.set_initializer(ctx, init);
            }
        } else {
            let factor = current_vec / state.max_vector_size;
            let mut results = vec![];
            for _ in 0..factor {
                let init = unrolled_init(ctx);
                let new_op = DeclareVariableOp::new(ctx, new_value_ty, addr_space, align, init);
                new_op
                    .get_operation()
                    .insert_before(ctx, self.get_operation());
                results.push(new_op.get_result(ctx));
            }
            state.mappings.insert(result, results);
            state.to_erase.push(self.get_operation());
        }
    }
}

#[op_interface_impl]
impl CustomUnrollOp for IndexOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let base = self.base(ctx);
        let checked = self.checked(ctx);
        let current_vec = try_get_vec(ctx, self.get_result(ctx));
        if current_vec > state.max_vector_size {
            state.result.ir_changed |= IRStatus::Changed;
            let unroll_factor = current_vec / state.max_vector_size;
            let unroll_const = const_usize(ctx, self, unroll_factor);

            let mul = IMulOp::new(ctx, self.index(ctx), unroll_const);
            mul.get_operation().insert_before(ctx, self.get_operation());
            let start_idx = mul.get_result(ctx);

            let new_results = (0..unroll_factor)
                .map(|i| {
                    let i = const_usize(ctx, self, i);
                    let add = IAddOp::new(ctx, start_idx, i);
                    add.get_operation().insert_before(ctx, self.get_operation());
                    let idx = add.get_result(ctx);

                    let op = IndexOp::maybe_checked(ctx, base, idx, checked);
                    op.get_operation().insert_before(ctx, self.get_operation());
                    op.get_result(ctx)
                })
                .collect();

            state.mappings.insert(self.get_result(ctx), new_results);
            state.to_erase.push(self.get_operation());
        }
    }
}

fn const_usize(ctx: &mut Context, anchor: &dyn Op, value: usize) -> Value {
    let op = ConstantOp::new(ctx, Box::new(IndexAttr::new(value)));
    op.get_operation()
        .insert_before(ctx, anchor.get_operation());
    op.get_result(ctx)
}

#[op_interface_impl]
impl CustomUnrollOp for CompositeExtractOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let vector = self.composite(ctx);
        if !vector.is_vector(ctx) {
            return;
        }

        let current_vec = vector.vector_size(ctx);
        if current_vec > state.max_vector_size {
            state.result.ir_changed |= IRStatus::Changed;
            let index = self.index(ctx).0;

            let unroll_idx = index / state.max_vector_size;
            let sub_idx = index % state.max_vector_size;

            let new_vector = state.mappings.get(&vector).expect("Should exist")[unroll_idx];
            vector.replace_use_with(ctx, self.composite_as_use(ctx), &new_vector);
            self.set_index(ctx, sub_idx);
        }
    }
}

#[op_interface_impl]
impl CustomUnrollOp for CompositeInsertOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let vector = self.composite(ctx);
        if !vector.is_vector(ctx) {
            return;
        }

        let value = self.value(ctx);
        let current_vec = vector.vector_size(ctx);
        if current_vec > state.max_vector_size {
            state.result.ir_changed |= IRStatus::Changed;
            let index = self.index(ctx).0;

            let unroll_idx = index / state.max_vector_size;
            let sub_idx = index % state.max_vector_size;

            let vectors = state.mappings.get(&vector).expect("Should exist");

            let new_results = vectors.iter().enumerate().map(|(i, vector)| {
                let op = if i == unroll_idx {
                    CompositeInsertOp::new(ctx, *vector, value, sub_idx).get_operation()
                } else {
                    CopyOp::new(ctx, *vector).get_operation()
                };
                op.insert_before(ctx, self.get_operation());
                op.deref(ctx).get_result(0)
            });
            let new_results = new_results.collect();
            state.mappings.insert(self.get_result(ctx), new_results);
            state.to_erase.push(self.get_operation());
        }
    }
}

/// A dynamic index reaches a lane the pass cannot name at compile time, so the
/// unrolled parts are all searched: the index splits into the part it lands in
/// and the lane within that part, and a select chain picks the part's answer.
/// The frontend documents both dynamic accessors as very slow already, and the
/// chain is over the (comptime) number of parts, not the lanes.
#[op_interface_impl]
impl CustomUnrollOp for VectorExtractDynamicOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let vector = self.vector(ctx);
        if try_get_vec(ctx, vector) <= state.max_vector_size {
            return;
        }

        state.result.ir_changed |= IRStatus::Changed;
        let (part, lane) = split_index(ctx, self, state.max_vector_size, self.index(ctx));
        let parts = state.mappings.get(&vector).expect("Should exist").clone();

        let mut extracted = None;
        for (i, part_vector) in parts.into_iter().enumerate() {
            let op = VectorExtractDynamicOp::new(ctx, part_vector, lane);
            op.get_operation().insert_before(ctx, self.get_operation());
            let value = op.get_result(ctx);
            extracted = Some(match extracted {
                Some(previous) => select_part(ctx, self, part, i, value, previous),
                None => value,
            });
        }

        let extracted = extracted.expect("Unrolled vector should have parts");
        self.get_result(ctx).replace_all_uses_with(ctx, &extracted);
        state.to_erase.push(self.get_operation());
    }
}

#[op_interface_impl]
impl CustomUnrollOp for VectorInsertDynamicOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let vector = self.vector(ctx);
        if try_get_vec(ctx, vector) <= state.max_vector_size {
            return;
        }

        state.result.ir_changed |= IRStatus::Changed;
        let value = self.value(ctx);
        let (part, lane) = split_index(ctx, self, state.max_vector_size, self.index(ctx));
        let parts = state.mappings.get(&vector).expect("Should exist").clone();

        let mut new_results = Vec::with_capacity(parts.len());
        for (i, part_vector) in parts.into_iter().enumerate() {
            // Selecting the *scalar* to write, rather than between the written
            // and untouched parts, keeps the condition off the vector operands:
            // a select over vectors with a scalar condition needs SPIR-V 1.4.
            let old = VectorExtractDynamicOp::new(ctx, part_vector, lane);
            old.get_operation().insert_before(ctx, self.get_operation());
            let written = select_part(ctx, self, part, i, value, old.get_result(ctx));

            let op = VectorInsertDynamicOp::new(ctx, part_vector, written, lane);
            op.get_operation().insert_before(ctx, self.get_operation());
            new_results.push(op.get_result(ctx));
        }

        state.mappings.insert(self.get_result(ctx), new_results);
        state.to_erase.push(self.get_operation());
    }
}

/// Split a dynamic vector index into the unrolled part it lands in and the lane
/// within that part.
fn split_index(
    ctx: &mut Context,
    anchor: &dyn Op,
    max_vector_size: VectorSize,
    index: Value,
) -> (Value, Value) {
    let size = const_usize(ctx, anchor, max_vector_size);
    let part = UDivOp::new(ctx, index, size);
    part.get_operation()
        .insert_before(ctx, anchor.get_operation());
    let lane = URemOp::new(ctx, index, size);
    lane.get_operation()
        .insert_before(ctx, anchor.get_operation());
    (part.get_result(ctx), lane.get_result(ctx))
}

/// `if part == i { on_match } else { otherwise }`.
fn select_part(
    ctx: &mut Context,
    anchor: &dyn Op,
    part: Value,
    i: usize,
    on_match: Value,
    otherwise: Value,
) -> Value {
    let i = const_usize(ctx, anchor, i);
    let is_part = IEqualOp::new(ctx, part, i);
    is_part
        .get_operation()
        .insert_before(ctx, anchor.get_operation());
    let select = SelectOp::new(ctx, is_part.get_result(ctx), on_match, otherwise);
    select
        .get_operation()
        .insert_before(ctx, anchor.get_operation());
    select.get_result(ctx)
}

/// A reinterpret changes the lane count with the lane width, so one side's pieces are not the
/// other's: a `Vector<u32, 2>` is a `Vector<e4m3, 8>`. When only one side exceeds the maximum, the
/// wide side is split at the unroll factor and the narrow side is cut into matching sub-vectors
/// with extracts and constructs. When both exceed it, the input's pieces are reinterpreted in
/// place and their lanes regrouped.
#[op_interface_impl]
impl CustomUnrollOp for ReinterpretCastOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let input = self.input(ctx);
        let result = self.get_result(ctx);
        let max = state.max_vector_size;
        let (in_vec, out_vec) = (try_get_vec(ctx, input), try_get_vec(ctx, result));
        if in_vec <= max && out_vec <= max {
            return;
        }
        // When both sides unroll, neither is a whole piece of the other: each `max`-lane input
        // piece reinterprets into the output lanes it shares bits with, and those lanes regroup
        // into the `max`-lane pieces the rest of the pass indexes into. Equal lane counts are the
        // common case here and stay one reinterpret per piece.
        if in_vec > max && out_vec > max {
            // A piece of the input covers `per_piece` lanes of the output, and both sides of the
            // reinterpret below have to be a legal vector: one lane at least, and no wider than
            // the maximum. Sub-dividing the input pieces would lift the upper bound, but nothing
            // unrolls what this pass emits, so an over-wide piece would reach the backend as-is.
            let per_piece = max * out_vec / in_vec;
            assert!(
                (1..=max).contains(&per_piece),
                "Cannot unroll a reinterpret between a {in_vec}-lane and a {out_vec}-lane vector \
                 when both exceed {max} lanes and the lane counts differ: reinterpret through a \
                 vector of at most {max} lanes"
            );
            state.result.ir_changed |= IRStatus::Changed;
            let op = self.get_operation();

            let piece_ty = lanes_type(ctx, result, per_piece);
            let pieces = state.mappings.get(&input).expect("Should exist").clone();
            let converted = pieces
                .into_iter()
                .map(|piece| {
                    let new_op = ReinterpretCastOp::new(ctx, piece_ty, piece);
                    new_op.get_operation().insert_before(ctx, op);
                    new_op.get_result(ctx)
                })
                .collect::<Vec<_>>();

            let new_results = if per_piece == max {
                converted
            } else {
                let lanes = converted
                    .into_iter()
                    .flat_map(|piece| split_lanes(ctx, piece, op))
                    .collect::<Vec<_>>();
                let result_ty = unroll_ty(ctx, result, max);
                lanes
                    .chunks(max)
                    .map(|lanes| {
                        let joined = CompositeConstructOp::new(ctx, result_ty, lanes.to_vec());
                        joined.get_operation().insert_before(ctx, op);
                        joined.get_result(ctx)
                    })
                    .collect()
            };
            state.mappings.insert(result, new_results);
            state.to_erase.push(op);
            return;
        }

        // The wide side unrolls into `factor` pieces, so the narrow side has to have at least
        // that many lanes to hand one to each. A 64-bit scalar reinterpreted as eight fp8 lanes
        // is the case that gets here: one lane cannot be cut in two.
        let (wide, narrow) = match in_vec > max {
            true => (in_vec, out_vec),
            false => (out_vec, in_vec),
        };
        let factor = wide / max;
        assert!(
            narrow >= factor,
            "Cannot unroll a reinterpret between a {in_vec}-lane and a {out_vec}-lane vector: the \
             {wide}-lane side unrolls into {factor} pieces and the {narrow}-lane side has no lane \
             to give each. Reinterpret through a vector of at least {factor} lanes"
        );
        state.result.ir_changed |= IRStatus::Changed;
        let op = self.get_operation();

        if in_vec > max {
            let piece_ty = lanes_type(ctx, result, out_vec / factor);
            let pieces = state.mappings.get(&input).expect("Should exist").clone();
            let converted = pieces
                .into_iter()
                .map(|piece| {
                    let new_op = ReinterpretCastOp::new(ctx, piece_ty, piece);
                    new_op.get_operation().insert_before(ctx, op);
                    new_op.get_result(ctx)
                })
                .collect::<Vec<_>>();
            let lanes = converted
                .iter()
                .flat_map(|piece| split_lanes(ctx, *piece, op))
                .collect();
            let result_ty = result.get_type(ctx);
            let joined = CompositeConstructOp::new(ctx, result_ty, lanes);
            joined.get_operation().insert_before(ctx, op);
            result.replace_all_uses_with(ctx, &joined.get_result(ctx));
        } else {
            let piece_ty = unroll_ty(ctx, result, max);
            let lanes = split_lanes(ctx, input, op);
            let new_results = lanes
                .chunks(in_vec / factor)
                .map(|chunk| {
                    let piece = match chunk {
                        [lane] => *lane,
                        lanes => {
                            let ty = lanes_type(ctx, input, lanes.len());
                            let joined = CompositeConstructOp::new(ctx, ty, lanes.to_vec());
                            joined.get_operation().insert_before(ctx, op);
                            joined.get_result(ctx)
                        }
                    };
                    let new_op = ReinterpretCastOp::new(ctx, piece_ty, piece);
                    new_op.get_operation().insert_before(ctx, op);
                    new_op.get_result(ctx)
                })
                .collect();
            state.mappings.insert(result, new_results);
        }
        state.to_erase.push(op);
    }
}

/// The type of `lanes` lanes of `value`'s element: the bare scalar for one lane.
fn lanes_type(ctx: &mut Context, value: Value, lanes: usize) -> TypeHandle {
    assert!(lanes > 0, "Reinterpret pieces must hold at least one lane");
    let scalar = value.scalar_ty(ctx);
    match lanes {
        1 => scalar,
        lanes => VectorType::get(ctx, scalar, lanes).into(),
    }
}

/// Extracts every lane of `value` before `anchor`; a scalar is its own single lane.
fn split_lanes(ctx: &mut Context, value: Value, anchor: Ptr<Operation>) -> Vec<Value> {
    if !value.is_vector(ctx) {
        return vec![value];
    }
    (0..value.vector_size(ctx))
        .map(|lane| {
            let extract = CompositeExtractOp::new(ctx, value, lane);
            extract.get_operation().insert_before(ctx, anchor);
            extract.get_result(ctx)
        })
        .collect()
}

#[op_interface_impl]
impl CustomUnrollOp for matrix::LoadOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let source = self.source(ctx);
        if source.vector_size(ctx) > state.max_vector_size {
            state.result.ir_changed |= IRStatus::Changed;
            let new_source = state.mappings.get(&source).expect("should exist")[0];
            source.replace_use_with(ctx, self.source_as_use(ctx), &new_source);
        }
    }
}

#[op_interface_impl]
impl CustomUnrollOp for matrix::StoreOp {
    fn unroll(&self, ctx: &mut Context, state: &mut UnrollState) {
        let dest = self.destination(ctx);
        if dest.vector_size(ctx) > state.max_vector_size {
            state.result.ir_changed |= IRStatus::Changed;
            let new_dest = state.mappings.get(&dest).expect("should exist")[0];
            dest.replace_use_with(ctx, self.destination_as_use(ctx), &new_dest);
        }
    }
}

pub struct UnrollState {
    mappings: Mappings,
    to_erase: Vec<Ptr<Operation>>,
    max_vector_size: VectorSize,
    result: PassResult,
    rewriter: PassRewriter,
}

#[pass_name]
impl Pass for UnrollPass {
    fn run(
        &mut self,
        op: Ptr<Operation>,
        ctx: &mut Context,
        _analyses: &mut AnalysisManager,
    ) -> Result<PassResult> {
        self.unroll_func(ctx, op);

        let mut state = UnrollState {
            mappings: Default::default(),
            to_erase: Default::default(),
            max_vector_size: self.max_vector_size,
            result: Default::default(),
            rewriter: PassRewriter::default(),
        };

        walk_op(
            ctx,
            &mut state,
            &WALKCONFIG_PREORDER_FORWARD,
            op,
            |ctx, state, node| {
                if let IRNode::Operation(op) = node {
                    let unroll_opds = op
                        .operands(ctx)
                        .iter()
                        .any(|it| should_unroll(ctx, it, state.max_vector_size));
                    let unroll_res = op
                        .results(ctx)
                        .iter()
                        .any(|it| should_unroll(ctx, it, state.max_vector_size));
                    let dyn_op = op.dyn_op(ctx);

                    if let Some(custom) = op_cast::<dyn CustomUnrollOp>(&*dyn_op) {
                        custom.unroll(ctx, state);
                    } else if unroll_opds || unroll_res {
                        state.result.ir_changed |= IRStatus::Changed;
                        unroll_default(ctx, state, op);
                    }
                }
            },
        );

        while !state.to_erase.is_empty() {
            // Pop the next op that no longer has uses. This ensures we always start at the end of
            // the def-use chain
            let next = state
                .to_erase
                .iter()
                .position(|it| !it.deref(ctx).has_use())
                .expect("Erased ops should only have uses in other erased ops");
            let op = state.to_erase.remove(next);
            state.rewriter.erase_operation(ctx, op);
        }

        Ok(state.result)
    }
}

impl UnrollPass {
    fn unroll_func(&self, ctx: &mut Context, op: Ptr<Operation>) {
        let func = op.as_op::<FuncOp>(ctx).expect("Should be func");
        let entry_block = func.get_entry_block(ctx);
        let func_ty = func.get_attr_func_type(ctx).unwrap().get_type(ctx);
        let func_ty = func_ty.deref(ctx);
        let func_ty = func_ty.downcast_ref::<FunctionType>().unwrap();

        let mut new_func_inputs = vec![];

        for (i, arg) in func_ty.arg_types().into_iter().enumerate() {
            if should_unroll(ctx, arg, self.max_vector_size) {
                let new_ty = unroll_ty(ctx, arg, self.max_vector_size);
                new_func_inputs.push(new_ty);
                let block_arg = entry_block.deref(ctx).get_argument(i);
                block_arg.set_type(ctx, new_ty);
            } else {
                new_func_inputs.push(arg);
            }
        }

        let new_func_ty = FunctionType::get(ctx, new_func_inputs, func_ty.res_types()).to_handle();
        func.set_attr_func_type(ctx, new_func_ty.into());
    }
}

fn unroll_default(ctx: &mut Context, state: &mut UnrollState, op: Ptr<Operation>) {
    let values = op.operands(ctx).into_iter().chain(op.results(ctx));
    let current_vec = values.map(|it| try_get_vec(ctx, it)).max().unwrap();
    let factor = current_vec / state.max_vector_size;
    let dyn_op = op.dyn_op(ctx);
    let rematerialize = try_cast_op!(dyn_op, ctx, dyn TriviallyUnrollable);
    let new_out_ty = op
        .results(ctx)
        .into_iter()
        .map(|it| unroll_ty(ctx, it, state.max_vector_size))
        .collect::<Vec<_>>();
    let mut new_results = vec![];

    for unroll_idx in 0..factor {
        let opds = op.operands(ctx).into_iter().map(|opd| {
            if should_unroll(ctx, opd, state.max_vector_size) {
                state.mappings.get(&opd).expect("Should have mapping")[unroll_idx]
            } else {
                opd
            }
        });
        let attrs = op.deref(ctx).attributes.clone();
        let new_op = rematerialize.materialize(ctx, new_out_ty.clone(), opds.collect(), attrs);
        new_results.extend(new_op.deref(ctx).results());
        new_op.insert_before(ctx, op);
    }

    if !new_results.is_empty() {
        state
            .mappings
            .insert(op.deref(ctx).get_result(0), new_results);
    }
    state.to_erase.push(op);
}

fn should_unroll(ctx: &Context, value: impl Typed, max_vector_size: usize) -> bool {
    let ty = value.get_type(ctx).deref(ctx);
    if !type_impls::<dyn UnrollableType>(&*ty) {
        return false;
    }
    let Some(vector_size) = value.try_get_vector_size(ctx) else {
        return false;
    };
    vector_size > max_vector_size
}

fn try_get_vec(ctx: &Context, value: impl Typed) -> usize {
    value.try_get_vector_size(ctx).unwrap_or(1)
}

fn unroll_ty(ctx: &Context, ty: impl Typed, vectorization: usize) -> TypeHandle {
    let ty = ty.get_type(ctx).deref(ctx);
    type_cast::<dyn UnrollableType>(&*ty)
        .expect("Should be unrollable")
        .with_vector_size(ctx, vectorization)
}