vyre-lower 0.4.1

//! Loop unrolling for small constant-bound loops.
//!
//! Source-of-truth: `PERF_ROADMAP_2026-05-01.md` section A item A29
//! (loop strip-mining family). This is the unconditional-unroll
//! special case: when both bounds are compile-time constants AND the
//! iteration count is small (≤ 4 by default), inline N copies of the
//! body in sequence and strip the loop.
//!
//! Rules:
//! - Both `lo` and `hi` operands must point at `Literal(U32)` ops.
//! - `hi - lo` must be ≤ `MAX_UNROLL_COUNT` (default 4).
//! - Nested child bodies are duplicated, remapped into the parent
//!   `child_bodies` table, and result ids are freshened across the
//!   whole duplicated subtree.

use std::collections::BTreeMap;

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

pub const MAX_UNROLL_COUNT: u32 = 4;

#[must_use]
pub fn loop_unroll(desc: &KernelDescriptor) -> KernelDescriptor {
    let mut out = desc.clone();
    out.body = unroll_body(&out.body);
    out
}

fn unroll_body(body: &KernelBody) -> KernelBody {
    // Map result-id → constant U32 value (for ops whose op is Literal(U32)).
    let lit_u32: BTreeMap<u32, u32> = body
        .ops
        .iter()
        .filter_map(|op| match (&op.kind, op.result, op.operands.first()) {
            (KernelOpKind::Literal, Some(r), Some(pool_idx)) => {
                match body.literals.get(*pool_idx as usize) {
                    Some(LiteralValue::U32(v)) => Some((r, *v)),
                    _ => None,
                }
            }
            _ => None,
        })
        .collect();

    // Compute the next free result-id (highest + 1).
    let mut next_id: u32 = body
        .ops
        .iter()
        .flat_map(KernelOp::result_ids)
        .max()
        .map(|m| m + 1)
        .unwrap_or(0);

    // Also consider child bodies' result-ids when allocating new ids,
    // since unrolled bodies may reuse them.
    for child in &body.child_bodies {
        for op in &child.ops {
            for r in op.result_ids() {
                if r >= next_id {
                    next_id = r + 1;
                }
            }
        }
    }

    let mut new_ops: Vec<KernelOp> = Vec::with_capacity(body.ops.len());
    let mut new_children = body.child_bodies.clone();
    // Literal-pool fix: when we inline ops from `child` into `body`,
    // every `Literal` op's first operand is a pool index into the
    // CHILD's literal pool. We merge the child's literals into the
    // parent's pool and rewrite the inlined ops' pool indices on the
    // way through. Without this the inlined `Literal` op points at a
    // pool slot that doesn't exist in the parent (LiteralPoolOutOfRange
    // at verify time). Surfaced on `c11_build_vast_nodes` with nt=1.
    let mut new_literals: Vec<LiteralValue> = body.literals.clone();

    for op in &body.ops {
        if let KernelOpKind::StructuredForLoop { .. } = &op.kind {
            if op.operands.len() != 3 {
                new_ops.push(op.clone());
                continue;
            }
            let lo_id = op.operands[0];
            let hi_id = op.operands[1];
            let body_idx = op.operands[2] as usize;
            let lo = lit_u32.get(&lo_id).copied();
            let hi = lit_u32.get(&hi_id).copied();
            let child = body.child_bodies.get(body_idx).cloned();
            let unrollable = match (lo, hi, &child) {
                (Some(lo), Some(hi), Some(c)) => {
                    let count = hi.saturating_sub(lo);
                    count <= MAX_UNROLL_COUNT && safe_to_unroll(c)
                }
                _ => false,
            };
            if unrollable {
                let lo = lo.unwrap();
                let hi = hi.unwrap();
                let child = child.unwrap();
                // Match the `unrollable` check exactly (saturating_sub
                // also returns 0 when hi < lo). A plain `hi - lo` here
                // would underflow in release mode and produce a near-4B
                // iteration count — OOM-killed the fuzz harness on the
                // first generator run that hit this shape.
                let count = hi.saturating_sub(lo);
                // Build a per-iteration literal-pool map: child's
                // pool index → parent's pool index after merging.
                // Two child Literal ops referencing the same child
                // pool slot share a parent pool slot (de-duplicated).
                for _iter in 0..count {
                    let (renumbered, new_next) = renumber_body(&child, next_id);
                    next_id = new_next;
                    let child_offset = new_children.len() as u32;
                    new_children.extend(renumbered.child_bodies);
                    let mut pool_map: BTreeMap<u32, u32> = BTreeMap::new();
                    let child_literals = child.literals.clone();
                    new_ops.extend(renumbered.ops.into_iter().map(|mut op| {
                        remap_top_level_child_body_operands(&mut op, child_offset);
                        if matches!(op.kind, KernelOpKind::Literal) {
                            if let Some(child_idx) = op.operands.first().copied() {
                                let parent_idx = *pool_map.entry(child_idx).or_insert_with(|| {
                                    let next_pool = new_literals.len() as u32;
                                    if let Some(value) =
                                        child_literals.get(child_idx as usize).cloned()
                                    {
                                        new_literals.push(value);
                                        next_pool
                                    } else {
                                        // Source literal missing — fall back to
                                        // the original index. Verifier will
                                        // catch a malformed source body.
                                        child_idx
                                    }
                                });
                                if !op.operands.is_empty() {
                                    op.operands[0] = parent_idx;
                                }
                            }
                        }
                        op
                    }));
                }
                continue;
            }
        }
        new_ops.push(op.clone());
    }

    // Recursively unroll child bodies that weren't already inlined.
    let mut final_children: Vec<KernelBody> = Vec::with_capacity(new_children.len());
    for c in new_children.drain(..) {
        final_children.push(unroll_body(&c));
    }

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

/// Unroll-safety check. Rejects (a) malformed child-body references and
/// (b) bodies that reference SSA ids defined in the body that contained
/// the loop. (b) is the scope-leak case: inlining the child into the
/// grandparent yanks the loop variable's outer-scope refs out of scope.
/// On `c11_build_vast_nodes` with nt=1 this fires constantly.
fn safe_to_unroll(child: &KernelBody) -> bool {
    let valid_child_refs = child.ops.iter().all(|op| {
        child_body_operands(&op.kind).all(|pos| {
            op.operands
                .get(pos)
                .is_some_and(|idx| (*idx as usize) < child.child_bodies.len())
        })
    }) && child.child_bodies.iter().all(safe_to_unroll);
    if !valid_child_refs {
        return false;
    }
    let mut produced = std::collections::BTreeSet::new();
    collect_produced_ids_inclusive(child, &mut produced);
    body_refs_only(child, &produced)
}

fn collect_produced_ids_inclusive(body: &KernelBody, out: &mut std::collections::BTreeSet<u32>) {
    for op in &body.ops {
        for r in op.result_ids() {
            out.insert(r);
        }
    }
    for c in &body.child_bodies {
        collect_produced_ids_inclusive(c, out);
    }
}

fn body_refs_only(body: &KernelBody, produced: &std::collections::BTreeSet<u32>) -> bool {
    for op in &body.ops {
        for (pos, &operand) in op.operands.iter().enumerate() {
            if !operand_is_result_reference(&op.kind, pos) {
                continue;
            }
            if !produced.contains(&operand) {
                return false;
            }
        }
    }
    for c in &body.child_bodies {
        if !body_refs_only(c, produced) {
            return false;
        }
    }
    true
}

/// Renumber every result-id in `body` starting at `next_id`. Operand
/// references that match an old result-id are rewritten to the new id.
/// Returns the renumbered body + the next free id after the rename.
fn renumber_body(body: &KernelBody, mut next_id: u32) -> (KernelBody, u32) {
    let mut id_map = BTreeMap::<u32, u32>::new();
    collect_result_renames(body, &mut id_map, &mut next_id);
    (rewrite_body_with_renames(body, &id_map), next_id)
}

fn collect_result_renames(body: &KernelBody, id_map: &mut BTreeMap<u32, u32>, next_id: &mut u32) {
    for op in &body.ops {
        for result in op.result_ids() {
            id_map.insert(result, *next_id);
            *next_id += 1;
        }
    }
    for child in &body.child_bodies {
        collect_result_renames(child, id_map, next_id);
    }
}

fn rewrite_body_with_renames(body: &KernelBody, id_map: &BTreeMap<u32, u32>) -> KernelBody {
    let new_ops: Vec<KernelOp> = body
        .ops
        .iter()
        .map(|op| rewrite_op_with_renames(op, id_map))
        .collect();
    let child_bodies = body
        .child_bodies
        .iter()
        .map(|child| rewrite_body_with_renames(child, id_map))
        .collect();
    KernelBody {
        ops: new_ops,
        child_bodies,
        literals: body.literals.clone(),
    }
}

fn rewrite_op_with_renames(op: &KernelOp, id_map: &BTreeMap<u32, u32>) -> KernelOp {
    let 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,
        result: op.result.map(|r| *id_map.get(&r).unwrap_or(&r)),
    }
}

fn remap_top_level_child_body_operands(op: &mut KernelOp, child_offset: u32) {
    for pos in child_body_operands(&op.kind) {
        if let Some(operand) = op.operands.get_mut(pos) {
            *operand = operand.saturating_add(child_offset);
        }
    }
}

fn child_body_operands(kind: &KernelOpKind) -> impl Iterator<Item = usize> + '_ {
    use KernelOpKind::*;
    let positions: &'static [usize] = match kind {
        StructuredIfThen => &[1],
        StructuredIfThenElse => &[1, 2],
        StructuredForLoop { .. } => &[2],
        StructuredBlock | Region { .. } => &[0],
        _ => &[],
    };
    positions.iter().copied()
}

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 loop_with_body(
        lo: u32,
        hi: u32,
        body_ops: Vec<KernelOp>,
        body_lits: Vec<LiteralValue>,
    ) -> KernelDescriptor {
        KernelDescriptor {
            id: "loop".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![KernelBody {
                    ops: body_ops,
                    child_bodies: vec![],
                    literals: body_lits,
                }],
                literals: vec![LiteralValue::U32(lo), LiteralValue::U32(hi)],
            },
        }
    }

    #[test]
    fn empty_kernel_unchanged() {
        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 = loop_unroll(&desc);
        assert!(out.body.ops.is_empty());
    }

    #[test]
    fn loop_with_count_4_unrolled() {
        // for i in 0..4 { lit(99) }  →  4 inlined Literal ops
        let body_op = vec![KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![0],
            result: Some(10),
        }];
        let body_lits = vec![LiteralValue::U32(99)];
        let desc = loop_with_body(0, 4, body_op, body_lits);
        let out = loop_unroll(&desc);
        // Outer ops: [Lit(0), Lit(4)] + 4 inlined Literal copies = 6 total.
        assert_eq!(out.body.ops.len(), 6);
        assert!(out
            .body
            .ops
            .iter()
            .all(|o| !matches!(o.kind, KernelOpKind::StructuredForLoop { .. })));
    }

    #[test]
    fn loop_with_count_above_threshold_not_unrolled() {
        // for i in 0..10 { ... }  →  loop kept (count > 4)
        let body_op = vec![KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![0],
            result: Some(10),
        }];
        let body_lits = vec![LiteralValue::U32(99)];
        let desc = loop_with_body(0, 10, body_op, body_lits);
        let out = loop_unroll(&desc);
        assert_eq!(out.body.ops.len(), 3); // unchanged
        assert!(out
            .body
            .ops
            .iter()
            .any(|o| matches!(o.kind, KernelOpKind::StructuredForLoop { .. })));
    }

    #[test]
    fn loop_with_runtime_bounds_not_unrolled() {
        // for i in tid..hi { ... }  →  loop kept (lo not literal)
        let desc = KernelDescriptor {
            id: "runtime".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::LocalInvocationId,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StructuredForLoop {
                            loop_var: "i".into(),
                        },
                        operands: vec![0, 1, 0],
                        result: None,
                    },
                ],
                child_bodies: vec![KernelBody {
                    ops: vec![],
                    child_bodies: vec![],
                    literals: vec![],
                }],
                literals: vec![LiteralValue::U32(8)],
            },
        };
        let out = loop_unroll(&desc);
        assert!(out
            .body
            .ops
            .iter()
            .any(|o| matches!(o.kind, KernelOpKind::StructuredForLoop { .. })));
    }

    #[test]
    fn loop_with_zero_count_strips_loop() {
        // for i in 5..5 { ... }  →  empty
        let body_op = vec![KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![0],
            result: Some(10),
        }];
        let body_lits = vec![LiteralValue::U32(99)];
        let desc = loop_with_body(5, 5, body_op, body_lits);
        let out = loop_unroll(&desc);
        // Outer: [Lit(5), Lit(5)] only — no inlined body, no loop op.
        assert_eq!(out.body.ops.len(), 2);
        assert!(out
            .body
            .ops
            .iter()
            .all(|o| !matches!(o.kind, KernelOpKind::StructuredForLoop { .. })));
    }

    #[test]
    fn loop_with_count_1_inlines_once() {
        let body_op = vec![KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![0],
            result: Some(10),
        }];
        let body_lits = vec![LiteralValue::U32(99)];
        let desc = loop_with_body(0, 1, body_op, body_lits);
        let out = loop_unroll(&desc);
        // Outer: [Lit(0), Lit(1), inlined Literal] = 3 ops
        assert_eq!(out.body.ops.len(), 3);
    }

    #[test]
    fn loop_with_nested_if_is_unrolled_and_child_indices_are_remapped() {
        let desc = KernelDescriptor {
            id: "nested_if".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![KernelBody {
                    ops: vec![
                        KernelOp {
                            kind: KernelOpKind::Literal,
                            operands: vec![0],
                            result: Some(10),
                        },
                        KernelOp {
                            kind: KernelOpKind::StructuredIfThen,
                            operands: vec![10, 0],
                            result: None,
                        },
                    ],
                    child_bodies: vec![KernelBody {
                        ops: vec![KernelOp {
                            kind: KernelOpKind::Literal,
                            operands: vec![0],
                            result: Some(20),
                        }],
                        child_bodies: vec![],
                        literals: vec![LiteralValue::U32(9)],
                    }],
                    literals: vec![LiteralValue::Bool(true)],
                }],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(2)],
            },
        };
        let out = loop_unroll(&desc);
        assert!(out
            .body
            .ops
            .iter()
            .all(|o| !matches!(o.kind, KernelOpKind::StructuredForLoop { .. })));
        let if_indices: Vec<u32> = out
            .body
            .ops
            .iter()
            .filter(|op| matches!(op.kind, KernelOpKind::StructuredIfThen))
            .map(|op| op.operands[1])
            .collect();
        assert_eq!(if_indices.len(), 2);
        assert_ne!(if_indices[0], if_indices[1]);
        assert!(if_indices
            .iter()
            .all(|idx| (*idx as usize) < out.body.child_bodies.len()));
    }

    #[test]
    fn unrolled_body_renumbers_result_ids() {
        // for i in 0..3 { lit; binop }
        // Each iteration produces 2 fresh result-ids; pre-loop produced 0..1.
        let body_ops = vec![
            KernelOp {
                kind: KernelOpKind::Literal,
                operands: vec![0],
                result: Some(10),
            },
            KernelOp {
                kind: KernelOpKind::BinOpKind(BinOp::Add),
                operands: vec![10, 10],
                result: Some(11),
            },
        ];
        let body_lits = vec![LiteralValue::U32(7)];
        let desc = loop_with_body(0, 3, body_ops, body_lits);
        let out = loop_unroll(&desc);
        // 3 iterations × 2 ops each = 6 inlined ops; outer adds 2 → 8 total.
        assert_eq!(out.body.ops.len(), 8);
        // Collect all result-ids of the inlined ops; they should all be distinct.
        let inlined_ids: Vec<u32> = out.body.ops[2..].iter().filter_map(|o| o.result).collect();
        let mut sorted = inlined_ids.clone();
        sorted.sort();
        sorted.dedup();
        assert_eq!(
            inlined_ids.len(),
            sorted.len(),
            "all unrolled result-ids must be distinct: {inlined_ids:?}"
        );
    }

    #[test]
    fn loop_unroll_is_idempotent() {
        let body_op = vec![KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![0],
            result: Some(10),
        }];
        let body_lits = vec![LiteralValue::U32(99)];
        let desc = loop_with_body(0, 3, body_op, body_lits);
        let once = loop_unroll(&desc);
        let twice = loop_unroll(&once);
        assert_eq!(once.body.ops.len(), twice.body.ops.len());
    }

    #[test]
    fn max_unroll_count_constant_is_documented() {
        assert_eq!(MAX_UNROLL_COUNT, 4);
    }
}