wasm_pvm/llvm_backend/

emitter.rs

// Core PVM emitter and value management.
//
// Each SSA value gets a dedicated stack slot (correctness-first, no register allocation).
// Pattern: load operands from slots → temp regs, compute, store result to slot.

// We use 'as' casts extensively for:
// - PVM register indices (u8) from iterators
// - Address offsets (i32) from pointers
// - Immediate values (i32/i64) from LLVM constants
// These are intentional truncations/wraps where we know the range is safe or valid for PVM.
#![allow(
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::cast_sign_loss
)]

use std::collections::{HashMap, HashSet};

use inkwell::IntPredicate;
use inkwell::basic_block::BasicBlock;
use inkwell::values::{
    AnyValue, AsValueRef, BasicValueEnum, FunctionValue, InstructionValue, IntValue, PhiValue,
};

use crate::pvm::Instruction;
use crate::{Error, Result};

use crate::abi::{FRAME_HEADER_SIZE, STACK_PTR_REG};
use crate::translate::OptimizationFlags;

use super::regalloc::RegAllocResult;

/// Context for lowering functions from a single WASM module.
pub struct LoweringContext {
    pub wasm_memory_base: i32,
    pub num_globals: usize,
    /// Base address for the parameter overflow area (5th+ args for `call_indirect`).
    pub param_overflow_base: i32,
    pub function_signatures: Vec<(usize, bool)>,
    pub type_signatures: Vec<(usize, usize)>,
    pub function_table: Vec<u32>,
    pub num_imported_funcs: usize,
    pub imported_func_names: Vec<String>,
    pub initial_memory_pages: u32,
    pub max_memory_pages: u32,
    pub stack_size: u32,
    /// Map from data segment index to offset in `RO_DATA` (for passive segments).
    pub data_segment_offsets: HashMap<u32, u32>,
    /// Map from data segment index to byte length (for passive segments bounds checking).
    pub data_segment_lengths: HashMap<u32, u32>,
    /// Map from data segment index to PVM address storing the effective length at runtime.
    /// Used by `memory.init` for bounds checking and by `data.drop` to zero the length.
    pub data_segment_length_addrs: HashMap<u32, i32>,
    /// User-provided mapping from WASM import names to actions (trap, nop).
    pub wasm_import_map: Option<HashMap<String, crate::translate::ImportAction>>,
    /// Optimization flags controlling which compiler passes are enabled.
    pub optimizations: OptimizationFlags,
}

/// Result of lowering one LLVM function to PVM instructions.
pub struct LlvmFunctionTranslation {
    pub instructions: Vec<Instruction>,
    pub call_fixups: Vec<LlvmCallFixup>,
    pub indirect_call_fixups: Vec<LlvmIndirectCallFixup>,
    /// Number of call return addresses allocated by this function (for jump table indexing).
    pub num_call_returns: usize,
    /// Stats collected during lowering (for compiler output).
    pub lowering_stats: FunctionLoweringStats,
}

/// Stats collected during function lowering for the compile stats output.
#[derive(Debug, Clone, Default)]
pub struct FunctionLoweringStats {
    pub frame_size: i32,
    pub is_leaf: bool,
    /// Instruction count before dead store elimination.
    pub pre_dse_instructions: usize,
    /// Instruction count after DSE but before peephole.
    pub pre_peephole_instructions: usize,
    // ── Regalloc stats ──
    pub regalloc_total_values: usize,
    pub regalloc_allocated_values: usize,
    /// Physical register numbers allocated (e.g. [9, 10]).
    pub regalloc_registers_used: Vec<u8>,
    pub regalloc_skipped_reason: Option<&'static str>,
    // ── Regalloc usage stats ──
    pub regalloc_load_hits: usize,
    pub regalloc_load_reloads: usize,
    pub regalloc_load_moves: usize,
    pub regalloc_store_hits: usize,
    pub regalloc_store_moves: usize,
}

#[derive(Debug, Clone)]
pub struct LlvmCallFixup {
    pub return_addr_instr: usize,
    pub jump_instr: usize,
    pub target_func: u32,
}

#[derive(Debug, Clone)]
pub struct LlvmIndirectCallFixup {
    pub return_addr_instr: usize,
    // Instruction index whose end offset is used as the call return PC.
    // For `LoadImmJumpInd`, this is the same index as `return_addr_instr`.
    pub jump_ind_instr: usize,
}

/// Per-function configuration for the PVM emitter.
///
/// These fields are set once when lowering begins and are never mutated during codegen.
/// Separating them from mutable state makes it clear what's fixed vs what changes.
#[allow(clippy::struct_excessive_bools)]
pub struct EmitterConfig {
    /// WASM memory base address (for converting WASM addresses to PVM addresses).
    pub wasm_memory_base: i32,

    /// Base address for the parameter overflow area (5th+ args for `call_indirect`).
    pub param_overflow_base: i32,

    /// Whether the register cache (store-load forwarding) is enabled.
    pub register_cache_enabled: bool,

    /// Whether constant propagation (redundant `LoadImm` elimination) is enabled.
    pub constant_propagation_enabled: bool,

    /// Whether ICmp+Branch fusion is enabled.
    pub icmp_fusion_enabled: bool,

    /// Whether callee-save shrink wrapping is enabled.
    pub shrink_wrap_enabled: bool,

    /// Whether cross-block register cache propagation is enabled.
    pub cross_block_cache_enabled: bool,

    /// Whether register allocation (long-lived values in allocatable regs) is enabled.
    pub register_allocation_enabled: bool,

    /// Whether fallthrough jump elimination is enabled.
    pub fallthrough_jumps_enabled: bool,

    /// Whether lazy spill is enabled (skip stack stores for register-allocated values).
    pub lazy_spill_enabled: bool,
}

/// PVM code emitter for a single function.
pub struct PvmEmitter<'ctx> {
    /// Per-function configuration (immutable after construction).
    pub(crate) config: EmitterConfig,

    pub(crate) instructions: Vec<Instruction>,
    pub(crate) labels: Vec<Option<usize>>,
    pub(crate) fixups: Vec<(usize, usize)>,

    /// LLVM basic block → PVM label.
    pub(crate) block_labels: HashMap<BasicBlock<'ctx>, usize>,

    /// LLVM int values (params + instruction results) → stack slot offset from SP.
    pub(crate) value_slots: HashMap<ValKey, i32>,

    /// Next available slot offset (bump allocator, starts after frame header).
    pub(crate) next_slot_offset: i32,

    /// Total frame size (set after pre-scan).
    pub(crate) frame_size: i32,

    pub(crate) call_fixups: Vec<LlvmCallFixup>,
    pub(crate) indirect_call_fixups: Vec<LlvmIndirectCallFixup>,

    /// Current byte offset of the emitted code (for O(1) offset calculations).
    pub(crate) byte_offset: usize,

    /// Maps stack slot offset → register that currently holds this slot's value.
    slot_cache: HashMap<i32, u8>,
    /// Reverse: register → slot offset it holds (for fast invalidation).
    reg_to_slot: [Option<i32>; 13],

    /// Pending fused `ICmp`: when an `ICmp` has a single use (by a branch), we defer
    /// it and store the predicate + operands here. The branch will emit a fused
    /// branch instruction instead of loading the boolean result.
    pub(crate) pending_fused_icmp: Option<FusedIcmp<'ctx>>,

    /// Maps register → known constant value currently held (for constant propagation).
    /// When a `LoadImm`/`LoadImm64` is about to be emitted, we check if the target
    /// register already holds the same constant and skip the load if so.
    reg_to_const: [Option<u64>; 13],

    /// Next jump table index for call return addresses.
    /// Incremented each time a call is emitted, starting from `call_return_base_idx`.
    pub(crate) next_call_return_idx: usize,

    /// Base index for call return addresses (set from the global counter).
    call_return_base_idx: usize,

    /// For each block with exactly one predecessor, maps block → its single predecessor.
    /// Used for cross-block register cache propagation.
    pub(crate) block_single_pred: HashMap<BasicBlock<'ctx>, BasicBlock<'ctx>>,

    /// Which callee-saved registers (r9-r12) are actually used by this function.
    /// Index 0 = r9, 1 = r10, 2 = r11, 3 = r12.
    pub(crate) used_callee_regs: [bool; 4],

    /// Frame offset for each callee-saved register (r9-r12), if saved.
    /// Index 0 = r9, 1 = r10, 2 = r11, 3 = r12.
    pub(crate) callee_save_offsets: [Option<i32>; 4],

    /// Whether the function contains any call instructions.
    /// If false (leaf function), we can skip saving/restoring the return address register.
    pub(crate) has_calls: bool,

    /// Register allocation results (empty if disabled).
    pub(crate) regalloc: RegAllocResult,

    /// Which stack slot is currently materialized in each allocated register.
    /// Multiple slots may map to the same physical register across disjoint
    /// intervals; this tracks runtime ownership so stale values are reloaded.
    alloc_reg_slot: [Option<i32>; 13],

    /// Dirty bit per register: true when the register holds a value not yet
    /// written to its stack slot (lazy spill). Only meaningful when
    /// `config.lazy_spill_enabled` is true.
    alloc_dirty: [bool; 13],

    /// Runtime usage counters for register-allocated mappings.
    pub(crate) regalloc_usage: RegAllocUsageStats,

    /// Label of the next basic block in layout order.
    /// When set, `emit_jump_to_label()` skips the Jump if the target matches,
    /// letting execution fall through to the next block naturally.
    pub next_block_label: Option<usize>,
}

/// Snapshot of the register cache state for cross-block propagation.
#[derive(Clone)]
pub struct CacheSnapshot {
    pub slot_cache: HashMap<i32, u8>,
    pub reg_to_slot: [Option<i32>; 13],
    pub reg_to_const: [Option<u64>; 13],
    pub alloc_reg_slot: [Option<i32>; 13],
    pub alloc_dirty: [bool; 13],
}

/// Instrumentation counters describing how much allocated mappings are used by codegen.
#[derive(Debug, Clone, Default)]
pub struct RegAllocUsageStats {
    /// `load_operand` looked up a value in `val_to_reg`.
    pub load_hits: usize,
    /// The allocated register had been invalidated and had to be reloaded from stack.
    pub load_reloads: usize,
    /// `load_operand` emitted `MoveReg` from allocated register to target temp register.
    pub load_moves: usize,
    /// `store_to_slot` wrote a value into an allocated destination register.
    pub store_hits: usize,
    /// `store_to_slot` needed a `MoveReg` to copy source into the allocated register.
    pub store_moves: usize,
}

impl CacheSnapshot {
    /// Invalidate a register's cache entries in this snapshot.
    /// Used to remove entries for registers that a terminator may clobber.
    pub fn invalidate_reg(&mut self, reg: u8) {
        let idx = reg as usize;
        if let Some(slot) = self.reg_to_slot[idx].take() {
            self.slot_cache.remove(&slot);
        }
        self.reg_to_const[idx] = None;
        self.invalidate_alloc_reg(reg);
    }

    /// Invalidate allocated-register slot ownership in this snapshot.
    pub fn invalidate_alloc_reg(&mut self, reg: u8) {
        let idx = reg as usize;
        self.alloc_reg_slot[idx] = None;
        self.alloc_dirty[idx] = false;
    }
}

/// Deferred `ICmp` info for branch fusion.
pub struct FusedIcmp<'ctx> {
    pub predicate: IntPredicate,
    pub lhs: BasicValueEnum<'ctx>,
    pub rhs: BasicValueEnum<'ctx>,
}

/// Wrapper key for LLVM values in the slot map.
/// Uses the raw LLVM value pointer cast to usize for hashing.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ValKey(pub(crate) usize);

pub fn val_key_int(val: IntValue<'_>) -> ValKey {
    ValKey(val.as_value_ref() as usize)
}

pub fn val_key_basic(val: BasicValueEnum<'_>) -> ValKey {
    ValKey(val.as_value_ref() as usize)
}

pub fn val_key_instr(val: InstructionValue<'_>) -> ValKey {
    ValKey(val.as_value_ref() as usize)
}

impl<'ctx> PvmEmitter<'ctx> {
    pub fn new(config: EmitterConfig, call_return_base: usize) -> Self {
        Self {
            config,
            instructions: Vec::new(),
            labels: Vec::new(),
            fixups: Vec::new(),
            block_labels: HashMap::new(),
            value_slots: HashMap::new(),
            next_slot_offset: FRAME_HEADER_SIZE,
            frame_size: 0,
            call_fixups: Vec::new(),
            indirect_call_fixups: Vec::new(),
            byte_offset: 0,
            slot_cache: HashMap::new(),
            reg_to_slot: [None; 13],
            reg_to_const: [None; 13],
            pending_fused_icmp: None,
            block_single_pred: HashMap::new(),
            next_call_return_idx: call_return_base,
            call_return_base_idx: call_return_base,
            used_callee_regs: [true; 4],
            callee_save_offsets: [Some(8), Some(16), Some(24), Some(32)],
            has_calls: true, // conservative default
            regalloc: RegAllocResult::default(),
            alloc_reg_slot: [None; 13],
            alloc_dirty: [false; 13],
            regalloc_usage: RegAllocUsageStats::default(),
            next_block_label: None,
        }
    }

    /// Allocate a call return address (jump table address) for a direct call site.
    /// Returns the pre-computed jump table address `(index + 1) * 2`.
    pub fn alloc_call_return_addr(&mut self) -> i32 {
        let idx = self.next_call_return_idx;
        self.next_call_return_idx += 1;
        (idx + 1)
            .checked_mul(2)
            .and_then(|v| i32::try_from(v).ok())
            .expect("alloc_call_return_addr: jump table index overflow (exceeds i32 range)")
    }

    /// Returns how many call return addresses this function allocated.
    pub fn num_call_returns(&self) -> usize {
        self.next_call_return_idx - self.call_return_base_idx
    }

    pub fn alloc_label(&mut self) -> usize {
        let id = self.labels.len();
        self.labels.push(None);
        id
    }

    pub fn define_label(&mut self, label: usize) {
        if self
            .instructions
            .last()
            .is_some_and(|last| !last.is_terminating())
        {
            self.emit(Instruction::Fallthrough);
        }
        self.labels[label] = Some(self.current_offset());
        self.clear_reg_cache();
    }

    /// Define a label, clearing the general cache but preserving allocated register state.
    ///
    /// Used for leaf functions where allocated registers are never clobbered.
    pub fn define_label_preserving_alloc(&mut self, label: usize) {
        if self
            .instructions
            .last()
            .is_some_and(|last| !last.is_terminating())
        {
            self.emit(Instruction::Fallthrough);
        }
        self.labels[label] = Some(self.current_offset());
        self.clear_reg_cache_preserving_alloc();
    }

    pub fn current_offset(&self) -> usize {
        self.byte_offset
    }

    pub fn emit(&mut self, instr: Instruction) {
        // Constant propagation: skip LoadImm/LoadImm64 if register already holds the value.
        if self.config.constant_propagation_enabled {
            match &instr {
                Instruction::LoadImm { reg, value } => {
                    // LoadImm sign-extends i32 → i64, so the 64-bit value is the sign extension.
                    let val64 = i64::from(*value) as u64;
                    if self.reg_to_const[*reg as usize] == Some(val64) {
                        return; // Already holds this constant.
                    }
                }
                Instruction::LoadImm64 { reg, value } => {
                    if self.reg_to_const[*reg as usize] == Some(*value) {
                        return; // Already holds this constant.
                    }
                }
                _ => {}
            }
        }

        if let Some(reg) = instr.dest_reg() {
            self.invalidate_reg(reg);
        }

        // Track constants after emit.
        if self.config.constant_propagation_enabled {
            match &instr {
                Instruction::LoadImm { reg, value } => {
                    self.reg_to_const[*reg as usize] = Some(i64::from(*value) as u64);
                }
                Instruction::LoadImm64 { reg, value } => {
                    self.reg_to_const[*reg as usize] = Some(*value);
                }
                _ => {}
            }
        }

        self.byte_offset += instr.encode().len();
        self.instructions.push(instr);
    }

    pub fn emit_jump_to_label(&mut self, label: usize) {
        // If jumping to the next block in layout order, skip the Jump — execution
        // will fall through naturally. define_label() emits a Fallthrough if needed.
        if self.config.fallthrough_jumps_enabled && self.next_block_label == Some(label) {
            return;
        }
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::Jump { offset: 0 });
    }

    pub fn emit_branch_ne_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchNeImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_eq_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchEqImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_lt_u_to_label(&mut self, reg1: u8, reg2: u8, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchLtU {
            reg1,
            reg2,
            offset: 0,
        });
    }

    pub fn emit_branch_ge_u_to_label(&mut self, reg1: u8, reg2: u8, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchGeU {
            reg1,
            reg2,
            offset: 0,
        });
    }

    pub fn emit_branch_eq_to_label(&mut self, reg1: u8, reg2: u8, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchEq {
            reg1,
            reg2,
            offset: 0,
        });
    }

    pub fn emit_branch_ne_to_label(&mut self, reg1: u8, reg2: u8, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchNe {
            reg1,
            reg2,
            offset: 0,
        });
    }

    pub fn emit_branch_lt_s_to_label(&mut self, reg1: u8, reg2: u8, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchLtS {
            reg1,
            reg2,
            offset: 0,
        });
    }

    pub fn emit_branch_ge_s_to_label(&mut self, reg1: u8, reg2: u8, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchGeS {
            reg1,
            reg2,
            offset: 0,
        });
    }

    pub fn emit_branch_lt_u_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchLtUImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_le_u_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchLeUImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_ge_u_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchGeUImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_gt_u_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchGtUImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_lt_s_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchLtSImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_le_s_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchLeSImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_ge_s_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchGeSImm {
            reg,
            value,
            offset: 0,
        });
    }

    pub fn emit_branch_gt_s_imm_to_label(&mut self, reg: u8, value: i32, label: usize) {
        let fixup_idx = self.instructions.len();
        self.fixups.push((fixup_idx, label));
        self.emit(Instruction::BranchGtSImm {
            reg,
            value,
            offset: 0,
        });
    }

    // ── Slot allocation ──

    pub fn alloc_slot_for_key(&mut self, key: ValKey) -> i32 {
        let offset = self.next_slot_offset;
        self.value_slots.insert(key, offset);
        self.next_slot_offset += 8;
        offset
    }

    pub fn get_slot(&self, key: ValKey) -> Option<i32> {
        self.value_slots.get(&key).copied()
    }

    // ── Value load / store ──

    /// Emit a constant load into `reg` (prefer compact `LoadImm` when the value fits i32).
    /// Constant propagation in `emit()` will skip if the register already holds this value.
    fn emit_const_to_reg(&mut self, reg: u8, value: u64) {
        let signed = value as i64;
        if let Ok(v32) = i32::try_from(signed) {
            self.emit(Instruction::LoadImm { reg, value: v32 });
        } else {
            self.emit(Instruction::LoadImm64 { reg, value });
        }
    }

    /// Load a value into a temp register. Constants are inlined; SSA values are loaded from slots.
    /// Poison/undef values are materialized as 0 (any value is valid for undefined behavior).
    /// Returns an error for truly unknown values.
    pub fn load_operand(&mut self, val: BasicValueEnum<'ctx>, temp_reg: u8) -> Result<()> {
        match val {
            BasicValueEnum::IntValue(iv) => {
                let key = val_key_int(iv);
                if let Some(signed_val) = iv.get_sign_extended_constant() {
                    self.emit_const_to_reg(temp_reg, signed_val as u64);
                } else if let Some(const_val) = iv.get_zero_extended_constant() {
                    // Fallback for unsigned constants.
                    self.emit_const_to_reg(temp_reg, const_val);
                } else if let Some(&alloc_reg) = self.regalloc.val_to_reg.get(&key) {
                    self.regalloc_usage.load_hits += 1;
                    let slot = self.get_slot(key).ok_or_else(|| {
                        Error::Internal(format!("no slot for allocated int value {key:?}"))
                    })?;
                    // Value has an allocated register. If it was clobbered OR this
                    // register currently materializes a different slot, reload
                    // lazily from the canonical stack slot.
                    let reg_idx = alloc_reg as usize;
                    if self.alloc_reg_slot[reg_idx] == Some(slot) {
                        // Fast path: register holds the correct value.
                        if alloc_reg != temp_reg {
                            self.regalloc_usage.load_moves += 1;
                            self.emit(Instruction::MoveReg {
                                dst: temp_reg,
                                src: alloc_reg,
                            });
                        }
                    } else if alloc_reg == temp_reg {
                        // Reload directly into the target (which is the alloc reg).
                        self.regalloc_usage.load_reloads += 1;
                        self.emit(Instruction::LoadIndU64 {
                            dst: alloc_reg,
                            base: STACK_PTR_REG,
                            offset: slot,
                        });
                        self.alloc_reg_slot[reg_idx] = Some(slot);
                        self.cache_slot(slot, alloc_reg);
                    } else {
                        // Allocated register was invalidated and target is different.
                        // Load directly into temp_reg to avoid clobbering alloc_reg
                        // (which may hold a call argument or other transient value).
                        self.regalloc_usage.load_reloads += 1;
                        self.emit(Instruction::LoadIndU64 {
                            dst: temp_reg,
                            base: STACK_PTR_REG,
                            offset: slot,
                        });
                        self.cache_slot(slot, temp_reg);
                    }
                } else if let Some(slot) = self.get_slot(key) {
                    // Check register cache: skip load if value is already in a register.
                    if let Some(&cached_reg) = self.slot_cache.get(&slot) {
                        if cached_reg != temp_reg {
                            // Emit a register copy (cheaper than memory load).
                            self.emit(Instruction::MoveReg {
                                dst: temp_reg,
                                src: cached_reg,
                            });
                        }
                        // If cached_reg == temp_reg, skip entirely (0 instructions).
                    } else {
                        self.emit(Instruction::LoadIndU64 {
                            dst: temp_reg,
                            base: STACK_PTR_REG,
                            offset: slot,
                        });
                        self.cache_slot(slot, temp_reg);
                    }
                } else if iv.is_poison() || iv.is_undef() {
                    // Poison/undef values can be materialized as any value; use 0.
                    self.emit(Instruction::LoadImm {
                        reg: temp_reg,
                        value: 0,
                    });
                } else {
                    return Err(Error::Internal(format!(
                        "no slot for int value {:?} (opcode: {:?})",
                        val_key_int(iv),
                        iv.as_instruction().map(InstructionValue::get_opcode),
                    )));
                }
            }
            _ => {
                return Err(Error::Internal(format!(
                    "cannot load non-integer value type {:?}",
                    val.get_type()
                )));
            }
        }
        Ok(())
    }

    pub fn store_to_slot(&mut self, slot_offset: i32, src_reg: u8) {
        let alloc_reg = self.regalloc.slot_to_reg.get(&slot_offset).copied();
        if let Some(alloc_reg) = alloc_reg {
            // If the register currently holds a DIFFERENT slot's dirty value,
            // spill it first. This happens when multiple values share a register
            // via early interval expiration (loop phi register reuse).
            let reg_idx = alloc_reg as usize;
            if let Some(current_slot) = self.alloc_reg_slot[reg_idx]
                && current_slot != slot_offset
                && self.alloc_dirty[reg_idx]
            {
                self.spill_dirty_reg(alloc_reg);
            }

            self.regalloc_usage.store_hits += 1;
            if src_reg != alloc_reg {
                self.regalloc_usage.store_moves += 1;
                self.emit(Instruction::MoveReg {
                    dst: alloc_reg,
                    src: src_reg,
                });
            }
            self.alloc_reg_slot[reg_idx] = Some(slot_offset);

            if self.config.lazy_spill_enabled {
                // Lazy spill: mark dirty, skip the stack store.
                self.alloc_dirty[reg_idx] = true;
                self.cache_slot(slot_offset, alloc_reg);
                return;
            }
        }
        self.emit(Instruction::StoreIndU64 {
            base: STACK_PTR_REG,
            src: src_reg,
            offset: slot_offset,
        });
        self.cache_slot(slot_offset, alloc_reg.unwrap_or(src_reg));
    }

    // ── Register allocation spill/reload ──

    /// Flush a single dirty allocated register to its stack slot.
    ///
    /// Emits `StoreIndU64` directly (bypassing `emit()` to avoid `invalidate_reg`
    /// clearing `alloc_reg_slot`), then marks the register clean.
    fn spill_dirty_reg(&mut self, reg: u8) {
        let reg_idx = reg as usize;
        if !self.alloc_dirty[reg_idx] {
            return;
        }
        if let Some(slot) = self.alloc_reg_slot[reg_idx] {
            let instr = Instruction::StoreIndU64 {
                base: STACK_PTR_REG,
                src: reg,
                offset: slot,
            };
            self.byte_offset += instr.encode().len();
            self.instructions.push(instr);
        }
        self.alloc_dirty[reg_idx] = false;
    }

    /// Flush ALL dirty allocated registers to their stack slots.
    pub fn spill_all_dirty_regs(&mut self) {
        let regs: Vec<u8> = self.regalloc.reg_to_slot.keys().copied().collect();
        for reg in regs {
            self.spill_dirty_reg(reg);
        }
    }

    /// Set `alloc_reg_slot` for a specific register to a given slot and mark dirty.
    ///
    /// Used after register-aware phi copies to tell the emitter that the
    /// allocated register now holds the phi destination value.
    /// Must be called AFTER `emit(MoveReg)` since `emit()` triggers `invalidate_reg()`.
    pub fn set_alloc_reg_for_slot(&mut self, reg: u8, slot: i32) {
        let idx = reg as usize;
        self.alloc_reg_slot[idx] = Some(slot);
        self.alloc_dirty[idx] = true;
        self.cache_slot(slot, reg);
    }

    /// Check if a register currently owns its expected allocated slot.
    pub fn is_alloc_reg_valid(&self, reg: u8, slot: i32) -> bool {
        self.alloc_reg_slot[reg as usize] == Some(slot)
    }

    /// Emit a raw `MoveReg` bypassing `emit()` (no `invalidate_reg` on dst).
    /// Used for register-aware phi copies where we manage alloc state manually.
    pub fn emit_raw_move(&mut self, dst: u8, src: u8) {
        let instr = Instruction::MoveReg { dst, src };
        self.byte_offset += instr.encode().len();
        self.instructions.push(instr);
        // Clear stale constant cache for dst — the register no longer holds
        // a known constant after the move.
        self.reg_to_const[dst as usize] = None;
    }

    /// Public wrapper for spilling a single dirty register.
    /// Used by the parallel move resolver in phi copies.
    pub fn spill_dirty_reg_pub(&mut self, reg: u8) {
        self.spill_dirty_reg(reg);
    }

    /// Spill register-allocated values before instructions that clobber them.
    /// With lazy spill, flushes dirty registers to stack.
    /// Without lazy spill (write-through), this is a no-op.
    pub fn spill_allocated_regs(&mut self) {
        if self.config.lazy_spill_enabled {
            self.spill_all_dirty_regs();
        }
    }

    fn invalidate_allocated_regs_where(&mut self, mut pred: impl FnMut(u8) -> bool) {
        let regs: Vec<u8> = self
            .regalloc
            .reg_to_slot
            .keys()
            .copied()
            .filter(|&r| pred(r))
            .collect();
        for reg in regs {
            self.invalidate_reg(reg);
        }
    }

    /// Invalidate allocated registers clobbered by scratch-heavy helpers.
    /// This is used after intrinsics that overwrite r5/r6.
    pub fn reload_allocated_regs_after_scratch_clobber(&mut self) {
        self.invalidate_allocated_regs_where(|r| {
            r == crate::abi::SCRATCH1 || r == crate::abi::SCRATCH2
        });
    }

    /// Invalidate allocated registers clobbered by a function call.
    ///
    /// Call argument setup may overwrite local argument registers (r9+), so
    /// allocated local-register mappings are invalidated after each call.
    /// This blanket variant invalidates all r9-r12; prefer
    /// `reload_allocated_regs_after_call_with_arity` when the call's argument
    /// count is known.
    pub fn reload_allocated_regs_after_call(&mut self) {
        self.reload_allocated_regs_after_call_with_arity(crate::abi::MAX_LOCAL_REGS);
    }

    /// Invalidate only the allocated registers actually clobbered by a call
    /// with `num_args` arguments. Registers r9..r9+min(num_args,4)-1 are
    /// clobbered by argument setup; higher registers remain valid.
    /// r5/r6 (scratch) and r7/r8 (caller-saved) are always clobbered by calls.
    /// Note: `clear_reg_cache()` already clears all alloc state before this is
    /// called, so this is largely a documentation/safety net.
    pub fn reload_allocated_regs_after_call_with_arity(&mut self, num_args: usize) {
        let clobbered_locals = num_args.min(crate::abi::MAX_LOCAL_REGS);
        self.invalidate_allocated_regs_where(|r| {
            r == crate::abi::SCRATCH1
                || r == crate::abi::SCRATCH2
                || r == crate::abi::RETURN_VALUE_REG
                || r == crate::abi::ARGS_LEN_REG
                || (r >= crate::abi::FIRST_LOCAL_REG
                    && r < crate::abi::FIRST_LOCAL_REG + clobbered_locals as u8)
        });
    }

    // ── Register cache ──

    /// Record that `reg` now holds the value of `slot`.
    pub fn cache_slot(&mut self, slot: i32, reg: u8) {
        if !self.config.register_cache_enabled {
            return;
        }
        // Remove any previous slot cached in this register.
        if let Some(old_slot) = self.reg_to_slot[reg as usize].take() {
            self.slot_cache.remove(&old_slot);
        }
        // Remove any previous register caching this slot.
        if let Some(old_reg) = self.slot_cache.insert(slot, reg) {
            self.reg_to_slot[old_reg as usize] = None;
        }
        self.reg_to_slot[reg as usize] = Some(slot);
    }

    /// Invalidate a register's cache entry (called when the register is overwritten).
    /// With lazy spill, if the register holds a dirty value, flush it to the
    /// stack before losing it.
    fn invalidate_reg(&mut self, reg: u8) {
        let idx = reg as usize;
        if let Some(slot) = self.reg_to_slot[idx].take() {
            self.slot_cache.remove(&slot);
        }
        if self.regalloc.reg_to_slot.contains_key(&reg) {
            // Lazy spill: if the register is dirty, spill to stack before clearing.
            if self.config.lazy_spill_enabled
                && self.alloc_dirty[idx]
                && let Some(slot) = self.alloc_reg_slot[idx]
            {
                let instr = Instruction::StoreIndU64 {
                    base: STACK_PTR_REG,
                    src: reg,
                    offset: slot,
                };
                self.byte_offset += instr.encode().len();
                self.instructions.push(instr);
            }
            self.alloc_reg_slot[idx] = None;
            self.alloc_dirty[idx] = false;
        }
        self.reg_to_const[idx] = None;
    }

    /// Clear the entire register cache (at block boundaries).
    ///
    /// Clears the general cache (`slot_cache`, `reg_to_slot`, `reg_to_const`) and
    /// allocated register state (`alloc_reg_slot`, `alloc_dirty`).
    pub fn clear_reg_cache(&mut self) {
        self.slot_cache.clear();
        self.reg_to_slot = [None; 13];
        self.reg_to_const = [None; 13];
        self.clear_allocated_reg_state();
    }

    /// Clear the general register cache but preserve allocated register state.
    ///
    /// Used at block boundaries in leaf functions where allocated registers
    /// are never clobbered by calls. With lazy spill, dirty values are preserved
    /// (they remain valid in registers across blocks in leaf functions).
    pub fn clear_reg_cache_preserving_alloc(&mut self) {
        self.slot_cache.clear();
        self.reg_to_slot = [None; 13];
        self.reg_to_const = [None; 13];
        // alloc_reg_slot and alloc_dirty intentionally NOT cleared for leaf functions.
    }

    fn clear_allocated_reg_state(&mut self) {
        for &reg in self.regalloc.reg_to_slot.keys() {
            let idx = reg as usize;
            self.alloc_reg_slot[idx] = None;
            self.alloc_dirty[idx] = false;
        }
    }

    /// Set allocated register state from a snapshot's `alloc_reg_slot`.
    /// Only sets entries where the snapshot has Some — does not clear others.
    pub fn set_alloc_reg_slot_from(&mut self, alloc_reg_slot: &[Option<i32>; 13]) {
        for &reg in self.regalloc.reg_to_slot.keys() {
            self.alloc_reg_slot[reg as usize] = alloc_reg_slot[reg as usize];
        }
    }

    /// Set allocated register state from a snapshot, but only for registers
    /// matching a filter predicate. Used for back-edge propagation where only
    /// callee-saved registers not clobbered by calls are safe to propagate.
    pub fn set_alloc_reg_slot_filtered(
        &mut self,
        alloc_reg_slot: &[Option<i32>; 13],
        filter: impl Fn(u8) -> bool,
    ) {
        for &reg in self.regalloc.reg_to_slot.keys() {
            if filter(reg) {
                self.alloc_reg_slot[reg as usize] = alloc_reg_slot[reg as usize];
            }
        }
    }

    /// Intersect allocated register state with a snapshot.
    /// Only keeps entries where both the current state and the snapshot agree.
    pub fn intersect_alloc_reg_slot(&mut self, other: &[Option<i32>; 13]) {
        for &reg in self.regalloc.reg_to_slot.keys() {
            let idx = reg as usize;
            if self.alloc_reg_slot[idx] != other[idx] {
                self.alloc_reg_slot[idx] = None;
            }
        }
    }

    /// Take a snapshot of the current register cache state.
    pub fn snapshot_cache(&self) -> CacheSnapshot {
        CacheSnapshot {
            slot_cache: self.slot_cache.clone(),
            reg_to_slot: self.reg_to_slot,
            reg_to_const: self.reg_to_const,
            alloc_reg_slot: self.alloc_reg_slot,
            alloc_dirty: self.alloc_dirty,
        }
    }

    /// Restore register cache state from a snapshot.
    pub fn restore_cache(&mut self, snapshot: &CacheSnapshot) {
        self.slot_cache.clone_from(&snapshot.slot_cache);
        self.reg_to_slot = snapshot.reg_to_slot;
        self.reg_to_const = snapshot.reg_to_const;
        self.alloc_reg_slot = snapshot.alloc_reg_slot;
        self.alloc_dirty = snapshot.alloc_dirty;
    }

    /// Define a label without clearing the register cache.
    /// Used for cross-block cache propagation when the cache will be restored from a snapshot.
    pub fn define_label_preserving_cache(&mut self, label: usize) {
        if self
            .instructions
            .last()
            .is_some_and(|last| !last.is_terminating())
        {
            self.emit(Instruction::Fallthrough);
        }
        self.labels[label] = Some(self.current_offset());
    }

    // ── Fixup resolution ──

    pub fn resolve_fixups(&mut self) -> Result<()> {
        // Precompute byte offsets for each instruction to avoid O(n²) re-scanning.
        let mut offsets = Vec::with_capacity(self.instructions.len());
        let mut running = 0usize;
        for instr in &self.instructions {
            offsets.push(running);
            running += instr.encode().len();
        }

        for &(instr_idx, label_id) in &self.fixups {
            let target_offset = self.labels[label_id]
                .ok_or_else(|| Error::Unsupported("unresolved label".to_string()))?;

            // PVM jump offsets are relative to the instruction start.
            let instr_start = offsets[instr_idx];
            let relative_offset = (target_offset as i32) - (instr_start as i32);

            match &mut self.instructions[instr_idx] {
                Instruction::Jump { offset }
                | Instruction::LoadImmJump { offset, .. }
                | Instruction::BranchNeImm { offset, .. }
                | Instruction::BranchEqImm { offset, .. }
                | Instruction::BranchGeSImm { offset, .. }
                | Instruction::BranchLtUImm { offset, .. }
                | Instruction::BranchLeUImm { offset, .. }
                | Instruction::BranchGeUImm { offset, .. }
                | Instruction::BranchGtUImm { offset, .. }
                | Instruction::BranchLtSImm { offset, .. }
                | Instruction::BranchLeSImm { offset, .. }
                | Instruction::BranchGtSImm { offset, .. }
                | Instruction::BranchEq { offset, .. }
                | Instruction::BranchNe { offset, .. }
                | Instruction::BranchGeU { offset, .. }
                | Instruction::BranchLtU { offset, .. }
                | Instruction::BranchLtS { offset, .. }
                | Instruction::BranchGeS { offset, .. } => {
                    *offset = relative_offset;
                }
                _ => {
                    return Err(Error::Unsupported(
                        "cannot fixup non-jump instruction".to_string(),
                    ));
                }
            }
        }
        Ok(())
    }
}

/// Try to extract a constant integer value from a `BasicValueEnum` without emitting instructions.
/// Returns `Some(i64)` for compile-time constants, `None` for SSA values.
pub fn try_get_constant(val: BasicValueEnum<'_>) -> Option<i64> {
    if let BasicValueEnum::IntValue(iv) = val {
        iv.get_sign_extended_constant()
    } else {
        None
    }
}

/// Get the i-th operand of an instruction as a `BasicValueEnum`.
pub fn get_operand(instr: InstructionValue<'_>, i: u32) -> Result<BasicValueEnum<'_>> {
    instr
        .get_operand(i)
        .and_then(inkwell::values::Operand::value)
        .ok_or_else(|| Error::Internal(format!("missing operand {i} in {:?}", instr.get_opcode())))
}

/// Get the i-th operand of an instruction as a `BasicBlock`.
pub fn get_bb_operand(instr: InstructionValue<'_>, i: u32) -> Result<BasicBlock<'_>> {
    instr
        .get_operand(i)
        .and_then(inkwell::values::Operand::block)
        .ok_or_else(|| {
            Error::Internal(format!(
                "missing bb operand {i} in {:?}",
                instr.get_opcode()
            ))
        })
}

/// Get the slot offset for an instruction's result.
pub fn result_slot(e: &PvmEmitter<'_>, instr: InstructionValue<'_>) -> Result<i32> {
    let key = val_key_instr(instr);
    e.get_slot(key)
        .ok_or_else(|| Error::Internal(format!("no slot for {:?} result", instr.get_opcode())))
}

/// Get the register to use for this instruction's result.
/// Returns the allocated register (for store-side coalescing) or `TEMP_RESULT` as fallback.
/// When the result has an allocated register, computing directly into it avoids a
/// subsequent `MoveReg` in `store_to_slot`.
pub fn result_reg(e: &PvmEmitter<'_>, instr: InstructionValue<'_>) -> u8 {
    result_reg_or(e, instr, crate::abi::TEMP_RESULT)
}

/// Like `result_reg` but with a custom fallback register.
/// Used by lowering paths that naturally use a different working register
/// (e.g., zext/sext/trunc use `TEMP1` instead of `TEMP_RESULT`).
pub fn result_reg_or(e: &PvmEmitter<'_>, instr: InstructionValue<'_>, fallback: u8) -> u8 {
    let key = val_key_instr(instr);
    if let Some(&alloc_reg) = e.regalloc.val_to_reg.get(&key) {
        alloc_reg
    } else {
        fallback
    }
}

/// Get the allocated register for a value if it's currently valid, otherwise return fallback.
/// Used for load-side coalescing: callers can use the allocated register directly as an
/// instruction operand instead of copying to a temp register via `load_operand`.
pub fn operand_reg(e: &PvmEmitter<'_>, val: BasicValueEnum<'_>, fallback: u8) -> u8 {
    if let BasicValueEnum::IntValue(iv) = val {
        // Constants don't have allocated registers.
        if iv.get_sign_extended_constant().is_some() || iv.get_zero_extended_constant().is_some() {
            return fallback;
        }
        if iv.is_poison() || iv.is_undef() {
            return fallback;
        }
        let key = val_key_int(iv);
        if let Some(&alloc_reg) = e.regalloc.val_to_reg.get(&key)
            && let Some(slot) = e.get_slot(key)
            && e.alloc_reg_slot[alloc_reg as usize] == Some(slot)
        {
            return alloc_reg;
        }
    }
    fallback
}

/// Detect the bit width of an instruction's **result** type.
///
/// For most instructions (binary ops, comparisons), this is the correct width
/// to use for choosing 32-bit vs 64-bit PVM instructions.
///
/// **Warning**: For conversion instructions (`SExt`, `ZExt`, `Trunc`), the result
/// type differs from the source type. Use [`source_bit_width`] instead when
/// you need the source operand's width (e.g., to determine which sign extension
/// to emit).
pub fn operand_bit_width(instr: InstructionValue<'_>) -> u32 {
    // Prefer the instruction's result type.
    if let inkwell::types::AnyTypeEnum::IntType(ty) = instr.get_type() {
        return ty.get_bit_width();
    }
    // Fallback: check the first operand's type.
    if let Some(op) = instr
        .get_operand(0)
        .and_then(inkwell::values::Operand::value)
        && let BasicValueEnum::IntValue(iv) = op
    {
        return iv.get_type().get_bit_width();
    }
    64 // default
}

/// Detect the bit width of an instruction's **source** (first operand) type.
///
/// This is the correct function for conversion instructions (`SExt`, `ZExt`, `Trunc`)
/// where you need to know what width the value is being converted *from*.
pub fn source_bit_width(instr: InstructionValue<'_>) -> u32 {
    if let Some(op) = instr
        .get_operand(0)
        .and_then(inkwell::values::Operand::value)
        && let BasicValueEnum::IntValue(iv) = op
    {
        return iv.get_type().get_bit_width();
    }
    64 // default
}

/// Check whether `target_bb` has any phi nodes with incomings from `current_bb`.
pub fn has_phi_from<'ctx>(current_bb: BasicBlock<'ctx>, target_bb: BasicBlock<'ctx>) -> bool {
    for instr in target_bb.get_instructions() {
        use inkwell::values::InstructionOpcode;
        if instr.get_opcode() != InstructionOpcode::Phi {
            break;
        }
        let Ok(phi): std::result::Result<PhiValue<'ctx>, _> = instr.try_into() else {
            break;
        };
        let num_incomings = phi.count_incoming();
        for i in 0..num_incomings {
            if let Some((_, block)) = phi.get_incoming(i)
                && block == current_bb
            {
                return true;
            }
        }
    }
    false
}

/// Check whether a basic block starts with any phi nodes.
pub fn block_has_phis(bb: BasicBlock<'_>) -> bool {
    bb.get_first_instruction()
        .is_some_and(|instr| instr.get_opcode() == inkwell::values::InstructionOpcode::Phi)
}

/// Pre-scan function to allocate labels and slots.
///
/// Also determines which callee-saved registers are actually used (for shrink wrapping).
pub fn pre_scan_function<'ctx>(
    emitter: &mut PvmEmitter<'ctx>,
    function: FunctionValue<'ctx>,
    is_main: bool,
) {
    // Detect real calls to determine if this is a leaf function.
    // PVM intrinsics (__pvm_load_*, __pvm_store_*, etc.) and LLVM intrinsics
    // (llvm.*) are NOT real function calls — they don't use the calling
    // convention and don't clobber callee-saved registers (r9-r12).
    // Only wasm_func_* and __pvm_call_indirect are real calls.
    let mut has_calls = false;
    for bb in function.get_basic_blocks() {
        for instr in bb.get_instructions() {
            if instr.get_opcode() == inkwell::values::InstructionOpcode::Call && is_real_call(instr)
            {
                has_calls = true;
                break;
            }
        }
        if has_calls {
            break;
        }
    }
    emitter.has_calls = has_calls;

    // Determine which callee-saved registers are used (shrink wrapping).
    if !is_main && emitter.config.shrink_wrap_enabled {
        let num_params = function.count_params() as usize;
        let mut used = [false; 4];

        // Parameters mapped to r9-r12 count as used.
        for u in used
            .iter_mut()
            .take(crate::abi::MAX_LOCAL_REGS.min(num_params))
        {
            *u = true;
        }

        // If the function contains any call instruction, all callee-saved regs are used
        // (because the callee may clobber them and expects us to preserve them).
        if has_calls {
            used = [true; 4];
        } else {
            // Check usage for non-call instructions if we assume registers might be used.
            // But since we don't allocate r9-r12 as temps, they are only used for params.
            // So `used` is already correct.
        }

        emitter.used_callee_regs = used;

        // Compute frame offsets for saved callee-saved registers.
        // Layout: [ra (optional), then used callee-saved regs contiguously].
        let mut offset = if has_calls { 8i32 } else { 0i32 }; // Reserve space for ra only if needed
        let mut offsets = [None; 4];
        for i in 0..4 {
            if used[i] {
                offsets[i] = Some(offset);
                offset += 8;
            }
        }
        emitter.callee_save_offsets = offsets;
        emitter.next_slot_offset = offset;
    }
    // When shrink wrapping is disabled or is_main, keep defaults (all regs, FRAME_HEADER_SIZE).

    // Compute single-predecessor map for cross-block register cache.
    if emitter.config.cross_block_cache_enabled && emitter.config.register_cache_enabled {
        let blocks = function.get_basic_blocks();
        let mut pred_count: HashMap<BasicBlock<'ctx>, usize> = HashMap::new();
        let mut pred_from: HashMap<BasicBlock<'ctx>, BasicBlock<'ctx>> = HashMap::new();

        for bb in &blocks {
            if let Some(term) = bb.get_terminator() {
                let successors = super::successors::collect_successors(term);
                // Deduplicate successors per predecessor (e.g. switch cases targeting the same block)
                // so that multiple edges from the same bb don't inflate the predecessor count.
                let unique_succs: HashSet<_> = successors.into_iter().collect();
                for succ in unique_succs {
                    let count = pred_count.entry(succ).or_insert(0);
                    *count += 1;
                    pred_from.insert(succ, *bb);
                }
            }
        }

        for (bb, count) in &pred_count {
            if *count == 1
                && let Some(pred) = pred_from.get(bb)
            {
                emitter.block_single_pred.insert(*bb, *pred);
            }
        }
    }

    // Allocate slots for function parameters.
    for param in function.get_params() {
        let key = val_key_basic(param);
        emitter.alloc_slot_for_key(key);
    }

    // Allocate labels for all basic blocks.
    for bb in function.get_basic_blocks() {
        let label = emitter.alloc_label();
        emitter.block_labels.insert(bb, label);
    }

    // Allocate slots for instruction results that produce integer values.
    for bb in function.get_basic_blocks() {
        for instr in bb.get_instructions() {
            if instruction_produces_value(instr) {
                let key = val_key_instr(instr);
                emitter.alloc_slot_for_key(key);
            }
        }
    }
}

fn instruction_produces_value(instr: InstructionValue<'_>) -> bool {
    use inkwell::values::InstructionOpcode;
    matches!(
        instr.get_opcode(),
        InstructionOpcode::Add
            | InstructionOpcode::Sub
            | InstructionOpcode::Mul
            | InstructionOpcode::UDiv
            | InstructionOpcode::SDiv
            | InstructionOpcode::URem
            | InstructionOpcode::SRem
            | InstructionOpcode::And
            | InstructionOpcode::Or
            | InstructionOpcode::Xor
            | InstructionOpcode::Shl
            | InstructionOpcode::LShr
            | InstructionOpcode::AShr
            | InstructionOpcode::ICmp
            | InstructionOpcode::ZExt
            | InstructionOpcode::SExt
            | InstructionOpcode::Trunc
            | InstructionOpcode::Select
            | InstructionOpcode::Phi
            | InstructionOpcode::Load
            | InstructionOpcode::Call
    )
}

/// Check whether an LLVM Call instruction is a "real" function call that uses the
/// calling convention and may clobber callee-saved registers (r9-r12).
///
/// PVM intrinsics (`__pvm_load_*`, `__pvm_store_*`, `__pvm_memory_*`, etc.) and
/// LLVM intrinsics (`llvm.*`) are lowered inline — they only use temp/scratch
/// registers and never clobber callee-saved registers.
///
/// Real calls: `wasm_func_*` (direct) and `__pvm_call_indirect` (indirect).
pub(super) fn is_real_call(instr: InstructionValue<'_>) -> bool {
    let call_site: std::result::Result<inkwell::values::CallSiteValue, _> = instr.try_into();
    let Ok(call_site) = call_site else {
        return true; // Conservative: treat as call if we can't classify.
    };
    let Some(fn_val) = call_site.get_called_fn_value() else {
        return true; // Indirect call without known callee.
    };
    let name = fn_val.get_name().to_string_lossy();
    // PVM intrinsics and LLVM intrinsics are NOT real calls.
    if name.starts_with("__pvm_") && name != "__pvm_call_indirect" {
        return false;
    }
    if name.starts_with("llvm.") {
        return false;
    }
    true
}

/// Check whether any instruction in the function will lower to a path that
/// clobbers r5/r6 (`abi::SCRATCH1`/`abi::SCRATCH2`).
///
/// Returns `true` if the function is safe (no r5/r6 clobbers).
///
/// Operations that clobber r5/r6:
/// - `__pvm_memory_grow`, `__pvm_memory_fill`, `__pvm_memory_copy`, `__pvm_memory_init`
/// - `llvm.fshl.*`, `llvm.fshr.*` (funnel shifts — conservatively flagged even if
///   they might lower to a rotation which doesn't clobber)
pub fn scratch_regs_safe(function: FunctionValue<'_>) -> bool {
    use inkwell::values::InstructionOpcode;

    for bb in function.get_basic_blocks() {
        for instr in bb.get_instructions() {
            if instr.get_opcode() == InstructionOpcode::Call {
                let call_site: std::result::Result<inkwell::values::CallSiteValue, _> =
                    instr.try_into();
                let Ok(cs) = call_site else {
                    continue;
                };
                let Some(fn_val) = cs.get_called_fn_value() else {
                    continue;
                };
                let name = fn_val.get_name().to_string_lossy();
                // Bulk memory intrinsics clobber r5/r6.
                if matches!(
                    name.as_ref(),
                    "__pvm_memory_grow"
                        | "__pvm_memory_fill"
                        | "__pvm_memory_copy"
                        | "__pvm_memory_init"
                ) {
                    return false;
                }
                // Funnel shifts (conservatively including rotations).
                if name.starts_with("llvm.fshl.") || name.starts_with("llvm.fshr.") {
                    return false;
                }
            }
        }
    }
    true
}

// Re-export scratch registers for other modules
pub use crate::abi::{ARGS_LEN_REG as SCRATCH1, ARGS_PTR_REG as SCRATCH2};