wasm_pvm/translate/

mod.rs

// Address calculations and jump offsets often require wrapping/truncation.
#![allow(
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::cast_sign_loss
)]

pub mod adapter_merge;
pub mod dead_function_elimination;
pub mod memory_layout;
pub mod wasm_module;

use std::collections::HashMap;

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

pub use wasm_module::WasmModule;

/// Action to take when a WASM import is called.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ImportAction {
    /// Emit a trap (unreachable) instruction.
    Trap,
    /// Emit a no-op (return 0 for functions with return values).
    Nop,
}

/// Flags to enable/disable individual compiler optimizations.
/// All optimizations are enabled by default.
#[derive(Debug, Clone)]
#[allow(clippy::struct_excessive_bools)]
pub struct OptimizationFlags {
    /// Run LLVM optimization passes (mem2reg, instcombine, simplifycfg, gvn, dce).
    /// When false, also disables inlining (all LLVM passes are skipped).
    pub llvm_passes: bool,
    /// Run peephole optimizer (fallthrough removal, dead code elimination).
    pub peephole: bool,
    /// Enable per-block register cache (store-load forwarding).
    pub register_cache: bool,
    /// Fuse `ICmp` + Branch into a single PVM branch instruction.
    pub icmp_branch_fusion: bool,
    /// Only save/restore callee-saved registers (r9-r12) that are actually used.
    pub shrink_wrap_callee_saves: bool,
    /// Eliminate SP-relative stores whose target offset is never loaded from.
    pub dead_store_elimination: bool,
    /// Skip redundant `LoadImm`/`LoadImm64` when the register already holds the constant.
    pub constant_propagation: bool,
    /// Inline small functions at the LLVM IR level to eliminate call overhead.
    pub inlining: bool,
    /// Propagate register cache across single-predecessor block boundaries.
    pub cross_block_cache: bool,
    /// Allocate long-lived SSA values to physical registers (r5, r6) across block boundaries.
    pub register_allocation: bool,
    /// Eliminate unreachable functions not called from entry points or the function table.
    pub dead_function_elimination: bool,
    /// Eliminate unconditional jumps to the immediately following block (fallthrough).
    pub fallthrough_jumps: bool,
}

impl Default for OptimizationFlags {
    fn default() -> Self {
        Self {
            llvm_passes: true,
            peephole: true,
            register_cache: true,
            icmp_branch_fusion: true,
            shrink_wrap_callee_saves: true,
            dead_store_elimination: true,
            constant_propagation: true,
            inlining: true,
            cross_block_cache: true,
            register_allocation: true,
            dead_function_elimination: true,
            fallthrough_jumps: true,
        }
    }
}

/// Options for compilation.
#[derive(Debug, Clone, Default)]
pub struct CompileOptions {
    /// Mapping from import function names to actions.
    /// When provided, all imports (except known intrinsics like `host_call` and `pvm_ptr`)
    /// must have a mapping or compilation will fail with `UnresolvedImport`.
    pub import_map: Option<HashMap<String, ImportAction>>,
    /// WAT source for an adapter module whose exports replace matching main imports.
    /// Applied before the text-based import map, so the two compose.
    pub adapter: Option<String>,
    /// Metadata blob to prepend to the SPI output.
    /// Typically contains the source filename and compiler version.
    pub metadata: Vec<u8>,
    /// Optimization flags controlling which compiler passes are enabled.
    pub optimizations: OptimizationFlags,
}

// Re-export register constants from abi module
pub use crate::abi::{ARGS_LEN_REG, ARGS_PTR_REG, RETURN_ADDR_REG, STACK_PTR_REG};

// ── Call fixup types (shared with LLVM backend) ──

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

#[derive(Debug, Clone)]
pub struct IndirectCallFixup {
    pub return_addr_instr: usize,
    // For `LoadImmJumpInd`, this equals `return_addr_instr`.
    pub jump_ind_instr: usize,
}

/// `RO_DATA` region size is 64KB (0x10000 to 0x1FFFF)
const RO_DATA_SIZE: usize = 64 * 1024;

pub fn compile(wasm: &[u8]) -> Result<SpiProgram> {
    compile_with_options(wasm, &CompileOptions::default())
}

pub fn compile_with_options(wasm: &[u8], options: &CompileOptions) -> Result<SpiProgram> {
    // Known intrinsics that don't need import mappings.
    const KNOWN_INTRINSICS: &[&str] = &["host_call", "pvm_ptr"];
    // Default mappings applied when no explicit import map is provided.
    const DEFAULT_MAPPINGS: &[&str] = &["abort"];

    // Apply adapter merge if provided (produces a new WASM binary with fewer imports).
    let merged_wasm;
    let wasm = if let Some(adapter_wat) = &options.adapter {
        merged_wasm = adapter_merge::merge_adapter(wasm, adapter_wat)?;
        &merged_wasm
    } else {
        wasm
    };

    let module = WasmModule::parse(wasm)?;

    // Validate that all imports are resolved.
    for name in &module.imported_func_names {
        if KNOWN_INTRINSICS.contains(&name.as_str()) {
            continue;
        }
        if let Some(import_map) = &options.import_map {
            if import_map.contains_key(name) {
                continue;
            }
        } else if DEFAULT_MAPPINGS.contains(&name.as_str()) {
            continue;
        }
        return Err(Error::UnresolvedImport(format!(
            "import '{name}' has no mapping. Provide a mapping via --imports or add it to the import map."
        )));
    }

    compile_via_llvm(&module, options)
}

pub fn compile_via_llvm(module: &WasmModule, options: &CompileOptions) -> Result<SpiProgram> {
    use crate::llvm_backend::{self, LoweringContext};
    use crate::llvm_frontend;
    use inkwell::context::Context;

    // Phase 0: Dead function elimination — compute reachable set.
    let reachable_locals = if options.optimizations.dead_function_elimination {
        Some(dead_function_elimination::reachable_functions(module)?)
    } else {
        None
    };

    // Phase 1: WASM → LLVM IR
    let context = Context::create();
    let llvm_module = llvm_frontend::translate_wasm_to_llvm(
        &context,
        module,
        options.optimizations.llvm_passes,
        options.optimizations.inlining,
        reachable_locals.as_ref(),
    )?;

    // Calculate RO_DATA offsets and lengths for passive data segments
    let mut data_segment_offsets = std::collections::HashMap::new();
    let mut data_segment_lengths = std::collections::HashMap::new();
    let mut current_ro_offset = if module.function_table.is_empty() {
        1 // dummy byte if no function table
    } else {
        module.function_table.len() * 8 // jump_ref + type_idx per entry
    };

    let mut data_segment_length_addrs = std::collections::HashMap::new();
    let mut passive_ordinal = 0usize;

    for (idx, seg) in module.data_segments.iter().enumerate() {
        if seg.offset.is_none() {
            // Check that segment fits within RO_DATA region
            if current_ro_offset + seg.data.len() > RO_DATA_SIZE {
                return Err(Error::Internal(format!(
                    "passive data segment {} (size {}) would overflow RO_DATA region ({} bytes used of {})",
                    idx,
                    seg.data.len(),
                    current_ro_offset,
                    RO_DATA_SIZE
                )));
            }
            data_segment_offsets.insert(idx as u32, current_ro_offset as u32);
            data_segment_lengths.insert(idx as u32, seg.data.len() as u32);
            data_segment_length_addrs.insert(
                idx as u32,
                memory_layout::data_segment_length_offset(module.globals.len(), passive_ordinal),
            );
            current_ro_offset += seg.data.len();
            passive_ordinal += 1;
        }
    }

    // Phase 2: Build lowering context
    let ctx = LoweringContext {
        wasm_memory_base: module.wasm_memory_base,
        num_globals: module.globals.len(),
        function_signatures: module.function_signatures.clone(),
        type_signatures: module.type_signatures.clone(),
        function_table: module.function_table.clone(),
        num_imported_funcs: module.num_imported_funcs as usize,
        imported_func_names: module.imported_func_names.clone(),
        initial_memory_pages: module.memory_limits.initial_pages,
        max_memory_pages: module.max_memory_pages,
        stack_size: memory_layout::DEFAULT_STACK_SIZE,
        data_segment_offsets,
        data_segment_lengths,
        data_segment_length_addrs,
        wasm_import_map: options.import_map.clone(),
        optimizations: options.optimizations.clone(),
    };

    // Phase 3: LLVM IR → PVM bytecode for each function
    let mut all_instructions: Vec<Instruction> = Vec::new();
    let mut all_call_fixups: Vec<(usize, CallFixup)> = Vec::new();
    let mut all_indirect_call_fixups: Vec<(usize, IndirectCallFixup)> = Vec::new();
    let mut function_offsets: Vec<usize> = vec![0; module.functions.len()];
    let mut next_call_return_idx: usize = 0;

    // Entry header: Jump to main (PC=0) + Trap or secondary Jump (PC=5).
    // When there's no secondary entry, we omit the Fallthrough padding (6 bytes instead of 10).
    all_instructions.push(Instruction::Jump { offset: 0 });
    if module.has_secondary_entry {
        all_instructions.push(Instruction::Jump { offset: 0 });
    } else {
        all_instructions.push(Instruction::Trap);
    }

    // Build emission order: main first, then secondary (if any), then remaining in index order.
    // This places main immediately after the entry header, minimizing the entry Jump distance.
    let mut emission_order: Vec<usize> = Vec::with_capacity(module.functions.len());
    emission_order.push(module.main_func_local_idx);
    if let Some(secondary_idx) = module.secondary_entry_local_idx
        && secondary_idx != module.main_func_local_idx
    {
        emission_order.push(secondary_idx);
    }
    for idx in 0..module.functions.len() {
        if idx != module.main_func_local_idx && module.secondary_entry_local_idx != Some(idx) {
            emission_order.push(idx);
        }
    }

    for &local_func_idx in &emission_order {
        // Dead functions: emit a single Trap as a placeholder.
        // The function offset is still recorded so dispatch table indices stay valid.
        if reachable_locals
            .as_ref()
            .is_some_and(|r| !r.contains(&local_func_idx))
        {
            let func_start_offset: usize = all_instructions.iter().map(|i| i.encode().len()).sum();
            function_offsets[local_func_idx] = func_start_offset;
            all_instructions.push(Instruction::Trap);
            continue;
        }

        let global_func_idx = module.num_imported_funcs as usize + local_func_idx;
        let fn_name = format!("wasm_func_{global_func_idx}");
        let llvm_func = llvm_module
            .get_function(&fn_name)
            .ok_or_else(|| Error::Internal(format!("missing LLVM function: {fn_name}")))?;

        let is_main = local_func_idx == module.main_func_local_idx;
        let is_secondary = module.secondary_entry_local_idx == Some(local_func_idx);
        let is_entry = is_main || is_secondary;

        let func_start_offset: usize = all_instructions.iter().map(|i| i.encode().len()).sum();
        function_offsets[local_func_idx] = func_start_offset;

        // If entry function and there's a start function, call it first.
        if let Some(start_local_idx) = module.start_func_local_idx.filter(|_| is_entry) {
            // Save r7 and r8 to stack.
            all_instructions.push(Instruction::AddImm64 {
                dst: STACK_PTR_REG,
                src: STACK_PTR_REG,
                value: -16,
            });
            all_instructions.push(Instruction::StoreIndU64 {
                base: STACK_PTR_REG,
                src: ARGS_PTR_REG,
                offset: 0,
            });
            all_instructions.push(Instruction::StoreIndU64 {
                base: STACK_PTR_REG,
                src: ARGS_LEN_REG,
                offset: 8,
            });

            // Call start function using LoadImmJump (combined load + jump).
            let call_return_addr = ((next_call_return_idx + 1) * 2) as i32;
            next_call_return_idx += 1;
            let current_instr_idx = all_instructions.len();
            all_instructions.push(Instruction::LoadImmJump {
                reg: RETURN_ADDR_REG,
                value: call_return_addr,
                offset: 0, // patched during fixup resolution
            });

            all_call_fixups.push((
                current_instr_idx,
                CallFixup {
                    target_func: start_local_idx as u32,
                    return_addr_instr: 0,
                    jump_instr: 0, // same instruction for LoadImmJump
                },
            ));

            // Restore r7 and r8.
            all_instructions.push(Instruction::LoadIndU64 {
                dst: ARGS_PTR_REG,
                base: STACK_PTR_REG,
                offset: 0,
            });
            all_instructions.push(Instruction::LoadIndU64 {
                dst: ARGS_LEN_REG,
                base: STACK_PTR_REG,
                offset: 8,
            });
            all_instructions.push(Instruction::AddImm64 {
                dst: STACK_PTR_REG,
                src: STACK_PTR_REG,
                value: 16,
            });
        }

        let translation = llvm_backend::lower_function(
            llvm_func,
            &ctx,
            is_entry,
            global_func_idx,
            next_call_return_idx,
        )?;
        next_call_return_idx += translation.num_call_returns;

        let instr_base = all_instructions.len();
        for fixup in translation.call_fixups {
            all_call_fixups.push((
                instr_base,
                CallFixup {
                    return_addr_instr: fixup.return_addr_instr,
                    jump_instr: fixup.jump_instr,
                    target_func: fixup.target_func,
                },
            ));
        }
        for fixup in translation.indirect_call_fixups {
            all_indirect_call_fixups.push((
                instr_base,
                IndirectCallFixup {
                    return_addr_instr: fixup.return_addr_instr,
                    jump_ind_instr: fixup.jump_ind_instr,
                },
            ));
        }

        all_instructions.extend(translation.instructions);
    }

    // Phase 4: Resolve call fixups and build jump table.
    let (jump_table, func_entry_jump_table_base) = resolve_call_fixups(
        &mut all_instructions,
        &all_call_fixups,
        &all_indirect_call_fixups,
        &function_offsets,
    )?;

    // Patch entry header jumps.
    let main_offset = function_offsets[module.main_func_local_idx] as i32;
    if let Instruction::Jump { offset } = &mut all_instructions[0] {
        *offset = main_offset;
    }

    if let Some(secondary_idx) = module.secondary_entry_local_idx {
        let secondary_offset = function_offsets[secondary_idx] as i32 - 5;
        if let Instruction::Jump { offset } = &mut all_instructions[1] {
            *offset = secondary_offset;
        }
    }

    // Phase 5: Build dispatch table for call_indirect.
    let mut ro_data = vec![0u8];
    if !module.function_table.is_empty() {
        ro_data.clear();
        for &func_idx in &module.function_table {
            if func_idx == u32::MAX || (func_idx as usize) < module.num_imported_funcs as usize {
                ro_data.extend_from_slice(&u32::MAX.to_le_bytes());
                ro_data.extend_from_slice(&u32::MAX.to_le_bytes());
            } else {
                let local_func_idx = func_idx as usize - module.num_imported_funcs as usize;
                let jump_ref = 2 * (func_entry_jump_table_base + local_func_idx + 1) as u32;
                ro_data.extend_from_slice(&jump_ref.to_le_bytes());
                let type_idx = *module
                    .function_type_indices
                    .get(local_func_idx)
                    .unwrap_or(&u32::MAX);
                ro_data.extend_from_slice(&type_idx.to_le_bytes());
            }
        }
    }

    // Append passive data segments to RO_DATA.
    // NOTE: This loop must iterate data_segments in the same order as the offset
    // calculation loop above, since data_segment_offsets indices depend on it.
    for seg in &module.data_segments {
        if seg.offset.is_none() {
            ro_data.extend_from_slice(&seg.data);
        }
    }

    let blob = crate::pvm::ProgramBlob::new(all_instructions).with_jump_table(jump_table);
    let rw_data_section = build_rw_data(
        &module.data_segments,
        &module.global_init_values,
        module.memory_limits.initial_pages,
        module.wasm_memory_base,
        &ctx.data_segment_length_addrs,
        &ctx.data_segment_lengths,
    );

    let heap_pages = calculate_heap_pages(
        rw_data_section.len(),
        module.wasm_memory_base,
        module.memory_limits.initial_pages,
        module.functions.len(),
    )?;

    Ok(SpiProgram::new(blob)
        .with_heap_pages(heap_pages)
        .with_ro_data(ro_data)
        .with_rw_data(rw_data_section)
        .with_metadata(options.metadata.clone()))
}

/// Calculate the number of 4KB PVM heap pages needed after `rw_data`.
///
/// `heap_pages` tells the runtime how many zero-initialized writable pages to allocate
/// immediately after the `rw_data` blob. This covers the initial WASM linear memory,
/// globals, and spilled locals that aren't already covered by `rw_data`.
///
/// By computing this **after** `build_rw_data()`, we use the actual (trimmed) `rw_data`
/// length instead of guessing with headroom.
///
/// We add 1 extra page beyond the exact initial memory requirement. This ensures that
/// the first `memory.grow` / sbrk allocation has a pre-allocated page available at the
/// boundary of the initial WASM memory. Without it, PVM-in-PVM execution fails because
/// the inner interpreter's page-fault handling at the exact heap boundary doesn't
/// correctly propagate through the outer PVM.
fn calculate_heap_pages(
    rw_data_len: usize,
    wasm_memory_base: i32,
    initial_pages: u32,
    num_functions: usize,
) -> Result<u16> {
    use wasm_module::MIN_INITIAL_WASM_PAGES;

    let initial_pages = initial_pages.max(MIN_INITIAL_WASM_PAGES);
    let wasm_memory_initial_end = wasm_memory_base as usize + (initial_pages as usize) * 64 * 1024;

    let spilled_locals_end = memory_layout::SPILLED_LOCALS_BASE as usize
        + num_functions * memory_layout::SPILLED_LOCALS_PER_FUNC as usize;

    let end = spilled_locals_end.max(wasm_memory_initial_end);
    let total_bytes = end - memory_layout::GLOBAL_MEMORY_BASE as usize;
    let rw_pages = rw_data_len.div_ceil(4096);
    let total_pages = total_bytes.div_ceil(4096);
    let heap_pages = total_pages.saturating_sub(rw_pages) + 1;

    u16::try_from(heap_pages).map_err(|_| {
        Error::Internal(format!(
            "heap size {heap_pages} pages exceeds u16::MAX ({}) — module too large",
            u16::MAX
        ))
    })
}

/// Build the `rw_data` section from WASM data segments and global initializers.
pub(crate) fn build_rw_data(
    data_segments: &[wasm_module::DataSegment],
    global_init_values: &[i32],
    initial_memory_pages: u32,
    wasm_memory_base: i32,
    data_segment_length_addrs: &std::collections::HashMap<u32, i32>,
    data_segment_lengths: &std::collections::HashMap<u32, u32>,
) -> Vec<u8> {
    // Calculate the minimum size needed for globals
    // +1 for the compiler-managed memory size global, plus passive segment lengths
    let num_passive_segments = data_segment_length_addrs.len();
    let globals_end =
        memory_layout::globals_region_size(global_init_values.len(), num_passive_segments);

    // Calculate the size needed for data segments
    let wasm_to_rw_offset = wasm_memory_base as u32 - 0x30000;

    let data_end = data_segments
        .iter()
        .filter_map(|seg| {
            seg.offset
                .map(|off| wasm_to_rw_offset + off + seg.data.len() as u32)
        })
        .max()
        .unwrap_or(0) as usize;

    let total_size = globals_end.max(data_end);

    if total_size == 0 {
        return Vec::new();
    }

    let mut rw_data = vec![0u8; total_size];

    // Initialize user globals
    for (i, &value) in global_init_values.iter().enumerate() {
        let offset = i * 4;
        if offset + 4 <= rw_data.len() {
            rw_data[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
        }
    }

    // Initialize compiler-managed memory size global (right after user globals)
    let mem_size_offset = global_init_values.len() * 4;
    if mem_size_offset + 4 <= rw_data.len() {
        rw_data[mem_size_offset..mem_size_offset + 4]
            .copy_from_slice(&initial_memory_pages.to_le_bytes());
    }

    // Initialize passive data segment effective lengths (right after memory size global).
    // These are used by memory.init for bounds checking and zeroed by data.drop.
    for (&seg_idx, &addr) in data_segment_length_addrs {
        if let Some(&length) = data_segment_lengths.get(&seg_idx) {
            // addr is absolute PVM address; convert to rw_data offset
            let rw_offset = (addr - memory_layout::GLOBAL_MEMORY_BASE) as usize;
            if rw_offset + 4 <= rw_data.len() {
                rw_data[rw_offset..rw_offset + 4].copy_from_slice(&length.to_le_bytes());
            }
        }
    }

    // Copy data segments to their WASM memory locations
    for seg in data_segments {
        if let Some(offset) = seg.offset {
            let rw_offset = (wasm_to_rw_offset + offset) as usize;
            if rw_offset + seg.data.len() <= rw_data.len() {
                rw_data[rw_offset..rw_offset + seg.data.len()].copy_from_slice(&seg.data);
            }
        }
    }

    // Trim trailing zeros to reduce SPI size. Heap pages are zero-initialized,
    // so omitted high-address zero bytes are semantically equivalent.
    if let Some(last_non_zero) = rw_data.iter().rposition(|&b| b != 0) {
        rw_data.truncate(last_non_zero + 1);
    } else {
        rw_data.clear();
    }

    rw_data
}

/// Extract the pre-assigned jump-table index from a return-address load instruction.
///
/// Call return addresses are pre-assigned as `(idx + 1) * 2` at emission time.
/// This helper recovers `idx` so that `resolve_call_fixups` can write the byte
/// offset into the correct jump-table slot instead of appending in list order
/// (which would desync when a function mixes direct and indirect calls).
///
/// Direct calls use `LoadImmJump`, while indirect calls use either `LoadImm` (legacy
/// two-instruction sequence) or `LoadImmJumpInd` (combined return-addr load + jump).
fn return_addr_jump_table_idx(
    instructions: &[Instruction],
    return_addr_instr: usize,
) -> Result<usize> {
    let value = match instructions.get(return_addr_instr) {
        Some(
            Instruction::LoadImmJump { value, .. }
            | Instruction::LoadImm { value, .. }
            | Instruction::LoadImmJumpInd { value, .. },
        ) => Some(*value),
        _ => None,
    };
    match value {
        Some(v) if v > 0 && v % 2 == 0 => Ok((v as usize / 2) - 1),
        _ => Err(Error::Internal(format!(
            "expected LoadImmJump/LoadImm/LoadImmJumpInd((idx+1)*2) at return_addr_instr {return_addr_instr}, got {:?}",
            instructions.get(return_addr_instr)
        ))),
    }
}

fn resolve_call_fixups(
    instructions: &mut [Instruction],
    call_fixups: &[(usize, CallFixup)],
    indirect_call_fixups: &[(usize, IndirectCallFixup)],
    function_offsets: &[usize],
) -> Result<(Vec<u32>, usize)> {
    // Count total call-return entries by finding the maximum pre-assigned index.
    // Entries are written at their pre-assigned slot so mixed direct/indirect
    // call ordering within a function is preserved correctly.
    let mut num_call_returns: usize = 0;

    for (instr_base, fixup) in call_fixups {
        let idx = return_addr_jump_table_idx(instructions, instr_base + fixup.return_addr_instr)?;
        num_call_returns = num_call_returns.max(idx + 1);
    }
    for (instr_base, fixup) in indirect_call_fixups {
        let idx = return_addr_jump_table_idx(instructions, instr_base + fixup.return_addr_instr)?;
        num_call_returns = num_call_returns.max(idx + 1);
    }

    let mut jump_table: Vec<u32> = vec![0u32; num_call_returns];

    // Call return addresses (LoadImmJump/LoadImm/LoadImmJumpInd values) are pre-assigned at emission time,
    // so we only need to compute byte offsets for the jump table and patch Jump targets.
    // Write each entry at its pre-assigned index to keep values in sync.
    for (instr_base, fixup) in call_fixups {
        let target_offset = function_offsets
            .get(fixup.target_func as usize)
            .ok_or_else(|| {
                Error::Unsupported(format!("call to unknown function {}", fixup.target_func))
            })?;

        let jump_idx = instr_base + fixup.jump_instr;

        // Return address = byte offset after the LoadImmJump instruction.
        let return_addr_offset: usize = instructions[..=jump_idx]
            .iter()
            .map(|i| i.encode().len())
            .sum();

        let slot = return_addr_jump_table_idx(instructions, instr_base + fixup.return_addr_instr)?;
        jump_table[slot] = return_addr_offset as u32;

        // Verify pre-assigned jump table address matches actual index.
        let expected_addr = ((slot + 1) * 2) as i32;
        debug_assert!(
            matches!(&instructions[jump_idx], Instruction::LoadImmJump { value, .. } if *value == expected_addr),
            "pre-assigned jump table address mismatch: expected {expected_addr}, got {:?}",
            &instructions[jump_idx]
        );

        // Patch the offset field of LoadImmJump.
        let jump_start_offset: usize = instructions[..jump_idx]
            .iter()
            .map(|i| i.encode().len())
            .sum();
        let relative_offset = (*target_offset as i32) - (jump_start_offset as i32);

        if let Instruction::LoadImmJump { offset, .. } = &mut instructions[jump_idx] {
            *offset = relative_offset;
        }
    }

    for (instr_base, fixup) in indirect_call_fixups {
        let jump_ind_idx = instr_base + fixup.jump_ind_instr;

        let return_addr_offset: usize = instructions[..=jump_ind_idx]
            .iter()
            .map(|i| i.encode().len())
            .sum();

        let slot = return_addr_jump_table_idx(instructions, instr_base + fixup.return_addr_instr)?;
        jump_table[slot] = return_addr_offset as u32;
    }

    let func_entry_base = jump_table.len();
    for &offset in function_offsets {
        jump_table.push(offset as u32);
    }

    Ok((jump_table, func_entry_base))
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::build_rw_data;
    use super::memory_layout;
    use super::wasm_module::DataSegment;

    #[test]
    fn build_rw_data_trims_all_zero_tail_to_empty() {
        let rw = build_rw_data(&[], &[], 0, 0x30000, &HashMap::new(), &HashMap::new());
        assert!(rw.is_empty());
    }

    #[test]
    fn build_rw_data_preserves_internal_zeros_and_trims_trailing_zeros() {
        let data_segments = vec![DataSegment {
            offset: Some(0),
            data: vec![1, 0, 2, 0, 0],
        }];

        let rw = build_rw_data(
            &data_segments,
            &[],
            0,
            0x30000,
            &HashMap::new(),
            &HashMap::new(),
        );

        assert_eq!(rw, vec![1, 0, 2]);
    }

    #[test]
    fn build_rw_data_keeps_non_zero_passive_length_bytes() {
        let mut addrs = HashMap::new();
        addrs.insert(0u32, memory_layout::GLOBAL_MEMORY_BASE + 4);
        let mut lengths = HashMap::new();
        lengths.insert(0u32, 7u32);

        let rw = build_rw_data(&[], &[], 0, 0x30000, &addrs, &lengths);

        assert_eq!(rw, vec![0, 0, 0, 0, 7]);
    }

    // ── calculate_heap_pages tests ──

    #[test]
    fn heap_pages_with_empty_rw_data_equals_total_pages_plus_one() {
        // wasm_memory_base = 0x33000 (typical), initial_pages = 0 (clamped to 16)
        // end = 0x33000 + 16*64*1024 = 0x33000 + 0x100000 = 0x133000
        // total_bytes = 0x133000 - 0x30000 = 0x103000 = 1060864
        // total_pages = ceil(1060864 / 4096) = 259
        // rw_pages = 0, heap_pages = 259 + 1 = 260
        let pages = super::calculate_heap_pages(0, 0x33000, 0, 10).unwrap();
        assert_eq!(pages, 260);
    }

    #[test]
    fn heap_pages_reduced_by_rw_data_pages() {
        // Same scenario but with 8192 bytes of rw_data (2 pages)
        let pages_no_rw = super::calculate_heap_pages(0, 0x33000, 0, 10).unwrap();
        let pages_with_rw = super::calculate_heap_pages(8192, 0x33000, 0, 10).unwrap();
        assert_eq!(pages_no_rw - pages_with_rw, 2);
    }

    #[test]
    fn heap_pages_saturates_at_one_for_large_rw_data() {
        // rw_data that covers more than total_pages still gets +1 headroom
        let pages = super::calculate_heap_pages(2 * 1024 * 1024, 0x33000, 0, 10).unwrap();
        assert_eq!(pages, 1);
    }

    #[test]
    fn heap_pages_respects_initial_pages() {
        // initial_pages = 32 (larger than MIN_INITIAL_WASM_PAGES=16)
        // end = 0x33000 + 32*64*1024 = 0x33000 + 0x200000 = 0x233000
        // total_bytes = 0x233000 - 0x30000 = 0x203000
        // total_pages = ceil(0x203000 / 4096) = 515
        // heap_pages = 515 + 1 = 516
        let pages = super::calculate_heap_pages(0, 0x33000, 32, 10).unwrap();
        assert_eq!(pages, 516);
    }
}