vyre-lower 0.4.1

//! Lower a `vyre_foundation::Program` into a substrate-neutral
//! `KernelDescriptor`.
//!
//! This module is the shared boundary between high-level vyre IR and
//! emitter input. It preserves supported IR semantics in descriptor
//! form and returns an explicit [`LowerError`] when the input is
//! invalid.

use crate::descriptor::{
    BindingLayout, BindingSlot, BindingVisibility, Dispatch, KernelBody, KernelDescriptor,
    KernelOp, KernelOpKind, LiteralValue, MemoryClass, TRAP_SIDECAR_NAME, TRAP_SIDECAR_WORDS,
};
use crate::error::LowerError;
mod scope;
use rustc_hash::{FxHashMap, FxHashSet};
use scope::VarScope;
use std::sync::Arc;
use vyre_foundation::ir::model::node::node_op_id;
use vyre_foundation::ir::{
    AtomicOp, BufferAccess, BufferDecl, DataType, Expr, Ident, MemoryKind, Node, Program,
};

/// Maximum nested-body depth before lowering refuses with
/// `LowerError::NestingTooDeep`. 64 levels is generous; real programs
/// rarely exceed 10.
#[allow(dead_code)]
const MAX_NESTING_DEPTH: usize = 64;

/// Lower a vyre Program to the substrate-neutral kernel descriptor.
///
/// # Errors
///
/// Returns [`LowerError`] when the input references undeclared buffers,
/// exceeds the supported structured nesting depth, or uses an IR
/// construct with invalid operands.
pub fn lower(program: &Program) -> Result<KernelDescriptor, LowerError> {
    let mut ctx = LowerCtx::new(program)?;
    let mut body = empty_body_with_capacity(estimated_root_op_capacity(program));
    ctx.lower_nodes(program.entry(), &mut body, 0)?;
    if body_contains_trap(&body) {
        ctx.add_trap_sidecar_binding()?;
    }

    Ok(KernelDescriptor {
        id: fingerprint_id(program),
        bindings: BindingLayout {
            slots: ctx.bindings,
        },
        dispatch: Dispatch {
            workgroup_size: program.workgroup_size(),
        },
        body,
    })
}

struct LowerCtx {
    bindings: Vec<BindingSlot>,
    buffer_slots: FxHashMap<Ident, u32>,
    slot_memory_classes: FxHashMap<u32, MemoryClass>,
    scope: VarScope,
    next_value: u32,
}

impl LowerCtx {
    fn new(program: &Program) -> Result<Self, LowerError> {
        let mut bindings = Vec::with_capacity(program.buffers().len());
        let mut buffer_slots = FxHashMap::default();
        let mut slot_memory_classes = FxHashMap::default();
        // Soundness: split slot allocation by memory class. Host-bound
        // buffers (Global, Constant) keep their declared binding ids
        // because the dispatch path looks them up by the same
        // BufferDecl::binding() value. Workgroup/Scratch buffers are
        // SM-local — they don't get bound by the host — so they live
        // in a high range starting at WORKGROUP_SLOT_BASE that cannot
        // collide with host-bound slots. Without this split,
        // multiple `BufferDecl::workgroup(...)` calls (which all
        // default to binding=0) collided with the host-bound input
        // and forced the output's slot to be auto-renumbered, which
        // then broke the dispatch path's slot-id-keyed lookup.
        const WORKGROUP_SLOT_BASE: u32 = 1 << 24;
        let mut host_used_slots = FxHashSet::default();
        let mut host_next_free_slot = 0u32;
        let mut shared_next_slot = WORKGROUP_SLOT_BASE;
        for buffer in program.buffers() {
            let mc = memory_class(buffer)?;
            let slot = match mc {
                MemoryClass::Shared | MemoryClass::Scratch => {
                    let s = shared_next_slot;
                    shared_next_slot = shared_next_slot
                        .checked_add(1)
                        .ok_or(LowerError::OperandIdOverflow)?;
                    s
                }
                MemoryClass::Global | MemoryClass::Constant => {
                    let requested = buffer.binding();
                    let s = if host_used_slots.insert(requested) {
                        requested
                    } else {
                        while host_used_slots.contains(&host_next_free_slot)
                            || host_next_free_slot >= WORKGROUP_SLOT_BASE
                        {
                            host_next_free_slot = host_next_free_slot
                                .checked_add(1)
                                .ok_or(LowerError::OperandIdOverflow)?;
                        }
                        host_used_slots.insert(host_next_free_slot);
                        host_next_free_slot
                    };
                    while host_used_slots.contains(&host_next_free_slot)
                        || host_next_free_slot >= WORKGROUP_SLOT_BASE
                    {
                        host_next_free_slot = host_next_free_slot
                            .checked_add(1)
                            .ok_or(LowerError::OperandIdOverflow)?;
                    }
                    s
                }
            };
            buffer_slots.insert(Ident::from(Arc::clone(&buffer.name)), slot);
            slot_memory_classes.insert(slot, mc);
            bindings.push(BindingSlot {
                slot,
                element_type: buffer.element.clone(),
                element_count: (buffer.count != 0).then_some(buffer.count),
                memory_class: mc,
                visibility: binding_visibility(&buffer.access),
                name: buffer.name().to_owned(),
            });
        }
        bindings.sort_by_key(|slot| slot.slot);
        Ok(Self {
            bindings,
            buffer_slots,
            slot_memory_classes,
            scope: VarScope::default(),
            next_value: 0,
        })
    }

    fn lower_nodes(
        &mut self,
        nodes: &[Node],
        body: &mut KernelBody,
        depth: usize,
    ) -> Result<(), LowerError> {
        if depth > MAX_NESTING_DEPTH {
            return Err(LowerError::NestingTooDeep(depth));
        }
        for node in nodes {
            self.lower_node(node, body, depth)?;
        }
        Ok(())
    }

    fn lower_node(
        &mut self,
        node: &Node,
        body: &mut KernelBody,
        depth: usize,
    ) -> Result<(), LowerError> {
        match node {
            Node::Region {
                generator,
                body: region,
                ..
            } => self.lower_child_node(
                body,
                depth,
                region.as_ref(),
                KernelOpKind::Region {
                    generator: generator.shared_text(),
                },
            ),
            Node::Block(region) => {
                self.lower_child_node(body, depth, region, KernelOpKind::StructuredBlock)
            }
            Node::Let { name, value } | Node::Assign { name, value } => {
                let id = self.lower_expr(value, body)?;
                self.scope.bind(name.clone(), id);
                Ok(())
            }
            Node::Store {
                buffer,
                index,
                value,
            } => {
                let slot = self.buffer_slot(buffer)?;
                let index_id = self.lower_expr(index, body)?;
                let value_id = self.lower_expr(value, body)?;
                body.ops.push(KernelOp {
                    kind: self.store_kind(slot),
                    operands: vec![slot, index_id, value_id],
                    result: None,
                });
                Ok(())
            }
            Node::If {
                cond,
                then,
                otherwise,
            } => {
                let cond_id = self.lower_expr(cond, body)?;
                let incoming_scope = self.scope.snapshot();
                let mut then_body = empty_body_for_nodes(then);
                self.scope.restore(incoming_scope.clone());
                self.lower_nodes(then, &mut then_body, depth + 1)?;
                let then_exit_scope = self.scope.snapshot();
                let then_id = push_child(body, then_body)?;
                if otherwise.is_empty() {
                    self.scope.restore(incoming_scope.clone());
                    body.ops.push(KernelOp {
                        kind: KernelOpKind::StructuredIfThen,
                        operands: vec![cond_id, then_id],
                        result: None,
                    });
                    // Phi merge for vars rebound in the then-branch: emit a
                    // Select(cond, then_id, incoming_id) in the parent body
                    // so post-if reads see the conditionally-updated value.
                    // Inside loops the inner phi target is itself the loop
                    // carrier (LoopCarrier output), so the Select correctly
                    // computes "either prior carrier or carrier+delta", and
                    // the surrounding loop's `LoopCarrierEnd` commits the
                    // merged value back at iter end.
                    self.merge_if_then_scope(body, &incoming_scope, &then_exit_scope, cond_id)?;
                } else {
                    let mut else_body = empty_body_for_nodes(otherwise);
                    self.scope.restore(incoming_scope.clone());
                    self.lower_nodes(otherwise, &mut else_body, depth + 1)?;
                    let else_exit_scope = self.scope.snapshot();
                    let else_id = push_child(body, else_body)?;
                    self.scope.restore(incoming_scope.clone());
                    body.ops.push(KernelOp {
                        kind: KernelOpKind::StructuredIfThenElse,
                        operands: vec![cond_id, then_id, else_id],
                        result: None,
                    });
                    self.merge_if_else_scope(
                        body,
                        &incoming_scope,
                        &then_exit_scope,
                        &else_exit_scope,
                        cond_id,
                    )?;
                }
                Ok(())
            }
            Node::Loop {
                var,
                from,
                to,
                body: loop_body,
            } => {
                let from_id = self.lower_expr(from, body)?;
                let to_id = self.lower_expr(to, body)?;
                let mut child = empty_body_for_nodes(loop_body);
                let incoming_scope = self.scope.snapshot();
                let loop_index_id = self.alloc_value()?;
                child.ops.push(KernelOp {
                    kind: KernelOpKind::LoopIndex {
                        loop_var: var.shared_text(),
                    },
                    operands: Vec::new(),
                    result: Some(loop_index_id),
                });
                self.scope.bind(var.clone(), loop_index_id);

                // Loop-carry: any name that's reassigned inside the loop
                // body AND was bound in the enclosing scope needs a
                // per-iteration carrier SSA so in-loop reads see the
                // accumulated value (not the pre-loop literal). The
                // emitter routes each `LoopCarrier` op through a
                // function-local; the matching `LoopCarrierEnd` writes
                // the iter's final value back so the next iteration
                // reads it. Without this, `assign tok_idx = tok_idx + 1`
                // inside a loop body always re-evaluates from the
                // pre-loop seed (e.g. literal 0), so the lex kernel
                // returned n_tokens=0 regardless of input length.
                let assigned = collect_assigned_names(loop_body);
                let mut carriers: Vec<Ident> = Vec::new();
                for name in &assigned {
                    if name == var {
                        continue;
                    }
                    let Some(&seed_ssa) = incoming_scope.get(name) else {
                        continue;
                    };
                    let carrier_ssa = self.alloc_value()?;
                    child.ops.push(KernelOp {
                        kind: KernelOpKind::LoopCarrier {
                            name: name.shared_text(),
                        },
                        operands: vec![seed_ssa],
                        result: Some(carrier_ssa),
                    });
                    self.scope.bind(name.clone(), carrier_ssa);
                    carriers.push(name.clone());
                }

                self.lower_nodes(loop_body, &mut child, depth + 1)?;

                // Write each carrier's final per-iter value back so the
                // next iteration's `LoopCarrier` Load picks it up.
                for name in &carriers {
                    if let Some(final_ssa) = self.scope.get(name) {
                        child.ops.push(KernelOp {
                            kind: KernelOpKind::LoopCarrierEnd {
                                name: name.shared_text(),
                            },
                            operands: vec![final_ssa],
                            result: None,
                        });
                    }
                }

                let loop_exit_scope = self.scope.snapshot();
                self.scope
                    .restore_loop_exit(incoming_scope, &loop_exit_scope, var);
                let child_id = push_child(body, child)?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::StructuredForLoop {
                        loop_var: var.shared_text(),
                    },
                    operands: vec![from_id, to_id, child_id],
                    result: None,
                });
                // Bind each carrier name's post-loop SSA to a
                // `LoopCarrierFinal` op the emitter routes to a Load
                // from the same function-local that `LoopCarrier` /
                // `LoopCarrierEnd` shared. Without this, post-loop
                // reads of `name` would still point at the in-loop
                // SSA (the BinOp result inside the body), which is
                // out of scope at the parent body.
                for name in &carriers {
                    let final_ssa = self.alloc_value()?;
                    body.ops.push(KernelOp {
                        kind: KernelOpKind::LoopCarrierFinal {
                            name: name.shared_text(),
                        },
                        operands: vec![],
                        result: Some(final_ssa),
                    });
                    self.scope.bind(name.clone(), final_ssa);
                }
                Ok(())
            }
            Node::Barrier { ordering } => {
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Barrier {
                        ordering: *ordering,
                    },
                    operands: Vec::new(),
                    result: None,
                });
                Ok(())
            }
            Node::IndirectDispatch {
                count_buffer,
                count_offset,
            } => {
                let slot = self.buffer_slot(count_buffer)?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::IndirectDispatch {
                        count_offset: *count_offset,
                    },
                    operands: vec![slot],
                    result: None,
                });
                Ok(())
            }
            Node::AsyncLoad {
                source,
                destination,
                offset,
                size,
                tag,
            } => self.lower_async_copy(
                body,
                KernelOpKind::AsyncLoad {
                    tag: tag.shared_text(),
                },
                source,
                destination,
                offset,
                size,
            ),
            Node::AsyncStore {
                source,
                destination,
                offset,
                size,
                tag,
            } => self.lower_async_copy(
                body,
                KernelOpKind::AsyncStore {
                    tag: tag.shared_text(),
                },
                source,
                destination,
                offset,
                size,
            ),
            Node::AsyncWait { tag } => {
                body.ops.push(KernelOp {
                    kind: KernelOpKind::AsyncWait {
                        tag: tag.shared_text(),
                    },
                    operands: Vec::new(),
                    result: None,
                });
                Ok(())
            }
            Node::Trap { address, tag } => {
                let address_id = self.lower_expr(address, body)?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Trap {
                        tag: tag.shared_text(),
                    },
                    operands: vec![address_id],
                    result: None,
                });
                Ok(())
            }
            Node::Resume { tag } => {
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Resume {
                        tag: tag.shared_text(),
                    },
                    operands: Vec::new(),
                    result: None,
                });
                Ok(())
            }
            Node::Return => {
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Return,
                    operands: Vec::new(),
                    result: None,
                });
                Ok(())
            }
            Node::Opaque(extension) => {
                body.ops.push(KernelOp {
                    kind: KernelOpKind::OpaqueNode {
                        extension_kind: extension.extension_kind().to_owned(),
                        payload: extension.wire_payload(),
                    },
                    operands: Vec::new(),
                    result: None,
                });
                Ok(())
            }
            other => Err(LowerError::UnsupportedConstruct(format!(
                "node `{}` has no KernelDescriptor lowering. Fix: add a KernelOpKind mapping before routing this program through vyre-lower.",
                node_op_id(other)
            ))),
        }
    }

    fn lower_child_node(
        &mut self,
        body: &mut KernelBody,
        depth: usize,
        nodes: &[Node],
        kind: KernelOpKind,
    ) -> Result<(), LowerError> {
        let mut child = empty_body_for_nodes(nodes);
        self.lower_nodes(nodes, &mut child, depth + 1)?;
        let child_id = push_child(body, child)?;
        body.ops.push(KernelOp {
            kind,
            operands: vec![child_id],
            result: None,
        });
        Ok(())
    }

    fn lower_async_copy(
        &mut self,
        body: &mut KernelBody,
        kind: KernelOpKind,
        source: &Ident,
        destination: &Ident,
        offset: &Expr,
        size: &Expr,
    ) -> Result<(), LowerError> {
        let source_slot = self.buffer_slot(source)?;
        let destination_slot = self.buffer_slot(destination)?;
        let offset_id = self.lower_expr(offset, body)?;
        let size_id = self.lower_expr(size, body)?;
        body.ops.push(KernelOp {
            kind,
            operands: vec![source_slot, destination_slot, offset_id, size_id],
            result: None,
        });
        Ok(())
    }

    fn lower_expr(&mut self, expr: &Expr, body: &mut KernelBody) -> Result<u32, LowerError> {
        match expr {
            Expr::LitU32(value) => self.literal(body, LiteralValue::U32(*value)),
            Expr::LitI32(value) => self.literal(body, LiteralValue::I32(*value)),
            Expr::LitF32(value) => self.literal(body, LiteralValue::F32(*value)),
            Expr::LitBool(value) => self.literal(body, LiteralValue::Bool(*value)),
            Expr::Var(name) => self.scope.get(name).ok_or_else(|| {
                LowerError::InvalidProgram(format!(
                    "variable `{name}` is referenced before binding. Fix: emit a Let/Assign before use."
                ))
            }),
            Expr::Load { buffer, index } => {
                let slot = self.buffer_slot(buffer)?;
                let index_id = self.lower_expr(index, body)?;
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: self.load_kind(slot),
                    operands: vec![slot, index_id],
                    result: Some(result),
                });
                Ok(result)
            }
            Expr::BufLen { buffer } => {
                let slot = self.buffer_slot(buffer)?;
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::BufferLength,
                    operands: vec![slot],
                    result: Some(result),
                });
                Ok(result)
            }
            Expr::InvocationId { axis } => {
                self.builtin_axis(body, KernelOpKind::GlobalInvocationId, *axis)
            }
            Expr::WorkgroupId { axis } => {
                self.builtin_axis(body, KernelOpKind::WorkgroupId, *axis)
            }
            Expr::LocalId { axis } => {
                self.builtin_axis(body, KernelOpKind::LocalInvocationId, *axis)
            }
            Expr::BinOp { op, left, right } => {
                let left_id = self.lower_expr(left, body)?;
                let right_id = self.lower_expr(right, body)?;
                self.binary(body, KernelOpKind::BinOpKind(*op), left_id, right_id)
            }
            Expr::UnOp { op, operand } => {
                let operand_id = self.lower_expr(operand, body)?;
                self.unary(body, KernelOpKind::UnOpKind(op.clone()), operand_id)
            }
            Expr::Call { op_id, args } => {
                let mut operands = Vec::with_capacity(args.len());
                for arg in args {
                    operands.push(self.lower_expr(arg, body)?);
                }
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Call {
                        op_id: op_id.shared_text(),
                    },
                    operands,
                    result: Some(result),
                });
                Ok(result)
            }
            Expr::Select {
                cond,
                true_val,
                false_val,
            } => {
                let cond_id = self.lower_expr(cond, body)?;
                let true_id = self.lower_expr(true_val, body)?;
                let false_id = self.lower_expr(false_val, body)?;
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Select,
                    operands: vec![cond_id, true_id, false_id],
                    result: Some(result),
                });
                Ok(result)
            }
            Expr::Cast { target, value } => {
                let value_id = self.lower_expr(value, body)?;
                self.unary(
                    body,
                    KernelOpKind::Cast {
                        target: target.clone(),
                    },
                    value_id,
                )
            }
            Expr::Fma { a, b, c } => {
                let a_id = self.lower_expr(a, body)?;
                let b_id = self.lower_expr(b, body)?;
                let c_id = self.lower_expr(c, body)?;
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Fma,
                    operands: vec![a_id, b_id, c_id],
                    result: Some(result),
                });
                Ok(result)
            }
            Expr::Atomic {
                op,
                buffer,
                index,
                expected,
                value,
                ordering,
            } => {
                let slot = self.buffer_slot(buffer)?;
                let index_id = self.lower_expr(index, body)?;
                let value_id = self.lower_expr(value, body)?;
                let operands = if matches!(
                    op,
                    AtomicOp::CompareExchange | AtomicOp::CompareExchangeWeak
                ) {
                    let Some(expected) = expected else {
                        return Err(LowerError::InvalidProgram(
                            "atomic compare-exchange is missing expected value. Fix: set Expr::Atomic.expected.".into(),
                        ));
                    };
                    let expected_id = self.lower_expr(expected, body)?;
                    vec![slot, index_id, expected_id, value_id]
                } else {
                    vec![slot, index_id, value_id]
                };
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::Atomic {
                        op: *op,
                        ordering: *ordering,
                    },
                    operands,
                    result: Some(result),
                });
                Ok(result)
            }
            Expr::SubgroupBallot { cond } => {
                let cond_id = self.lower_expr(cond, body)?;
                self.unary(body, KernelOpKind::SubgroupBallot, cond_id)
            }
            Expr::SubgroupShuffle { value, lane } => {
                let value_id = self.lower_expr(value, body)?;
                let lane_id = self.lower_expr(lane, body)?;
                self.binary(body, KernelOpKind::SubgroupShuffle, value_id, lane_id)
            }
            Expr::SubgroupAdd { value } => {
                let value_id = self.lower_expr(value, body)?;
                self.unary(body, KernelOpKind::SubgroupAdd, value_id)
            }
            Expr::SubgroupLocalId => self.simple_result(body, KernelOpKind::SubgroupLocalId),
            Expr::SubgroupSize => self.simple_result(body, KernelOpKind::SubgroupSize),
            Expr::Opaque(extension) => {
                let result = self.alloc_value()?;
                body.ops.push(KernelOp {
                    kind: KernelOpKind::OpaqueExpr {
                        extension_id: opaque_extension_id(&**extension),
                        extension_kind: extension.extension_kind().to_owned(),
                        payload: extension.wire_payload(),
                    },
                    operands: Vec::new(),
                    result: Some(result),
                });
                Ok(result)
            }
            other => Err(LowerError::UnsupportedConstruct(format!(
                "expression `{other:?}` has no KernelDescriptor lowering. Fix: add a descriptor op mapping."
            ))),
        }
    }

    fn buffer_slot(&self, buffer: &Ident) -> Result<u32, LowerError> {
        self.buffer_slots
            .get(buffer)
            .copied()
            .ok_or_else(|| LowerError::UndeclaredBuffer(buffer.to_string()))
    }

    fn load_kind(&self, slot: u32) -> KernelOpKind {
        self.slot_memory_classes
            .get(&slot)
            .copied()
            .map(|memory_class| match memory_class {
                MemoryClass::Shared => KernelOpKind::LoadShared,
                MemoryClass::Constant => KernelOpKind::LoadConstant,
                MemoryClass::Global | MemoryClass::Scratch => KernelOpKind::LoadGlobal,
            })
            .unwrap_or(KernelOpKind::LoadGlobal)
    }

    fn store_kind(&self, slot: u32) -> KernelOpKind {
        self.slot_memory_classes
            .get(&slot)
            .copied()
            .map(|memory_class| match memory_class {
                MemoryClass::Shared => KernelOpKind::StoreShared,
                MemoryClass::Global | MemoryClass::Scratch | MemoryClass::Constant => {
                    KernelOpKind::StoreGlobal
                }
            })
            .unwrap_or(KernelOpKind::StoreGlobal)
    }

    fn add_trap_sidecar_binding(&mut self) -> Result<(), LowerError> {
        if self
            .buffer_slots
            .contains_key(&Ident::from(TRAP_SIDECAR_NAME))
        {
            return Err(LowerError::UnsupportedConstruct(format!(
                "program declares reserved trap sidecar buffer `{TRAP_SIDECAR_NAME}`. Fix: choose a user buffer name outside the `__vyre_*` namespace."
            )));
        }
        let next_slot = self
            .bindings
            .iter()
            .map(|binding| binding.slot)
            .max()
            .map_or(Ok(0), |slot| {
                slot.checked_add(1).ok_or(LowerError::OperandIdOverflow)
            })?;
        self.buffer_slots
            .insert(Ident::from(TRAP_SIDECAR_NAME), next_slot);
        self.slot_memory_classes
            .insert(next_slot, MemoryClass::Global);
        self.bindings.push(BindingSlot {
            slot: next_slot,
            element_type: DataType::U32,
            element_count: Some(TRAP_SIDECAR_WORDS),
            memory_class: MemoryClass::Global,
            visibility: BindingVisibility::ReadWrite,
            name: TRAP_SIDECAR_NAME.to_owned(),
        });
        self.bindings.sort_by_key(|slot| slot.slot);
        Ok(())
    }

    fn literal(&mut self, body: &mut KernelBody, literal: LiteralValue) -> Result<u32, LowerError> {
        let literal_index = push_literal(body, literal)?;
        let result = self.alloc_value()?;
        body.ops.push(KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![literal_index],
            result: Some(result),
        });
        Ok(result)
    }

    fn builtin_axis(
        &mut self,
        body: &mut KernelBody,
        kind: KernelOpKind,
        axis: u8,
    ) -> Result<u32, LowerError> {
        if axis > 2 {
            return Err(LowerError::InvalidProgram(format!(
                "builtin axis {axis} is out of range. Fix: use axis 0, 1, or 2."
            )));
        }
        let result = self.alloc_value()?;
        body.ops.push(KernelOp {
            kind,
            operands: vec![u32::from(axis)],
            result: Some(result),
        });
        Ok(result)
    }

    fn simple_result(
        &mut self,
        body: &mut KernelBody,
        kind: KernelOpKind,
    ) -> Result<u32, LowerError> {
        let result = self.alloc_value()?;
        body.ops.push(KernelOp {
            kind,
            operands: Vec::new(),
            result: Some(result),
        });
        Ok(result)
    }

    fn unary(
        &mut self,
        body: &mut KernelBody,
        kind: KernelOpKind,
        operand: u32,
    ) -> Result<u32, LowerError> {
        let result = self.alloc_value()?;
        body.ops.push(KernelOp {
            kind,
            operands: vec![operand],
            result: Some(result),
        });
        Ok(result)
    }

    fn binary(
        &mut self,
        body: &mut KernelBody,
        kind: KernelOpKind,
        left: u32,
        right: u32,
    ) -> Result<u32, LowerError> {
        let result = self.alloc_value()?;
        body.ops.push(KernelOp {
            kind,
            operands: vec![left, right],
            result: Some(result),
        });
        Ok(result)
    }

    fn alloc_value(&mut self) -> Result<u32, LowerError> {
        let id = self.next_value;
        self.next_value = self
            .next_value
            .checked_add(1)
            .ok_or(LowerError::OperandIdOverflow)?;
        Ok(id)
    }

    /// Merge variable bindings after a `Node::If` (no else): for each
    /// name in `incoming` whose ssa id was rebound in the then-branch,
    /// emit a `Select(cond, then_id, incoming_id)` in the parent body
    /// and bind name → new_id. Names that didn't change pass through.
    /// New bindings introduced inside the then-branch don't escape
    /// (they aren't in `incoming`).
    fn merge_if_then_scope(
        &mut self,
        body: &mut KernelBody,
        incoming: &scope::ScopeSnapshot,
        then_exit: &scope::ScopeSnapshot,
        cond_id: u32,
    ) -> Result<(), LowerError> {
        for (name, &incoming_id) in incoming.iter() {
            let then_id = then_exit.get(name).copied().unwrap_or(incoming_id);
            if then_id == incoming_id {
                continue;
            }
            let result = self.alloc_value()?;
            body.ops.push(KernelOp {
                kind: KernelOpKind::Select,
                operands: vec![cond_id, then_id, incoming_id],
                result: Some(result),
            });
            self.scope.bind(name.clone(), result);
        }
        Ok(())
    }

    /// Merge variable bindings after a `Node::If { otherwise: ..}`:
    /// for each name in `incoming` whose ssa id changed in either
    /// branch, emit a `Select(cond, then_id, else_id)` and bind name
    /// → new_id. Falls through when both branches kept the incoming id.
    fn merge_if_else_scope(
        &mut self,
        body: &mut KernelBody,
        incoming: &scope::ScopeSnapshot,
        then_exit: &scope::ScopeSnapshot,
        else_exit: &scope::ScopeSnapshot,
        cond_id: u32,
    ) -> Result<(), LowerError> {
        for (name, &incoming_id) in incoming.iter() {
            let then_id = then_exit.get(name).copied().unwrap_or(incoming_id);
            let else_id = else_exit.get(name).copied().unwrap_or(incoming_id);
            if then_id == incoming_id && else_id == incoming_id {
                continue;
            }
            if then_id == else_id {
                self.scope.bind(name.clone(), then_id);
                continue;
            }
            let result = self.alloc_value()?;
            body.ops.push(KernelOp {
                kind: KernelOpKind::Select,
                operands: vec![cond_id, then_id, else_id],
                result: Some(result),
            });
            self.scope.bind(name.clone(), result);
        }
        Ok(())
    }
}

fn body_contains_trap(body: &KernelBody) -> bool {
    body.ops
        .iter()
        .any(|op| matches!(op.kind, KernelOpKind::Trap { .. }))
        || body.child_bodies.iter().any(body_contains_trap)
}

fn opaque_extension_id(extension: &dyn vyre_foundation::ir::ExprNode) -> u32 {
    u32::from_le_bytes(
        extension.stable_fingerprint()[0..4]
            .try_into()
            .expect("slice length is fixed"),
    )
}

fn empty_body_for_nodes(nodes: &[Node]) -> KernelBody {
    empty_body_with_capacity(estimated_node_slice_op_capacity(nodes))
}

/// Walk a node slice and collect every `Node::Assign` target name
/// (recursively into nested If/Loop/Block/Region bodies). Used by
/// `Node::Loop` lowering to decide which incoming-scope names need a
/// loop-carrier SSA so their in-loop reads see the accumulated value.
fn collect_assigned_names(nodes: &[Node]) -> std::collections::BTreeSet<Ident> {
    let mut out = std::collections::BTreeSet::new();
    for n in nodes {
        walk_assigns(n, &mut out);
    }
    out
}

fn walk_assigns(node: &Node, out: &mut std::collections::BTreeSet<Ident>) {
    match node {
        Node::Assign { name, .. } => {
            out.insert(name.clone());
        }
        Node::If {
            then, otherwise, ..
        } => {
            for n in then {
                walk_assigns(n, out);
            }
            for n in otherwise {
                walk_assigns(n, out);
            }
        }
        Node::Loop { body, .. } => {
            for n in body {
                walk_assigns(n, out);
            }
        }
        Node::Block(body) => {
            for n in body {
                walk_assigns(n, out);
            }
        }
        Node::Region { body, .. } => {
            for n in body.iter() {
                walk_assigns(n, out);
            }
        }
        _ => {}
    }
}

fn empty_body_with_capacity(op_capacity: usize) -> KernelBody {
    KernelBody {
        ops: Vec::with_capacity(op_capacity),
        child_bodies: Vec::with_capacity(estimated_child_body_capacity(op_capacity)),
        literals: Vec::with_capacity(op_capacity / 3),
    }
}

fn estimated_root_op_capacity(program: &Program) -> usize {
    let stats = program.stats();
    stats
        .instruction_count
        .saturating_add(stats.node_count as u64)
        .saturating_add(4)
        .min(usize::MAX as u64) as usize
}

fn estimated_node_slice_op_capacity(nodes: &[Node]) -> usize {
    nodes
        .len()
        .saturating_mul(2)
        .saturating_add(estimated_child_body_capacity(nodes.len()))
}

fn estimated_child_body_capacity(parent_ops: usize) -> usize {
    parent_ops.min(16)
}

fn push_literal(body: &mut KernelBody, literal: LiteralValue) -> Result<u32, LowerError> {
    let index = u32::try_from(body.literals.len()).map_err(|_| LowerError::OperandIdOverflow)?;
    body.literals.push(literal);
    Ok(index)
}

fn push_child(body: &mut KernelBody, child: KernelBody) -> Result<u32, LowerError> {
    let index =
        u32::try_from(body.child_bodies.len()).map_err(|_| LowerError::OperandIdOverflow)?;
    body.child_bodies.push(child);
    Ok(index)
}

fn memory_class(buffer: &BufferDecl) -> Result<MemoryClass, LowerError> {
    match (buffer.kind, &buffer.access) {
        (MemoryKind::Persistent, _) => Err(LowerError::UnsupportedConstruct(format!(
            "Persistent memory buffer `{}` cannot be lowered as a direct GPU binding. Fix: stage Persistent data through the host transfer path using AsyncLoad/AsyncStore into Global/Readonly memory before concrete GPU emission.",
            buffer.name()
        ))),
        (MemoryKind::Shared, _) | (_, BufferAccess::Workgroup) => Ok(MemoryClass::Shared),
        (MemoryKind::Local, _) => Ok(MemoryClass::Scratch),
        (MemoryKind::Uniform | MemoryKind::Push | MemoryKind::Readonly, _)
        | (_, BufferAccess::Uniform | BufferAccess::ReadOnly) => Ok(MemoryClass::Constant),
        (MemoryKind::Global, _) => Ok(MemoryClass::Global),
        (other, _) => Err(LowerError::UnsupportedConstruct(format!(
            "MemoryKind::{other:?} for buffer `{}` is not supported by neutral lowering. Fix: map the buffer to Global, Shared, Uniform, Readonly, Push, or Local before emission.",
            buffer.name()
        ))),
    }
}

fn binding_visibility(access: &BufferAccess) -> BindingVisibility {
    match access {
        BufferAccess::ReadOnly | BufferAccess::Uniform => BindingVisibility::ReadOnly,
        BufferAccess::WriteOnly => BindingVisibility::WriteOnly,
        _ => BindingVisibility::ReadWrite,
    }
}

fn fingerprint_id(program: &Program) -> String {
    let fingerprint = program.fingerprint();
    let mut out = String::with_capacity(fingerprint.len() * 2);
    for byte in fingerprint {
        use std::fmt::Write;
        let _ = write!(&mut out, "{byte:02x}");
    }
    out
}

/// Public re-export of the binding-construction shape so emitters can
/// build descriptors in tests without going through `lower()`.
pub fn binding_slot(
    slot: u32,
    name: impl Into<String>,
    element_type: DataType,
    element_count: Option<u32>,
    memory_class: MemoryClass,
    visibility: BindingVisibility,
) -> BindingSlot {
    BindingSlot {
        slot,
        element_type,
        element_count,
        memory_class,
        visibility,
        name: name.into(),
    }
}

/// Public helper for a scalar-store op (used by descriptor-building
/// tests in this crate and by emitter integration tests).
pub fn store_global(
    slot_operand_id: u32,
    index_operand_id: u32,
    value_operand_id: u32,
) -> KernelOp {
    KernelOp {
        kind: KernelOpKind::StoreGlobal,
        operands: vec![slot_operand_id, index_operand_id, value_operand_id],
        result: None,
    }
}

/// Public helper for a u32 literal op.
pub fn literal_u32(literal_pool_index: u32, result_id: u32) -> KernelOp {
    KernelOp {
        kind: KernelOpKind::Literal,
        operands: vec![literal_pool_index],
        result: Some(result_id),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn lower_empty_wrapped_program_preserves_region() {
        let program = vyre_foundation::ir::Program::wrapped(vec![], [1, 1, 1], vec![]);
        let desc = lower(&program).unwrap();
        assert_eq!(desc.dispatch.workgroup_size, [1, 1, 1]);
        assert!(desc.bindings.slots.is_empty());
        assert_eq!(desc.body.ops.len(), 1);
        assert!(matches!(desc.body.ops[0].kind, KernelOpKind::Region { .. }));
    }

    #[test]
    fn binding_slot_helper_records_inputs() {
        let s = binding_slot(
            3,
            "scratch",
            DataType::F32,
            Some(64),
            MemoryClass::Shared,
            BindingVisibility::ReadWrite,
        );
        assert_eq!(s.slot, 3);
        assert_eq!(s.name, "scratch");
        assert_eq!(s.element_type, DataType::F32);
        assert_eq!(s.element_count, Some(64));
        assert_eq!(s.memory_class, MemoryClass::Shared);
        assert_eq!(s.visibility, BindingVisibility::ReadWrite);
    }

    #[test]
    fn store_global_helper_packs_three_operands() {
        let op = store_global(0, 1, 2);
        assert_eq!(op.kind, KernelOpKind::StoreGlobal);
        assert_eq!(op.operands, vec![0, 1, 2]);
        assert_eq!(op.result, None);
    }

    #[test]
    fn literal_u32_helper_assigns_result_id() {
        let op = literal_u32(5, 42);
        assert_eq!(op.kind, KernelOpKind::Literal);
        assert_eq!(op.operands, vec![5]);
        assert_eq!(op.result, Some(42));
    }

    #[test]
    fn lower_assigns_unique_descriptor_slots_for_duplicate_program_bindings() {
        use vyre_foundation::ir::{BufferDecl, Expr, Node};

        let program = Program::wrapped(
            vec![
                BufferDecl::workgroup("scratch", 16, DataType::U32),
                BufferDecl::output("out", 0, DataType::U32).with_count(1),
            ],
            [64, 1, 1],
            vec![Node::store("out", Expr::u32(0), Expr::u32(1))],
        );

        let desc = lower(&program).expect("duplicate Program bindings must descriptor-lower");

        assert_eq!(desc.bindings.slots.len(), 2);
        assert_ne!(desc.bindings.slots[0].slot, desc.bindings.slots[1].slot);
        assert!(crate::verify::verify(&desc).is_ok());
    }

    #[test]
    fn lower_trap_inserts_descriptor_sidecar_binding() {
        use vyre_foundation::ir::{Expr, Node};

        let program = Program::wrapped(
            vec![],
            [64, 1, 1],
            vec![Node::trap(Expr::u32(7), "page-fault")],
        );

        let desc = lower(&program).expect("trap programs must descriptor-lower");
        let sidecar = desc
            .bindings
            .slots
            .iter()
            .find(|slot| slot.name == TRAP_SIDECAR_NAME)
            .expect("trap sidecar binding must be inserted");
        assert_eq!(sidecar.element_type, DataType::U32);
        assert_eq!(sidecar.element_count, Some(TRAP_SIDECAR_WORDS));
        assert!(matches!(sidecar.visibility, BindingVisibility::ReadWrite));
        assert!(crate::verify::verify(&desc).is_ok());
    }

    #[test]
    fn lower_opaque_expr_preserves_kind_and_payload() {
        use vyre_foundation::ir::{BufferDecl, Expr, Node};

        let program = Program::wrapped(
            vec![BufferDecl::output("out", 0, DataType::U32).with_count(1)],
            [1, 1, 1],
            vec![Node::store("out", Expr::u32(0), Expr::u64(42))],
        );

        let desc = lower(&program).expect("opaque literals must descriptor-lower");
        fn find_opaque_expr(body: &KernelBody) -> Option<(&String, &Vec<u8>)> {
            body.ops
                .iter()
                .find_map(|op| match &op.kind {
                    KernelOpKind::OpaqueExpr {
                        extension_kind,
                        payload,
                        ..
                    } => Some((extension_kind, payload)),
                    _ => None,
                })
                .or_else(|| body.child_bodies.iter().find_map(find_opaque_expr))
        }

        let opaque = find_opaque_expr(&desc.body).expect("opaque expression op must be present");
        assert_eq!(opaque.0, "vyre.literal.u64");
        assert_eq!(opaque.1, &42u64.to_le_bytes().to_vec());
    }

    #[test]
    fn loop_variable_lowers_to_child_loop_index_result() {
        use vyre_foundation::ir::{BufferDecl, Expr, Node};

        let program = Program::wrapped(
            vec![
                BufferDecl::storage("input", 0, BufferAccess::ReadOnly, DataType::U32),
                BufferDecl::output("out", 1, DataType::U32).with_count(1),
            ],
            [1, 1, 1],
            vec![Node::loop_for(
                "i",
                Expr::u32(0),
                Expr::buf_len("input"),
                vec![Node::store(
                    "out",
                    Expr::u32(0),
                    Expr::load("input", Expr::var("i")),
                )],
            )],
        );

        let desc = lower(&program).expect("loop variable must descriptor-lower");
        assert!(crate::verify::verify(&desc).is_ok());
        let (loop_body, loop_op) =
            find_loop(&desc.body).expect("structured loop op must be present");
        let child = &loop_body.child_bodies[loop_op.operands[2] as usize];
        assert!(
            matches!(
                child.ops.first().map(|op| &op.kind),
                Some(KernelOpKind::LoopIndex { loop_var }) if loop_var.as_ref() == "i"
            ),
            "loop body must materialize the induction value before lowering input[i]"
        );

        fn find_loop(body: &KernelBody) -> Option<(&KernelBody, &KernelOp)> {
            for op in &body.ops {
                if matches!(op.kind, KernelOpKind::StructuredForLoop { .. }) {
                    return Some((body, op));
                }
            }
            body.child_bodies.iter().find_map(find_loop)
        }
    }

    #[test]
    fn loop_variable_does_not_clobber_same_named_outer_binding() {
        use vyre_foundation::ir::{BufferDecl, Expr, Node};

        let program = Program::wrapped(
            vec![BufferDecl::output("out", 0, DataType::U32).with_count(1)],
            [1, 1, 1],
            vec![
                Node::let_bind("i", Expr::u32(9)),
                Node::loop_for("i", Expr::u32(0), Expr::u32(1), vec![]),
                Node::store("out", Expr::u32(0), Expr::var("i")),
            ],
        );

        let desc = lower(&program).expect("shadowed loop variable must descriptor-lower");
        assert!(crate::verify::verify(&desc).is_ok());
        let store = find_store(&desc.body).expect("post-loop store must be present");
        assert_eq!(
            store.operands[2], 0,
            "post-loop read must use the outer i binding, not the loop induction result"
        );

        fn find_store(body: &KernelBody) -> Option<&KernelOp> {
            body.ops
                .iter()
                .find(|op| matches!(op.kind, KernelOpKind::StoreGlobal))
                .or_else(|| body.child_bodies.iter().find_map(find_store))
        }
    }

    #[test]
    fn if_else_branches_lower_from_the_same_incoming_scope() {
        use vyre_foundation::ir::{BufferDecl, Expr, Node};

        let program = Program::wrapped(
            vec![BufferDecl::output("out", 0, DataType::U32).with_count(1)],
            [1, 1, 1],
            vec![
                Node::let_bind("x", Expr::u32(1)),
                Node::if_then_else(
                    Expr::bool(true),
                    vec![Node::assign("x", Expr::add(Expr::var("x"), Expr::u32(1)))],
                    vec![Node::store("out", Expr::u32(0), Expr::var("x"))],
                ),
            ],
        );

        let desc = lower(&program).expect("if/else must descriptor-lower");
        assert!(crate::verify::verify(&desc).is_ok());
        let (_, if_op) = find_if_else(&desc.body).expect("if/else op must be present");
        let parent = find_parent_body_containing_op(&desc.body, if_op as *const KernelOp)
            .expect("if op parent body must be found");
        let else_body = &parent.child_bodies[if_op.operands[2] as usize];
        let else_store = else_body
            .ops
            .iter()
            .find(|op| matches!(op.kind, KernelOpKind::StoreGlobal))
            .expect("else branch must contain the store");
        assert_eq!(
            else_store.operands[2], 0,
            "else branch must read the incoming x binding, not the result assigned only by then"
        );

        fn find_if_else(body: &KernelBody) -> Option<(&KernelBody, &KernelOp)> {
            for op in &body.ops {
                if matches!(op.kind, KernelOpKind::StructuredIfThenElse) {
                    return Some((body, op));
                }
            }
            body.child_bodies.iter().find_map(find_if_else)
        }

        fn find_parent_body_containing_op(
            body: &KernelBody,
            target: *const KernelOp,
        ) -> Option<&KernelBody> {
            if body.ops.iter().any(|op| std::ptr::eq(op, target)) {
                return Some(body);
            }
            body.child_bodies
                .iter()
                .find_map(|child| find_parent_body_containing_op(child, target))
        }
    }

    #[test]
    fn max_nesting_depth_constant_is_documented() {
        assert_eq!(MAX_NESTING_DEPTH, 64);
    }
}