vyre-lower 0.4.1

//! Loop-invariant code motion.
//!
//! Source-of-truth: `PERF_ROADMAP_2026-05-01.md` section A item A17.
//!
//! For each `StructuredForLoop` in the kernel body, walk the loop body
//! and identify ops whose operand chain doesn't depend on the loop
//! variable. Hoist them out of the loop body into the parent body
//! immediately before the loop op.
//!
//! Hoisted ops are assigned fresh parent-body ids, loop-body operands
//! are rewritten to those fresh ids, and child literal-pool references
//! are merged into the parent literal pool. That keeps the per-body id
//! invariant intact while still removing pure loop-invariant work from
//! hot loop bodies.

use std::collections::{BTreeMap, BTreeSet};

use crate::{KernelBody, KernelDescriptor, KernelOp, KernelOpKind};

/// LICM with a correct cross-body id rewrite.
///
/// For each loop, invariant ops are moved to the parent body before
/// the loop op. Hoisted ops receive fresh result ids that do not
/// collide with any parent-body id. Operand references inside the
/// remaining loop body are rewritten to use the new ids. Literal
/// pool entries referenced by hoisted `Literal` ops are merged into
/// the parent body's literal pool.
#[must_use]
pub fn licm(desc: &KernelDescriptor) -> KernelDescriptor {
    let mut out = desc.clone();
    out.body = licm_body(&out.body);
    out
}

fn licm_body(body: &KernelBody) -> KernelBody {
    let mut new_ops = Vec::with_capacity(body.ops.len());
    let mut new_children = body.child_bodies.clone();
    let mut new_literals = body.literals.clone();
    // Fresh ids for hoisted ops start above every id already in the
    // original parent body. Since we monotonically increase this
    // value, successive hoists in the same body also avoid collision.
    // Scan the body AND every transitive child for the max result id
    // so the fresh id we allocate for hoisted ops cannot collide with
    // an id produced inside a nested loop body. The pre-existing
    // `body.ops.iter().filter_map(|o| o.result).max()` form misses
    // child-body ids; on deep IR like `ast_shunting_yard` this
    // produced `DuplicateResultId(318)` at verify time.
    let mut next_free_id = {
        let mut max = 0u32;
        let mut found = false;
        fn walk(b: &KernelBody, max: &mut u32, found: &mut bool) {
            for op in &b.ops {
                if let Some(r) = op.result {
                    if !*found || r > *max {
                        *max = r;
                        *found = true;
                    }
                }
            }
            for child in &b.child_bodies {
                walk(child, max, found);
            }
        }
        walk(body, &mut max, &mut found);
        if found {
            max.wrapping_add(1)
        } else {
            0
        }
    };
    // Soundness: when a hoisted op produces a result, ALL downstream
    // operands in the SAME parent body that referenced the original
    // (loop-body) id need to point at the new parent-body id. Without
    // this, a `StoreGlobal` (or any op) appearing after the loop in
    // the parent body keeps referencing the dead loop-body id and
    // descriptor verify rejects with `DanglingResultRef`.
    //
    // Note that `remap_body_ids` already rewrites the LOOP body and
    // any ops we push AFTER the loop op need this same treatment in
    // the parent — that's what `parent_id_map` captures.
    let mut parent_id_map = BTreeMap::<u32, u32>::new();

    for op in &body.ops {
        if let KernelOpKind::StructuredForLoop { .. } = &op.kind {
            // Find the body child id (operand 2 per descriptor contract).
            if let Some(body_idx) = op.operands.get(2).copied() {
                if let Some(child) = body.child_bodies.get(body_idx as usize).cloned() {
                    // Hoist invariant ops from the child body.
                    let (hoisted, remaining) = hoist_invariants(&child);
                    if !hoisted.is_empty() {
                        // 1. Allocate fresh parent-body ids for every
                        //    hoisted result so we never collide.
                        let mut id_map = BTreeMap::<u32, u32>::new();
                        for h_op in &hoisted {
                            if let Some(r) = h_op.result {
                                id_map.insert(r, next_free_id);
                                next_free_id = next_free_id.wrapping_add(1);
                            }
                        }

                        // 2. Merge literals referenced by hoisted ops
                        //    into the parent literal pool and build a
                        //    child-index → parent-index map.
                        let mut lit_map = BTreeMap::<u32, u32>::new();
                        let mut renumbered_hoisted = Vec::with_capacity(hoisted.len());
                        for mut h_op in hoisted {
                            if matches!(h_op.kind, KernelOpKind::Literal) {
                                if let Some(&old_idx) = h_op.operands.first() {
                                    // Soundness: skip the rewrite when
                                    // the source literal is missing.
                                    // Previously the captured `idx` was
                                    // committed to `lit_map` even when
                                    // the conditional `push` was a no-op,
                                    // so the rewritten Literal op pointed
                                    // at a pool slot that was never
                                    // populated → `LiteralPoolOutOfRange`
                                    // at descriptor verify time.
                                    if let Some(val) = child.literals.get(old_idx as usize) {
                                        let new_idx =
                                            *lit_map.entry(old_idx).or_insert_with(|| {
                                                let idx = new_literals.len() as u32;
                                                new_literals.push(val.clone());
                                                idx
                                            });
                                        h_op.operands[0] = new_idx;
                                    }
                                    // When the source literal is missing
                                    // we leave the operand alone — the
                                    // op was already invalid in the
                                    // source body; LICM does not have to
                                    // fabricate a slot for it.
                                }
                            }
                            // Rewrite result-id refs inside the hoisted op.
                            h_op.operands = h_op
                                .operands
                                .iter()
                                .enumerate()
                                .map(|(pos, val)| {
                                    if operand_is_result_reference(&h_op.kind, pos) {
                                        *id_map.get(val).unwrap_or(val)
                                    } else {
                                        *val
                                    }
                                })
                                .collect();
                            h_op.result = h_op.result.map(|r| *id_map.get(&r).unwrap_or(&r));
                            renumbered_hoisted.push(h_op);
                        }

                        // 3. Recurse into the remaining body, then remap
                        //    any references to the old hoisted ids so the
                        //    child body still points at the parent results.
                        let recursed = licm_body(&remaining);
                        let remapped = remap_body_ids(&recursed, &id_map);

                        new_children[body_idx as usize] = remapped;
                        new_ops.extend(renumbered_hoisted);
                        new_ops.push(op.clone());
                        // Carry the id remap forward so any parent-body
                        // op AFTER this loop that referenced the old
                        // (now-hoisted) child id picks up the new
                        // parent-body result id. Without this, the
                        // post-loop reader (e.g. a `StoreGlobal`
                        // consuming the loop accumulator) would point
                        // at a result id that no longer exists in the
                        // descriptor → `DanglingResultRef` at verify.
                        for (&old, &new) in id_map.iter() {
                            parent_id_map.insert(old, new);
                        }
                        continue;
                    }
                    // No invariants to hoist — just recurse into child.
                    let recursed = licm_body(&remaining);
                    new_children[body_idx as usize] = recursed;
                    new_ops.push(op.clone());
                    continue;
                }
            }
        }
        // Recurse into other structured-control-flow children too.
        match &op.kind {
            KernelOpKind::StructuredIfThen
            | KernelOpKind::StructuredIfThenElse
            | KernelOpKind::StructuredBlock
            | KernelOpKind::Region { .. } => {
                for child_id in op.operands.iter() {
                    if let Some(child) = body.child_bodies.get(*child_id as usize) {
                        let recursed = licm_body(child);
                        new_children[*child_id as usize] = recursed;
                    }
                }
            }
            _ => {}
        }
        // Apply the accumulated parent_id_map to this op's
        // result-reference operands before pushing it. This rewrites
        // post-loop readers to point at the new parent-body result
        // ids that LICM created when it hoisted invariants out of an
        // earlier loop in this same parent body.
        let mut rewritten = op.clone();
        if !parent_id_map.is_empty() {
            for (pos, val) in rewritten.operands.iter_mut().enumerate() {
                if operand_is_result_reference(&op.kind, pos) {
                    if let Some(&new) = parent_id_map.get(val) {
                        *val = new;
                    }
                }
            }
        }
        new_ops.push(rewritten);
    }

    // Recurse into all OTHER children (those not touched above) so
    // nested loops in non-StructuredForLoop children also get LICM'd.
    let final_children: Vec<KernelBody> = new_children.into_iter().map(|c| licm_body(&c)).collect();

    KernelBody {
        ops: new_ops,
        child_bodies: final_children,
        literals: new_literals,
    }
}

/// Recursively apply `id_map` to every result-reference operand in
/// `body` and all nested child bodies. Result ids of ops themselves
/// are left unchanged — only operand refs are rewritten.
fn remap_body_ids(body: &KernelBody, id_map: &BTreeMap<u32, u32>) -> KernelBody {
    let new_ops: Vec<KernelOp> = body
        .ops
        .iter()
        .map(|op| {
            let new_operands: Vec<u32> = op
                .operands
                .iter()
                .enumerate()
                .map(|(pos, val)| {
                    if operand_is_result_reference(&op.kind, pos) {
                        *id_map.get(val).unwrap_or(val)
                    } else {
                        *val
                    }
                })
                .collect();
            KernelOp {
                kind: op.kind.clone(),
                operands: new_operands,
                result: op.result.map(|r| *id_map.get(&r).unwrap_or(&r)),
            }
        })
        .collect();
    let new_children: Vec<KernelBody> = body
        .child_bodies
        .iter()
        .map(|c| remap_body_ids(c, id_map))
        .collect();
    KernelBody {
        ops: new_ops,
        child_bodies: new_children,
        literals: body.literals.clone(),
    }
}

/// Split a loop body into (invariant_ops_to_hoist, loop_dependent_ops).
/// Invariants don't depend (transitively) on any value produced inside
/// the body, and have no side effects.
fn hoist_invariants(body: &KernelBody) -> (Vec<KernelOp>, KernelBody) {
    // Phase-1 conservative rule: an op is loop-dependent if any of its
    // result-id-position operands references a result-id produced
    // earlier in this body. The loop variable doesn't have a result-id
    // (it's implicit) so we treat ANY operand whose value is produced
    // inside the body as loop-dependent.
    // Include result ids produced ANYWHERE inside the loop body — top
    // level AND any descendant child bodies. Hoisting an op that
    // references a child-body-local id (e.g. a phi-Select after a
    // StructuredIfThen) would dangle the reference outside the loop.
    fn collect_all_ids(body: &KernelBody, out: &mut BTreeSet<u32>) {
        for op in &body.ops {
            for r in op.result_ids() {
                out.insert(r);
            }
        }
        for child in &body.child_bodies {
            collect_all_ids(child, out);
        }
    }
    let mut body_local_ids: BTreeSet<u32> = BTreeSet::new();
    collect_all_ids(body, &mut body_local_ids);
    let mut dependent_ids = BTreeSet::<u32>::new();

    let mut invariants = Vec::new();
    let mut remaining_ops = Vec::new();

    for op in &body.ops {
        if !is_pure(&op.kind) {
            // Side effects → never hoist.
            dependent_ids.extend(op.result_ids());
            remaining_ops.push(op.clone());
            continue;
        }
        let depends = op.operands.iter().enumerate().any(|(pos, val)| {
            operand_is_result_reference(&op.kind, pos)
                && (dependent_ids.contains(val) || body_local_ids.contains(val))
        });
        if !depends {
            // Op is invariant. Hoist.
            invariants.push(op.clone());
            // Its result-id is now produced OUTSIDE the loop body, so
            // it's NOT in the dependent set.
        } else {
            // Loop-dependent.
            dependent_ids.extend(op.result_ids());
            remaining_ops.push(op.clone());
        }
    }

    let new_body = KernelBody {
        ops: remaining_ops,
        child_bodies: body.child_bodies.clone(),
        literals: body.literals.clone(),
    };
    (invariants, new_body)
}

fn is_pure(kind: &KernelOpKind) -> bool {
    !matches!(
        kind,
        KernelOpKind::StoreGlobal
            | KernelOpKind::StoreShared
            | KernelOpKind::Barrier { .. }
            | KernelOpKind::Atomic { .. }
            | KernelOpKind::AsyncLoad { .. }
            | KernelOpKind::AsyncStore { .. }
            | KernelOpKind::AsyncWait { .. }
            | KernelOpKind::Trap { .. }
            | KernelOpKind::Resume { .. }
            | KernelOpKind::IndirectDispatch { .. }
            | KernelOpKind::Return
            | KernelOpKind::StructuredIfThen
            | KernelOpKind::StructuredIfThenElse
            | KernelOpKind::StructuredForLoop { .. }
            | KernelOpKind::StructuredBlock
            | KernelOpKind::Region { .. }
            | KernelOpKind::Call { .. }
            | KernelOpKind::OpaqueExpr { .. }
            | KernelOpKind::OpaqueNode { .. }
            // Loop induction variable is not invariant.
            | KernelOpKind::LoopIndex { .. }
            | KernelOpKind::LoopCarrier { .. }
            | KernelOpKind::LoopCarrierEnd { .. }
            | KernelOpKind::LoopCarrierFinal { .. }
            // Loads aren't safely hoistable — the underlying buffer
            // could be written by another thread between iterations.
            | KernelOpKind::LoadGlobal
            | KernelOpKind::LoadShared
            | KernelOpKind::LoadConstant
    )
}

fn operand_is_result_reference(kind: &KernelOpKind, pos: usize) -> bool {
    use KernelOpKind::*;
    match kind {
        Literal => false,
        LocalInvocationId | GlobalInvocationId | WorkgroupId => false,
        SubgroupLocalId | SubgroupSize | LoopIndex { .. } => false,
        LoopCarrier { .. } | LoopCarrierEnd { .. } => pos == 0,
        LoopCarrierFinal { .. } => false,
        LoadGlobal | LoadShared | LoadConstant => pos != 0,
        BufferLength => false,
        StoreGlobal | StoreShared => pos != 0,
        BinOpKind(_) | UnOpKind(_) | Fma | MatrixMma { .. } | Select | Cast { .. } => true,
        Atomic { .. } => pos != 0,
        SubgroupBallot | SubgroupShuffle | SubgroupAdd => true,
        StructuredIfThen | StructuredIfThenElse => pos == 0,
        StructuredForLoop { .. } => pos != 2,
        StructuredBlock | Region { .. } => false,
        Return | Barrier { .. } => false,
        AsyncLoad { .. } | AsyncStore { .. } => pos >= 2,
        AsyncWait { .. } => false,
        Trap { .. } => pos == 0,
        Resume { .. } => false,
        IndirectDispatch { .. } => false,
        Call { .. } => true,
        OpaqueExpr { .. } | OpaqueNode { .. } => true,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        BindingLayout, Dispatch, KernelBody, KernelDescriptor, KernelOp, KernelOpKind, LiteralValue,
    };
    use vyre_foundation::ir::BinOp;

    fn empty_kernel_with_loop(loop_body: KernelBody) -> KernelDescriptor {
        KernelDescriptor {
            id: "loopy".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StructuredForLoop {
                            loop_var: "i".into(),
                        },
                        operands: vec![0, 1, 0],
                        result: None,
                    },
                ],
                child_bodies: vec![loop_body],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(64)],
            },
        }
    }

    #[test]
    fn licm_on_empty_kernel_is_noop() {
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![],
                child_bodies: vec![],
                literals: vec![],
            },
        };
        let out = licm(&desc);
        assert!(out.body.ops.is_empty());
    }

    #[test]
    fn licm_kernel_with_no_loop_is_noop() {
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::BinOpKind(BinOp::Add),
                        operands: vec![0, 1],
                        result: Some(2),
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(3), LiteralValue::U32(4)],
            },
        };
        let out = licm(&desc);
        assert_eq!(out.body.ops.len(), 3);
    }

    #[test]
    fn licm_hoists_constant_out_of_loop() {
        // for i in 0..64 { lit(99); /* nothing else */ }
        // The Literal in the loop body has no operand dependencies; it's
        // a constant, hoistable.
        let desc = empty_kernel_with_loop(KernelBody {
            ops: vec![KernelOp {
                kind: KernelOpKind::Literal,
                operands: vec![0],
                result: Some(2),
            }],
            child_bodies: vec![],
            literals: vec![LiteralValue::U32(99)],
        });
        let out = licm(&desc);
        // Outer body started with 3 ops. After LICM, the hoisted Literal
        // appears before the StructuredForLoop, making it 4.
        assert_eq!(out.body.ops.len(), 4);
        // Hoisted Literal is at position 2 (before the loop op which is now at 3).
        assert!(matches!(out.body.ops[2].kind, KernelOpKind::Literal));
        assert!(matches!(
            out.body.ops[3].kind,
            KernelOpKind::StructuredForLoop { .. }
        ));
        // Loop body should now be empty (the hoisted op moved out).
        let loop_body = out.body.child_bodies[0].clone();
        assert!(loop_body.ops.is_empty());
    }

    #[test]
    fn licm_does_not_hoist_loop_dependent_op() {
        // Loop body: lit(5) (invariant); add(lit, lit) (invariant);
        //   plus a synthetic op that depends on the previous op
        //   (so the chain is loop-dependent).
        // Phase 1: lit and add are both hoistable (neither depends on
        // anything produced inside the body's chain).
        let desc = empty_kernel_with_loop(KernelBody {
            ops: vec![
                KernelOp {
                    kind: KernelOpKind::Literal,
                    operands: vec![0],
                    result: Some(10),
                },
                KernelOp {
                    kind: KernelOpKind::Literal,
                    operands: vec![1],
                    result: Some(11),
                },
                // This Add USES the prior Literals — but those are now
                // hoisted, so the Add itself can also be hoisted.
                KernelOp {
                    kind: KernelOpKind::BinOpKind(BinOp::Add),
                    operands: vec![10, 11],
                    result: Some(12),
                },
            ],
            child_bodies: vec![],
            literals: vec![LiteralValue::U32(5), LiteralValue::U32(7)],
        });
        let out = licm(&desc);
        // Phase-1 conservative: Add depends on result 10 (which was
        // produced inside the loop body) → the Add is NOT hoisted on
        // the first pass. Both literals ARE hoisted.
        // After: outer ops = init lits (0,1) + 2 hoisted lits (10,11)
        // + loop op = 5. Loop body has just the Add.
        assert_eq!(out.body.ops.len(), 5);
        let loop_body = &out.body.child_bodies[0];
        assert_eq!(loop_body.ops.len(), 1);
        assert!(matches!(
            loop_body.ops[0].kind,
            KernelOpKind::BinOpKind(BinOp::Add)
        ));
    }

    #[test]
    fn licm_does_not_hoist_load() {
        // Loads are unsafely hoistable (other threads may write between
        // iterations). Phase 1 forbids hoisting all loads.
        use crate::{BindingSlot, BindingVisibility, MemoryClass};
        use vyre_foundation::ir::DataType;
        let desc = KernelDescriptor {
            id: "load_loop".into(),
            bindings: BindingLayout {
                slots: vec![BindingSlot {
                    slot: 0,
                    element_type: DataType::U32,
                    element_count: None,
                    memory_class: MemoryClass::Global,
                    visibility: BindingVisibility::ReadOnly,
                    name: "buf".into(),
                }],
            },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StructuredForLoop {
                            loop_var: "i".into(),
                        },
                        operands: vec![0, 1, 0],
                        result: None,
                    },
                ],
                child_bodies: vec![KernelBody {
                    ops: vec![
                        KernelOp {
                            kind: KernelOpKind::Literal,
                            operands: vec![0],
                            result: Some(10),
                        },
                        // LoadGlobal — should NOT hoist.
                        KernelOp {
                            kind: KernelOpKind::LoadGlobal,
                            operands: vec![0, 10],
                            result: Some(11),
                        },
                    ],
                    child_bodies: vec![],
                    literals: vec![LiteralValue::U32(0)],
                }],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(8)],
            },
        };
        let out = licm(&desc);
        // The Literal hoists; the Load stays.
        let loop_body = &out.body.child_bodies[0];
        let has_load = loop_body
            .ops
            .iter()
            .any(|o| matches!(o.kind, KernelOpKind::LoadGlobal));
        assert!(has_load, "Load must stay inside the loop");
    }

    #[test]
    fn licm_does_not_hoist_store() {
        // Stores are side-effecting. Hoisting changes semantics.
        use crate::{BindingSlot, BindingVisibility, MemoryClass};
        use vyre_foundation::ir::DataType;
        let desc = KernelDescriptor {
            id: "store_loop".into(),
            bindings: BindingLayout {
                slots: vec![BindingSlot {
                    slot: 0,
                    element_type: DataType::U32,
                    element_count: None,
                    memory_class: MemoryClass::Global,
                    visibility: BindingVisibility::WriteOnly,
                    name: "out".into(),
                }],
            },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StructuredForLoop {
                            loop_var: "i".into(),
                        },
                        operands: vec![0, 1, 0],
                        result: None,
                    },
                ],
                child_bodies: vec![KernelBody {
                    ops: vec![
                        KernelOp {
                            kind: KernelOpKind::Literal,
                            operands: vec![0],
                            result: Some(10),
                        },
                        KernelOp {
                            kind: KernelOpKind::Literal,
                            operands: vec![1],
                            result: Some(11),
                        },
                        KernelOp {
                            kind: KernelOpKind::StoreGlobal,
                            operands: vec![0, 10, 11],
                            result: None,
                        },
                    ],
                    child_bodies: vec![],
                    literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
                }],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(8)],
            },
        };
        let out = licm(&desc);
        let loop_body = &out.body.child_bodies[0];
        let has_store = loop_body
            .ops
            .iter()
            .any(|o| matches!(o.kind, KernelOpKind::StoreGlobal));
        assert!(has_store, "Store must stay inside the loop");
    }

    #[test]
    fn licm_is_idempotent() {
        let desc = empty_kernel_with_loop(KernelBody {
            ops: vec![KernelOp {
                kind: KernelOpKind::Literal,
                operands: vec![0],
                result: Some(2),
            }],
            child_bodies: vec![],
            literals: vec![LiteralValue::U32(99)],
        });
        let once = licm(&desc);
        let twice = licm(&once);
        assert_eq!(once.body.ops.len(), twice.body.ops.len());
        assert_eq!(
            once.body.child_bodies[0].ops.len(),
            twice.body.child_bodies[0].ops.len()
        );
    }

    #[test]
    fn licm_handles_no_for_loop_op_gracefully() {
        // Body with StructuredIfThen but no for-loop — should be a noop on the loop side.
        let desc = KernelDescriptor {
            id: "if_only".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::StructuredIfThen,
                        operands: vec![0, 0],
                        result: None,
                    },
                ],
                child_bodies: vec![KernelBody {
                    ops: vec![],
                    child_bodies: vec![],
                    literals: vec![],
                }],
                literals: vec![LiteralValue::Bool(true)],
            },
        };
        let out = licm(&desc);
        // Outer ops stay at 2; nothing to hoist.
        assert_eq!(out.body.ops.len(), 2);
    }

    #[test]
    fn public_licm_hoists_invariants() {
        // The public `licm` API now performs a correct cross-body
        // hoist with id renumbering and literal-pool merging.
        let desc = empty_kernel_with_loop(KernelBody {
            ops: vec![KernelOp {
                kind: KernelOpKind::Literal,
                operands: vec![0],
                result: Some(2),
            }],
            child_bodies: vec![],
            literals: vec![LiteralValue::U32(99)],
        });
        let out = licm(&desc);
        // Outer body started with 3 ops. After LICM, the hoisted Literal
        // appears before the StructuredForLoop, making it 4.
        assert_eq!(out.body.ops.len(), 4, "hoisted literal adds one parent op");
        assert!(
            matches!(out.body.ops[2].kind, KernelOpKind::Literal),
            "hoisted op is at position 2"
        );
        assert!(
            matches!(out.body.ops[3].kind, KernelOpKind::StructuredForLoop { .. }),
            "loop op follows hoisted literal"
        );
        // Loop body should now be empty (the hoisted op moved out).
        assert!(
            out.body.child_bodies[0].ops.is_empty(),
            "child body empty after hoist"
        );
        // The parent literal pool gained the child's literal.
        assert_eq!(
            out.body.literals.len(),
            desc.body.literals.len() + 1,
            "parent literals grew by one"
        );
    }
}