metal2vulkan 0.1.0

//! Byte-neutral responsibility split of the former monolith; see the parent module.

use super::*;
use crate::passes::stage_input::layout_ty_size_align;

/// Candidate Workgroup variable for the flat-word remodel.
pub(in crate::passes) struct WgFlatwordCand {
    pub(in crate::passes) var: Word,
    pub(in crate::passes) elem_words: u32,
    pub(in crate::passes) total_words: u32,
}

/// Validate a candidate: every use of the variable across the entry function must be a single-index
/// access chain whose index is word-addressed (stride `elem_words`) and whose result flows ONLY into
/// `uint` loads/stores. Returns the `(block, inst)` positions of the chains to retype, or `None` if any
/// use disqualifies the variable.
pub(in crate::passes) fn validate_wg_flatword(
    ctx: &Ctx,
    func: &Function,
    cand: &WgFlatwordCand,
    uint_ty: Word,
) -> Option<Vec<(usize, usize)>> {
    let mut chains: Vec<(usize, usize)> = Vec::new();
    let mut chain_ids: HashSet<Word> = HashSet::new();
    for (bi, block) in func.blocks.iter().enumerate() {
        for (ii, inst) in block.instructions.iter().enumerate() {
            let mentions = inst
                .operands
                .iter()
                .any(|o| matches!(o, Operand::IdRef(id) if *id == cand.var));
            if !mentions {
                continue;
            }
            // The variable may ONLY appear as the base of a single-index chain.
            if !matches!(inst.class.opcode, Op::InBoundsAccessChain | Op::AccessChain)
                || inst.operands.len() != 2
                || operand_id(inst, 0) != Some(cand.var)
            {
                return None;
            }
            let cid = inst.result_id?;
            let idx = operand_id(inst, 1)?;
            if !index_is_word_addressed(ctx, func, idx, cand.elem_words) {
                return None;
            }
            chains.push((bi, ii));
            chain_ids.insert(cid);
        }
    }
    if chains.is_empty() {
        return None;
    }
    // Every chain result must be consumed ONLY by a `uint` OpLoad / OpStore(uint object).
    for block in &func.blocks {
        for inst in &block.instructions {
            if inst
                .result_id
                .map(|r| chain_ids.contains(&r))
                .unwrap_or(false)
                && matches!(inst.class.opcode, Op::InBoundsAccessChain | Op::AccessChain)
            {
                continue; // the chain definition itself
            }
            match inst.class.opcode {
                Op::Load => {
                    if let Some(ptr) = operand_id(inst, 0) {
                        if chain_ids.contains(&ptr) && inst.result_type != Some(uint_ty) {
                            return None;
                        }
                    }
                }
                Op::Store => {
                    if let Some(ptr) = operand_id(inst, 0) {
                        if chain_ids.contains(&ptr) {
                            let obj_ty =
                                operand_id(inst, 1).and_then(|o| value_result_type(ctx, o));
                            if obj_ty != Some(uint_ty) {
                                return None;
                            }
                        }
                    }
                }
                _ => {
                    if inst
                        .operands
                        .iter()
                        .any(|o| matches!(o, Operand::IdRef(id) if chain_ids.contains(id)))
                    {
                        return None; // chain pointer escapes to some other consumer
                    }
                }
            }
        }
    }
    Some(chains)
}

/// Remodel a Workgroup (threadgroup) aggregate variable that the AIR addresses EXCLUSIVELY by flat-WORD
/// scalar indices as a flat `[totalWords x uint]` array.
///
/// The AIR allocates threadgroup scratch as `[K x S]` (S a struct/aggregate) but addresses it with a
/// FLAT WORD index `c + elem*W` (W = S's word footprint) and reads/writes 32-bit scalars through that
/// single index, bitcasting each value to/from `uint` itself. metal2vulkan faithfully models the variable as
/// `[K x S]`, so a single-index chain `%v %wordIdx` selects ELEMENT `wordIdx` (a whole struct) and the
/// `OpStore`/`OpLoad` of a `uint` through that struct pointer is rejected
/// ("OpStore Pointer's type does not match Object's type").
///
/// Threadgroup memory is shader-internal scratch (no host/device ABI, no golden) and SPIR-V Workgroup
/// layout is word-granular, so `[K x S]` and `[K*W x uint]` cover byte-IDENTICAL storage. The fix
/// re-declares the variable as `[K*W x uint]` and re-types each chain result to `_ptr_Workgroup_uint`;
/// the now-uint pointee matches the AIR's own uint loads/stores directly (the AIR already bitcast the
/// values to/from uint), so word `c+elem*W` lands at byte `4*(c+elem*W)` — the exact byte the typed
/// `[K x S]` model places at element `elem`, struct-byte `4c`.
///
/// Byte-EXACT by construction (same byte address; the flat-uint element type is what the AIR's own
/// bitcasts already produce) AND floor-SAFE by construction: only remodels a Workgroup `[K x S]` (S an
/// aggregate) variable whose EVERY use is a single-index chain feeding a `uint` load/store — a valid
/// module's threadgroup struct access uses multi-index chains whose load/store type matches the member,
/// so it never matches. The index of every chain must be a proven word-addressed form (`const`,
/// `OpIMul(v, W)`, or `OpIAdd(const, OpIMul(v, W))`) with stride `W` equal to S's word footprint
/// (`layout_ty_size_align(S)/4`); any other shape leaves the whole variable untouched rather than guess.
pub(in crate::passes) fn remodel_workgroup_flatword_aggregate(ctx: &mut Ctx, entry_idx: usize) {
    let uint_ty = ctx.ty_uint();

    let mut ptr_info: HashMap<Word, (StorageClass, Word)> = HashMap::new();
    let mut defs: HashMap<Word, Instruction> = HashMap::new();
    for inst in ctx
        .new_globals
        .iter()
        .chain(ctx.module.types_global_values.iter())
    {
        if let Some(result) = inst.result_id {
            defs.insert(result, inst.clone());
        }
        if inst.class.opcode == Op::TypePointer {
            if let (Some(id), Some(Operand::StorageClass(s)), Some(Operand::IdRef(p))) =
                (inst.result_id, inst.operands.first(), inst.operands.get(1))
            {
                ptr_info.insert(id, (*s, *p));
            }
        }
    }

    // Collect Workgroup `[K x E]` variables with E an aggregate of a whole number of words.
    let mut cands: Vec<WgFlatwordCand> = Vec::new();
    for inst in ctx
        .new_globals
        .iter()
        .chain(ctx.module.types_global_values.iter())
    {
        if inst.class.opcode != Op::Variable {
            continue;
        }
        let (Some(var), Some(ptr_ty)) = (inst.result_id, inst.result_type) else {
            continue;
        };
        let Some(&(StorageClass::Workgroup, pointee)) = ptr_info.get(&ptr_ty) else {
            continue;
        };
        let Some(arr_def) = defs.get(&pointee) else {
            continue;
        };
        if arr_def.class.opcode != Op::TypeArray {
            continue;
        }
        let (Some(Operand::IdRef(elem_ty)), Some(Operand::IdRef(len_c))) =
            (arr_def.operands.first(), arr_def.operands.get(1))
        else {
            continue;
        };
        let Some(k) = const_u32(ctx, *len_c) else {
            continue;
        };
        // E must be an aggregate — a scalar array is already flat-word addressable and never mismatches.
        if direct_scalar_width(ctx, *elem_ty).is_some() {
            continue;
        }
        let (elem_bytes, _) = layout_ty_size_align(ctx, *elem_ty, &defs);
        if elem_bytes == 0 || elem_bytes % 4 != 0 {
            continue;
        }
        let elem_words = elem_bytes / 4;
        let Some(total_words) = k.checked_mul(elem_words) else {
            continue;
        };
        cands.push(WgFlatwordCand {
            var,
            elem_words,
            total_words,
        });
    }

    for cand in cands {
        let Some(chains) =
            validate_wg_flatword(ctx, &ctx.module.functions[entry_idx], &cand, uint_ty)
        else {
            continue;
        };
        // Commit: build `[total_words x uint]` + its Workgroup pointer types, retype the variable and
        // every chain result. The AIR's own uint loads/stores then match the uint pointee directly.
        //
        // The new array/pointer/length-constant defs must DEFINE-BEFORE-USE the variable's OpVariable
        // (an existing global). `finalize` appends `new_globals` at the END of the type section — fine
        // for types used only in function bodies, but the retyped variable would then forward-reference
        // them. So drain exactly the defs these `ty_*` calls synthesized and splice them in immediately
        // BEFORE the variable. (`new_chain_ptr` is used only by the chains in the body, but ordering it
        // here too is harmless — every dependency, `uint`, is defined far earlier in the section.)
        let new_glob_start = ctx.new_globals.len();
        let arr_ty = ctx.ty_array(uint_ty, cand.total_words);
        let new_var_ptr = ctx.ty_ptr(StorageClass::Workgroup, arr_ty);
        let new_chain_ptr = ctx.ty_ptr(StorageClass::Workgroup, uint_ty);
        let synthesized: Vec<Instruction> = ctx.new_globals.drain(new_glob_start..).collect();
        // The OpVariable may live in EITHER the already-emitted `types_global_values` or the pending
        // `new_globals` (finalize drains `new_globals` to the end of the type section). Splice the
        // synthesized defs immediately before it IN THE SAME vector, so they define-before-use it. (Its
        // only new dependency, `uint`, is defined far earlier; the array/pointer/length defs all sit in
        // `synthesized` in dependency order.)
        let is_var =
            |i: &Instruction| i.class.opcode == Op::Variable && i.result_id == Some(cand.var);
        if let Some(p) = ctx.module.types_global_values.iter().position(is_var) {
            ctx.module.types_global_values[p].result_type = Some(new_var_ptr);
            let tail = ctx.module.types_global_values.split_off(p);
            ctx.module.types_global_values.extend(synthesized);
            ctx.module.types_global_values.extend(tail);
        } else if let Some(p) = ctx.new_globals.iter().position(is_var) {
            ctx.new_globals[p].result_type = Some(new_var_ptr);
            let tail = ctx.new_globals.split_off(p);
            ctx.new_globals.extend(synthesized);
            ctx.new_globals.extend(tail);
        } else {
            // Variable vanished (unexpected) — drop the synthesized defs back so nothing dangles.
            ctx.new_globals.extend(synthesized);
        }
        for &(bi, ii) in &chains {
            ctx.module.functions[entry_idx].blocks[bi].instructions[ii].result_type =
                Some(new_chain_ptr);
        }
    }
}

/// Whether `op` is an atomic instruction whose FIRST operand is the pointer it operates on (the whole
/// `OpAtomic*` family relevant to the float-as-uint idiom — min/max/exchange and the integer RMWs).
pub(in crate::passes) fn is_atomic_pointer_op(op: Op) -> bool {
    matches!(
        op,
        Op::AtomicLoad
            | Op::AtomicStore
            | Op::AtomicExchange
            | Op::AtomicCompareExchange
            | Op::AtomicCompareExchangeWeak
            | Op::AtomicIIncrement
            | Op::AtomicIDecrement
            | Op::AtomicIAdd
            | Op::AtomicISub
            | Op::AtomicSMin
            | Op::AtomicUMin
            | Op::AtomicSMax
            | Op::AtomicUMax
            | Op::AtomicAnd
            | Op::AtomicOr
            | Op::AtomicXor
    )
}

/// `Some(id)` of the existing 32-bit `OpTypeFloat` / unsigned `OpTypeInt` in the module, without
/// synthesizing one (so a module with no such type — and therefore no candidate — is left untouched).
pub(in crate::passes) fn existing_scalar_type(ctx: &Ctx, opcode: Op) -> Option<Word> {
    ctx.new_globals
        .iter()
        .chain(ctx.module.types_global_values.iter())
        .find(|inst| {
            inst.class.opcode == opcode && inst.operands.first() == Some(&Operand::LiteralBit32(32))
        })
        .and_then(|inst| inst.result_id)
}

/// The validated retype plan for one Workgroup `array<float, K>` atomically accessed as a 32-bit int.
pub(in crate::passes) struct WgFloatAtomicPlan {
    pub(in crate::passes) var: Word,
    pub(in crate::passes) k: u32,
    pub(in crate::passes) int_ty: Word, // the 32-bit int type (signed or unsigned) the array is reinterpreted as
    pub(in crate::passes) chain_ids: HashSet<Word>, // chain results to retype to `_ptr_Workgroup_<int_ty>`
    pub(in crate::passes) chains: Vec<(usize, usize)>, // their (block, instr) positions
    pub(in crate::passes) bitcast_ids: HashSet<Word>, // `OpBitcast %_ptr_Workgroup_<int_ty>` results to DROP
    pub(in crate::passes) bitcast_to_chain: HashMap<Word, Word>, // bitcast result -> the chain it aliased (atomic repoint)
}

/// Retype a Workgroup `array<float, K>` that is atomically accessed AS uint (Metal's float-as-signed-int
/// `atomic_fetch_min/max` idiom) to `array<uint, K>`, repoint the atomics directly at the uint chain, and
/// value-bitcast its non-atomic float load/store accesses. The AIR lowers a `threadgroup float bins[K]`
/// reduced with `atomic_fetch_min/max` to:
/// ```text
///   %c  = OpInBoundsAccessChain %_ptr_Workgroup_float %V %idx
///   %cu = OpBitcast %_ptr_Workgroup_uint %c            ; ILLEGAL: pointer OpBitcast under Logical
///   %r  = OpAtomicSMin %uint %cu %scope %sem %val      ; float-as-signed-int min
///   ; — and plain reads/writes of the same bins —
///   %f  = OpLoad %float %c        |    OpStore %c %fval
/// ```
/// The pointer `OpBitcast` is illegal under Logical addressing ("Instruction may not have a logical
/// pointer operand"). Threadgroup memory is shader-internal scratch (no host/device ABI, never compared
/// against a golden) and a 32-bit float and uint share bit-IDENTICAL storage, so retyping the array to
/// `array<uint, K>` is BYTE-EXACT: the atomics then point straight at the uint element (no pointer
/// bitcast), and each plain float load/store is rewritten to load/store the uint and `OpBitcast` the
/// VALUE (a pure bit reinterpret).
///
/// **Floor-safe by construction:** fires ONLY on a Workgroup `array<float, K>` variable at least one of
/// whose element chains is consumed by an `OpBitcast %_ptr_Workgroup_uint` — a pointer bitcast a
/// valid/banked module never contains (it is itself illegal) — and only when EVERY use of every chain is
/// one of {float-as-uint atomic via such a bitcast, plain float `OpLoad`, plain float `OpStore`}; any
/// other consumer leaves the variable untouched. Decides purely from IR structure, never a shader name.
pub(in crate::passes) fn remodel_workgroup_floatarray_atomic_as_uint(
    ctx: &mut Ctx,
    entry_idx: usize,
) -> Result<(), String> {
    let Some(float_ty) = existing_scalar_type(ctx, Op::TypeFloat) else {
        return Ok(());
    };

    // id -> (storage class, pointee) for every pointer type; id -> def for everything.
    let mut ptr_info: HashMap<Word, (StorageClass, Word)> = HashMap::new();
    let mut defs: HashMap<Word, Instruction> = HashMap::new();
    for inst in ctx
        .new_globals
        .iter()
        .chain(ctx.module.types_global_values.iter())
    {
        if let Some(result) = inst.result_id {
            defs.insert(result, inst.clone());
        }
        if inst.class.opcode == Op::TypePointer {
            if let (Some(id), Some(Operand::StorageClass(s)), Some(Operand::IdRef(p))) =
                (inst.result_id, inst.operands.first(), inst.operands.get(1))
            {
                ptr_info.insert(id, (*s, *p));
            }
        }
    }

    // The 32-bit integer types (signed AND unsigned) the array may be atomically reinterpreted as.
    // Metal's `air.atomic.local.add.s.i32` lowers the bitcast to a SIGNED `int` pointer, so the retype
    // target is whichever 32-bit int the atomics/loads actually use — not necessarily the canonical uint.
    let int32_types: HashSet<Word> = defs
        .values()
        .filter(|inst| {
            inst.class.opcode == Op::TypeInt
                && inst.operands.first() == Some(&Operand::LiteralBit32(32))
        })
        .filter_map(|inst| inst.result_id)
        .collect();

    // Workgroup `array<float, K>` variables.
    let mut cands: Vec<(Word, u32)> = Vec::new();
    for inst in ctx
        .new_globals
        .iter()
        .chain(ctx.module.types_global_values.iter())
    {
        if inst.class.opcode != Op::Variable {
            continue;
        }
        let (Some(var), Some(ptr_ty)) = (inst.result_id, inst.result_type) else {
            continue;
        };
        let Some(&(StorageClass::Workgroup, pointee)) = ptr_info.get(&ptr_ty) else {
            continue;
        };
        let Some(arr_def) = defs.get(&pointee) else {
            continue;
        };
        if arr_def.class.opcode != Op::TypeArray {
            continue;
        }
        let (Some(Operand::IdRef(elem_ty)), Some(Operand::IdRef(len_c))) =
            (arr_def.operands.first(), arr_def.operands.get(1))
        else {
            continue;
        };
        if *elem_ty != float_ty {
            continue;
        }
        let Some(k) = const_u32(ctx, *len_c) else {
            continue;
        };
        cands.push((var, k));
    }

    for (var, k) in cands {
        let Some(plan) = validate_wg_float_atomic(
            ctx,
            &ctx.module.functions[entry_idx],
            var,
            k,
            float_ty,
            &int32_types,
            &ptr_info,
        ) else {
            continue;
        };

        // Retype the variable to `array<int_ty, K>` (int_ty = the detected 32-bit int the atomics use,
        // signed or unsigned). The synthesized array/pointer defs must DEFINE-BEFORE-USE the existing
        // OpVariable, so splice them immediately before it (mirrors `remodel_workgroup_flatword_aggregate`).
        let int_ty = plan.int_ty;
        let new_glob_start = ctx.new_globals.len();
        let arr_ty = ctx.ty_array(int_ty, plan.k);
        let new_var_ptr = ctx.ty_ptr(StorageClass::Workgroup, arr_ty);
        let chain_ptr_uint = ctx.ty_ptr(StorageClass::Workgroup, int_ty);
        let synthesized: Vec<Instruction> = ctx.new_globals.drain(new_glob_start..).collect();
        let is_var =
            |i: &Instruction| i.class.opcode == Op::Variable && i.result_id == Some(plan.var);
        if let Some(p) = ctx.module.types_global_values.iter().position(is_var) {
            ctx.module.types_global_values[p].result_type = Some(new_var_ptr);
            let tail = ctx.module.types_global_values.split_off(p);
            ctx.module.types_global_values.extend(synthesized);
            ctx.module.types_global_values.extend(tail);
        } else if let Some(p) = ctx.new_globals.iter().position(is_var) {
            ctx.new_globals[p].result_type = Some(new_var_ptr);
            let tail = ctx.new_globals.split_off(p);
            ctx.new_globals.extend(synthesized);
            ctx.new_globals.extend(tail);
        } else {
            ctx.new_globals.extend(synthesized);
            continue;
        }

        // Retype each chain result to a uint element pointer.
        for &(bi, ii) in &plan.chains {
            ctx.module.functions[entry_idx].blocks[bi].instructions[ii].result_type =
                Some(chain_ptr_uint);
        }

        // Rewrite the body: drop the pointer bitcasts, repoint atomics straight at the uint chain, and
        // value-bitcast the plain float loads/stores.
        let n_blocks = ctx.module.functions[entry_idx].blocks.len();
        for bi in 0..n_blocks {
            let insts = ctx.module.functions[entry_idx].blocks[bi]
                .instructions
                .clone();
            let mut out = Vec::with_capacity(insts.len());
            for mut inst in insts {
                // Drop the now-redundant `OpBitcast %_ptr_Workgroup_uint %chain`.
                if inst.class.opcode == Op::Bitcast
                    && inst
                        .result_id
                        .map(|r| plan.bitcast_ids.contains(&r))
                        .unwrap_or(false)
                {
                    continue;
                }
                // Plain load through a retyped (now int_ty) chain. A FLOAT load reinterprets: load int_ty,
                // bitcast the value to float. An int_ty load is already native through the retyped chain —
                // leave it (a `OpBitcast int_ty->int_ty` would be a no-op spirv-val rejects).
                if inst.class.opcode == Op::Load {
                    if let Some(ptr) = operand_id(&inst, 0) {
                        if plan.chain_ids.contains(&ptr) && inst.result_type != Some(int_ty) {
                            let (rt, rid) = (
                                inst.result_type
                                    .ok_or("wg float atomic: load missing result type")?,
                                inst.result_id
                                    .ok_or("wg float atomic: load missing result id")?,
                            );
                            let tmp = ctx.module.fresh_id();
                            out.push(Instruction::new(
                                Op::Load,
                                Some(int_ty),
                                Some(tmp),
                                vec![Operand::IdRef(ptr)],
                            ));
                            out.push(Instruction::new(
                                Op::Bitcast,
                                Some(rt),
                                Some(rid),
                                vec![Operand::IdRef(tmp)],
                            ));
                            continue;
                        }
                    }
                }
                // Plain store through a retyped chain. A FLOAT object reinterprets: bitcast float->int_ty,
                // store it. An int_ty object stores natively through the retyped chain — leave it.
                if inst.class.opcode == Op::Store {
                    if let Some(ptr) = operand_id(&inst, 0) {
                        if plan.chain_ids.contains(&ptr) {
                            let fval = operand_id(&inst, 1)
                                .ok_or("wg float atomic: store missing value operand")?;
                            let stored_int = value_result_type(ctx, fval) == Some(int_ty);
                            if !stored_int {
                                let tmp = ctx.module.fresh_id();
                                out.push(Instruction::new(
                                    Op::Bitcast,
                                    Some(int_ty),
                                    Some(tmp),
                                    vec![Operand::IdRef(fval)],
                                ));
                                out.push(Instruction::new(
                                    Op::Store,
                                    None,
                                    None,
                                    vec![Operand::IdRef(ptr), Operand::IdRef(tmp)],
                                ));
                                continue;
                            }
                        }
                    }
                }
                // Repoint any operand that referenced a dropped bitcast onto the underlying int_ty chain
                // (the atomic's pointer operand).
                for op in inst.operands.iter_mut() {
                    if let Operand::IdRef(id) = op {
                        if let Some(&c) = plan.bitcast_to_chain.get(id) {
                            *op = Operand::IdRef(c);
                        }
                    }
                }
                out.push(inst);
            }
            ctx.module.functions[entry_idx].blocks[bi].instructions = out;
        }
    }
    Ok(())
}

/// Validate that Workgroup `array<float, K>` variable `var` is accessed ONLY as the float-as-int-atomic
/// idiom (+ plain float and plain int_ty load/store), returning the retype plan or `None`. `int_ty` (the
/// 32-bit int the array is reinterpreted as) is DETECTED from the pointer bitcasts: Metal's
/// `atomic_fetch_min/max` lower through an UNSIGNED `uint` bitcast, but `air.atomic.local.add.s.i32`
/// (an integer histogram living in a threadgroup `float` array) lowers through a SIGNED `int` bitcast —
/// both are byte-exact 32-bit reinterprets of the float element.
pub(in crate::passes) fn validate_wg_float_atomic(
    ctx: &Ctx,
    func: &Function,
    var: Word,
    k: u32,
    float_ty: Word,
    int32_types: &HashSet<Word>,
    ptr_info: &HashMap<Word, (StorageClass, Word)>,
) -> Option<WgFloatAtomicPlan> {
    // Pass 1: collect the single-index float chains rooted at `var`.
    let mut chains: Vec<(usize, usize)> = Vec::new();
    let mut chain_ids: HashSet<Word> = HashSet::new();
    for (bi, block) in func.blocks.iter().enumerate() {
        for (ii, inst) in block.instructions.iter().enumerate() {
            let mentions_var = inst
                .operands
                .iter()
                .any(|o| matches!(o, Operand::IdRef(id) if *id == var));
            if !mentions_var {
                continue;
            }
            // The variable may ONLY appear as the base of a single-index chain reaching a float.
            if !matches!(inst.class.opcode, Op::InBoundsAccessChain | Op::AccessChain)
                || inst.operands.len() != 2
                || operand_id(inst, 0) != Some(var)
            {
                return None;
            }
            let result_ty = inst.result_type?;
            if ptr_info.get(&result_ty) != Some(&(StorageClass::Workgroup, float_ty)) {
                return None;
            }
            chains.push((bi, ii));
            chain_ids.insert(inst.result_id?);
        }
    }
    if chains.is_empty() {
        return None;
    }

    // Pass 2a: collect the chain->`_ptr_Workgroup_<int>` pointer bitcasts and PIN the single 32-bit int
    // type they reinterpret to (`int_ty`). All such bitcasts must agree on `int_ty`; a chain bitcast to a
    // non-(Workgroup,int32) pointer disqualifies the variable.
    let mut bitcast_ids: HashSet<Word> = HashSet::new();
    let mut bitcast_to_chain: HashMap<Word, Word> = HashMap::new();
    let mut int_ty: Option<Word> = None;
    for block in &func.blocks {
        for inst in &block.instructions {
            if inst.class.opcode != Op::Bitcast {
                continue;
            }
            let Some(src) = operand_id(inst, 0) else {
                continue;
            };
            if !chain_ids.contains(&src) {
                continue;
            }
            let rt = inst.result_type?;
            let &(sc, pointee) = ptr_info.get(&rt)?;
            if sc != StorageClass::Workgroup || !int32_types.contains(&pointee) {
                return None;
            }
            if *int_ty.get_or_insert(pointee) != pointee {
                return None; // two different reinterpret int types — not modelled
            }
            let cu = inst.result_id?;
            bitcast_ids.insert(cu);
            bitcast_to_chain.insert(cu, src);
        }
    }
    // Floor gate: at least one pointer bitcast must exist — a valid/banked threadgroup float array never
    // contains a pointer `OpBitcast` (it is itself illegal under Logical addressing).
    let int_ty = int_ty?;
    if bitcast_ids.is_empty() {
        return None;
    }

    // Pass 2b: classify every NON-bitcast use of every chain. A chain may feed only a plain `OpLoad` /
    // `OpStore` whose value type is `float` (a reinterpret) or `int_ty` (native after retype). Anything
    // else that references a chain directly is an unmodelled escape.
    for block in &func.blocks {
        for inst in &block.instructions {
            if inst
                .result_id
                .map(|r| chain_ids.contains(&r))
                .unwrap_or(false)
                && matches!(inst.class.opcode, Op::InBoundsAccessChain | Op::AccessChain)
            {
                continue; // the chain definitions themselves
            }
            if inst.class.opcode == Op::Bitcast
                && inst
                    .result_id
                    .map(|r| bitcast_ids.contains(&r))
                    .unwrap_or(false)
            {
                continue; // the chain->int_ty pointer bitcasts (handled in Pass 2a)
            }
            match inst.class.opcode {
                Op::Load => {
                    if let Some(ptr) = operand_id(inst, 0) {
                        if chain_ids.contains(&ptr)
                            && inst.result_type != Some(float_ty)
                            && inst.result_type != Some(int_ty)
                        {
                            return None; // load of an unmodelled type through the chain
                        }
                    }
                }
                Op::Store => {
                    if let Some(ptr) = operand_id(inst, 0) {
                        if chain_ids.contains(&ptr) {
                            let obj_ty =
                                operand_id(inst, 1).and_then(|o| value_result_type(ctx, o));
                            if obj_ty != Some(float_ty) && obj_ty != Some(int_ty) {
                                return None; // store of an unmodelled object type through the chain
                            }
                        }
                    }
                }
                _ => {
                    if inst
                        .operands
                        .iter()
                        .any(|o| matches!(o, Operand::IdRef(id) if chain_ids.contains(id)))
                    {
                        return None;
                    }
                }
            }
        }
    }

    // Pass 3: the bitcast results may be consumed ONLY as the pointer (operand 0) of (a) an atomic (the
    // float-as-int idiom), (b) a plain `OpLoad %int_ty` (the histogram read-back), or (c) a plain
    // `OpStore` of an `int_ty` object. All are byte-exact after the array is retyped to `array<int_ty,K>`:
    // the dropped bitcast's operands are repointed onto the now-int_ty chain, so the atomic / load / store
    // runs natively on the int element. A non-int_ty load/store would not be valid post-retype.
    for block in &func.blocks {
        for inst in &block.instructions {
            if inst.class.opcode == Op::Bitcast
                && inst
                    .result_id
                    .map(|r| bitcast_ids.contains(&r))
                    .unwrap_or(false)
            {
                continue; // the bitcast definitions themselves
            }
            let references_bitcast = inst
                .operands
                .iter()
                .any(|o| matches!(o, Operand::IdRef(id) if bitcast_ids.contains(id)));
            if !references_bitcast {
                continue;
            }
            let pointer_is_bitcast =
                matches!(operand_id(inst, 0), Some(id) if bitcast_ids.contains(&id));
            let ok = match inst.class.opcode {
                op if is_atomic_pointer_op(op) => pointer_is_bitcast,
                Op::Load => pointer_is_bitcast && inst.result_type == Some(int_ty),
                Op::Store => {
                    pointer_is_bitcast
                        && operand_id(inst, 1).and_then(|o| value_result_type(ctx, o))
                            == Some(int_ty)
                }
                _ => false,
            };
            if !ok {
                return None;
            }
        }
    }

    Some(WgFloatAtomicPlan {
        var,
        k,
        int_ty,
        chain_ids,
        chains,
        bitcast_ids,
        bitcast_to_chain,
    })
}

/// `Some(M)` when `ty` is `vector<elem, M>` (element type EQUALS `elem`), or `Some(1)` when `ty` IS
/// `elem` itself — the lane count of a scalar-or-vector value over a fixed element type. `None`
/// otherwise (a different element type, or a non-scalar/non-vector).
pub(in crate::passes) fn vector_lanes_of(ctx: &Ctx, ty: Word, elem: Word) -> Option<u32> {
    if ty == elem {
        return Some(1);
    }
    let def = type_def_of(ctx, ty)?;
    if def.class.opcode != Op::TypeVector {
        return None;
    }
    let Some(Operand::IdRef(e)) = def.operands.first() else {
        return None;
    };
    if *e != elem {
        return None;
    }
    match def.operands.get(1) {
        Some(Operand::LiteralBit32(n)) => Some(*n),
        _ => None,
    }
}