use alloc::{boxed::Box, sync::Arc, vec::Vec};
use core::borrow::{Borrow, BorrowMut};

use itertools::Itertools;
use miden_air::{
    AIRS, CoreCols, Felt, MIDEN_AIR_COUNT, MidenAir, StackCols, SystemCols, config, memory,
    trace::{
        DECODER_TRACE_WIDTH, MIN_TRACE_LEN, MainTrace, RANGE_CHECK_TRACE_WIDTH, RowIndex,
        STACK_TRACE_WIDTH, SYS_TRACE_WIDTH, chiplets::bitwise::OP_CYCLE_LEN, decoder::NUM_OP_BITS,
    },
};
use miden_core::{
    ONE, Word, ZERO,
    field::{PrimeCharacteristicRing, batch_inversion_allow_zeros},
    mast::{MastForestId, OpBatch, SparseMastForest},
    operations::opcodes,
    program::{KernelDescriptor, MIN_STACK_DEPTH},
    utils::Idx,
};
use rayon::prelude::*;
use tracing::{info_span, instrument};

use super::{
    chiplets::Chiplets,
    execution_tracer::TraceReplay,
    trace_state::{
        AceReplay, BitwiseOp, BitwiseReplay, CoreTraceFragmentContext, CoreTraceState,
        ExecutionReplay, HasherRequestReplay, KernelReplay, MemoryWritesReplay, RangeCheckerReplay,
        ResolvedBasicBlockGroups, ResolvedHasherOp,
    },
};
use crate::{
    ContextId, ExecutionError,
    continuation_stack::{Continuation, ContinuationStack},
    errors::MapExecErrNoCtx,
    trace::{
        ChipletsLengths, TraceLenSummary, VmTrace, VmWitness,
        chiplets::{Ace, Bitwise, Hasher, KernelRom, Memory},
        parallel::{processor::ReplayProcessor, tracer::CoreTraceGenerationTracer},
        range::RangeChecker,
        utils::RowMajorTraceWriter,
    },
};

/// Per-row payload written by the core tracer (system + decoder + stack).
pub const CORE_TRACE_WIDTH: usize = SYS_TRACE_WIDTH + DECODER_TRACE_WIDTH + STACK_TRACE_WIDTH;

/// Physical row width of the core buffer: the [`CORE_TRACE_WIDTH`] payload plus the two
/// trailing range-checker columns, which together form the per-AIR Core matrix
/// (`NUM_CORE_COLS`) consumed directly by proving. The range columns are filled in-place
/// after padding (see `write_range_into_core`).
pub const CORE_STORAGE_WIDTH: usize = CORE_TRACE_WIDTH + RANGE_CHECK_TRACE_WIDTH;

/// `build_trace()` uses this as a hard cap on trace rows, independent of the memory budget.
///
/// The code checks `core_trace_contexts.len() * fragment_size` before allocation. It checks the
/// same cap again while replaying chiplet activity, and once more against each padded per-AIR
/// height once they're known (see [`validate_heights_within_max_trace_len`]). Row indices are
/// `u32`-backed, so this bound must hold regardless of budget; the actual memory bound is the
/// tiered check in `build_trace_inner` (see [`memory::max_any_height_for_budget`] and
/// [`memory::prover_peak_bytes`]).
pub(crate) const MAX_TRACE_LEN: usize = 1 << 29;

/// Default maximum memory, in bytes, [`build_trace`] assumes when no budget is given explicitly.
/// Set to 64 GiB, comfortably above every workload in-repo and far below what the previous
/// row-count cap admitted; callers that own actual proving policy (e.g. `miden-prover`'s
/// `Prover`) are expected to set their own via [`build_trace_with_budget`].
pub const DEFAULT_MAX_PROVER_MEMORY_BYTES: u64 = 64 << 30;

pub(crate) mod core_trace_fragment;

mod processor;
mod tracer;

#[cfg(test)]
mod tests;

// BUILD TRACE
// ================================================================================================

/// Builds the main trace from the provided trace states in parallel.
///
/// # Example
/// ```
/// use miden_assembly::Assembler;
/// use miden_processor::{DefaultHost, FastProcessor, StackInputs};
///
/// let program = Assembler::default()
///     .assemble_program("prg", "begin push.1 drop end")
///     .unwrap()
///     .unwrap_program();
/// let mut host = DefaultHost::default();
///
/// let execution_witness = FastProcessor::new(StackInputs::default())
///     .execute_for_proving_sync(&program, &mut host)
///     .unwrap();
/// let (vm_witness, _) = execution_witness.into_parts();
/// let trace = miden_processor::trace::build_trace(vm_witness).unwrap();
///
/// assert_eq!(*trace.program_hash(), program.hash());
/// ```
#[instrument(name = "build_trace", skip_all)]
pub fn build_trace(witness: VmWitness) -> Result<VmTrace, ExecutionError> {
    build_trace_inner(witness, None, DEFAULT_MAX_PROVER_MEMORY_BYTES)
}

/// Same as [`build_trace`], but with an explicit memory budget instead of the default.
pub fn build_trace_with_budget(
    witness: VmWitness,
    max_prover_memory_bytes: u64,
) -> Result<VmTrace, ExecutionError> {
    build_trace_inner(witness, None, max_prover_memory_bytes)
}

/// Same as [`build_trace_with_budget`], but with a hasher chiplet that was already built — used
/// by the streaming path, where the hasher builder runs concurrently with program execution
/// (`FastProcessor::execute_and_build_trace_sync`, std-only).
#[cfg(feature = "std")]
pub(crate) fn build_trace_with_prebuilt_hasher(
    witness: VmWitness,
    prebuilt_hasher: Hasher,
    max_prover_memory_bytes: u64,
) -> Result<VmTrace, ExecutionError> {
    build_trace_inner(witness, Some(prebuilt_hasher), max_prover_memory_bytes)
}

fn build_trace_inner(
    witness: VmWitness,
    prebuilt_hasher: Option<Hasher>,
    max_prover_memory_bytes: u64,
) -> Result<VmTrace, ExecutionError> {
    let VmWitness {
        program_info,
        stack_inputs,
        stack_outputs,
        trace,
        precompile_root,
    } = witness;

    let TraceReplay {
        core_trace_contexts,
        mast_forest_store,
        range_checker_replay,
        memory_writes,
        bitwise_replay: bitwise,
        kernel_replay,
        hasher_for_chiplet,
        ace_replay,
        fragment_size,
        max_stack_depth,
    } = trace;

    let pcs_params = config::pcs_params();

    // Tier 1: a permissive per-AIR row cap derived from the cheapest AIR's marginal cost, so it
    // never rejects a shape that the exact budget check below would accept. `MAX_TRACE_LEN` is a
    // hard ceiling independent of the budget (row indices are `u32`-backed).
    let max_trace_len =
        MAX_TRACE_LEN.min(memory::max_any_height_for_budget(max_prover_memory_bytes, &pcs_params));

    // Before any trace generation, check that the core trace buffer `generate_core_trace_row_major`
    // is about to allocate (`core_trace_contexts.len() * fragment_size` rows, before any padding or
    // proving-time blowup) doesn't itself exceed the byte budget. This is deliberately a separate,
    // more permissive cap than `max_trace_len`: that one prices in the full proving pipeline
    // (blowup, quotient, Merkle trees) that this raw buffer hasn't incurred yet, so reusing it here
    // would reject buffer allocations the exact budget check below would happily accept once the
    // real (usually much smaller) padded core height is known.
    //
    // Note that we add 1 to the total core trace rows to account for the additional HALT opcode row
    // that is pushed at the end of the last fragment.
    let max_core_alloc_len = MAX_TRACE_LEN.min(max_core_alloc_rows(max_prover_memory_bytes));
    let total_core_trace_rows = core_trace_contexts
        .len()
        .checked_mul(fragment_size)
        .and_then(|n| n.checked_add(1))
        .ok_or(ExecutionError::TraceLenExceeded(max_core_alloc_len))?;
    if total_core_trace_rows > max_core_alloc_len {
        return Err(ExecutionError::TraceLenExceeded(max_core_alloc_len));
    }

    if core_trace_contexts.is_empty() {
        return Err(ExecutionError::Internal("no trace fragments provided in the trace witness"));
    }

    let chiplets = info_span!("initialize_chiplets").in_scope(|| {
        initialize_chiplets(
            program_info.kernel().clone(),
            &core_trace_contexts,
            memory_writes,
            bitwise,
            kernel_replay,
            hasher_for_chiplet,
            prebuilt_hasher,
            ace_replay,
            &mast_forest_store,
            max_trace_len,
        )
    })?;

    let range_checker = info_span!("initialize_range_checker")
        .in_scope(|| initialize_range_checker(range_checker_replay, &chiplets));

    let mut core_trace_data = info_span!("generate_core_trace").in_scope(|| {
        generate_core_trace_row_major(
            core_trace_contexts,
            program_info.kernel().clone(),
            fragment_size,
            &mast_forest_store,
            max_stack_depth,
        )
    })?;

    let core_trace_len = core_trace_data.len() / CORE_STORAGE_WIDTH;

    // Get the number of rows for the range checker
    let range_table_len = range_checker.get_number_range_checker_rows();

    let core_height = pad_to_trace_length(core_trace_len.max(range_table_len));
    let chiplets_height = pad_to_trace_length(chiplets.trace_len());
    let poseidon2_permutation_trace_len = chiplets.poseidon2_permutation_trace_len();
    let poseidon2_permutation_height = pad_to_trace_length(poseidon2_permutation_trace_len);

    // Exact check against modelled peak prover memory: pad-up can push usage over budget even
    // when the permissive tier-1 row cap above passed.
    debug_assert_eq!(
        AIRS,
        [MidenAir::Core, MidenAir::Chiplets, MidenAir::Poseidon2Permutation],
        "heights below must be listed in AIRS order",
    );
    let heights: [usize; MIDEN_AIR_COUNT] =
        [core_height, chiplets_height, poseidon2_permutation_height];
    validate_heights_within_max_trace_len(&heights)?;
    let estimated_bytes = memory::prover_peak_bytes(&heights, &pcs_params).ok_or(
        ExecutionError::ProverMemoryExceeded {
            estimated_bytes: u64::MAX,
            budget_bytes: max_prover_memory_bytes,
        },
    )?;
    if estimated_bytes > max_prover_memory_bytes {
        return Err(ExecutionError::ProverMemoryExceeded {
            estimated_bytes,
            budget_bytes: max_prover_memory_bytes,
        });
    }

    let trace_len_summary = TraceLenSummary::new_with_padded(
        core_trace_len,
        range_table_len,
        ChipletsLengths::new(&chiplets),
        poseidon2_permutation_trace_len,
        heights,
    );

    // Each segment is built at its own per-AIR height (no cross-padding to the unified max).
    let ((chiplets_trace, poseidon2_permutation_trace), ()) = info_span!("chiplet_traces_core_pad")
        .in_scope(|| {
            rayon::join(
                || chiplets.into_traces(chiplets_height, poseidon2_permutation_height),
                || pad_core_row_major(&mut core_trace_data, core_height),
            )
        });

    // The range checker occupies the two trailing columns of the core buffer.
    info_span!("write_range_checker_columns").in_scope(|| {
        range_checker.write_range_into_core(
            &mut core_trace_data,
            CORE_STORAGE_WIDTH,
            CORE_TRACE_WIDTH,
            CORE_TRACE_WIDTH + 1,
            range_table_len,
            core_height,
        )
    });

    // Create the MainTrace
    let main_trace = {
        let last_program_row = RowIndex::from((core_trace_len as u32).saturating_sub(1));
        MainTrace::from_parts(
            core_trace_data,
            chiplets_trace.trace,
            poseidon2_permutation_trace.trace,
            last_program_row,
        )
    };

    Ok(VmTrace::new_from_parts(
        program_info,
        stack_inputs,
        stack_outputs,
        precompile_root,
        main_trace,
        trace_len_summary,
    ))
}

// HELPERS
// ================================================================================================

/// Pad a logical row count to a valid trace length: next power of two, clamped to `MIN_TRACE_LEN`.
fn pad_to_trace_length(logical_len: usize) -> usize {
    logical_len.next_power_of_two().max(MIN_TRACE_LEN)
}

/// The largest number of core-trace rows that may be allocated for
/// `generate_core_trace_row_major`'s raw buffer while staying within `max_prover_memory_bytes`.
///
/// The buffer is `rows * CORE_STORAGE_WIDTH` [`Felt`]s at its 1x size, with no blowup, quotient, or
/// Merkle-tree overhead yet, so this is priced directly off the buffer's own byte size rather than
/// through [`memory::max_any_height_for_budget`]'s full-pipeline model.
fn max_core_alloc_rows(max_prover_memory_bytes: u64) -> usize {
    let bytes_per_row = (CORE_STORAGE_WIDTH * size_of::<Felt>()) as u64;
    usize::try_from(max_prover_memory_bytes / bytes_per_row).unwrap_or(usize::MAX)
}

/// Rejects any padded per-AIR height above the hard [`MAX_TRACE_LEN`] row cap, independent of the
/// memory budget: row indices are `u32`-backed, so a height must stay well under `2^32` regardless
/// of how permissive the budget is.
fn validate_heights_within_max_trace_len(
    heights: &[usize; MIDEN_AIR_COUNT],
) -> Result<(), ExecutionError> {
    if heights.iter().any(|&height| height > MAX_TRACE_LEN) {
        return Err(ExecutionError::TraceLenExceeded(MAX_TRACE_LEN));
    }
    Ok(())
}

/// Generates row-major core trace in parallel from the provided trace fragment contexts.
fn generate_core_trace_row_major(
    core_trace_contexts: Vec<CoreTraceFragmentContext>,
    kernel: KernelDescriptor,
    fragment_size: usize,
    mast_forest_store: &[Arc<SparseMastForest>],
    max_stack_depth: usize,
) -> Result<Vec<Felt>, ExecutionError> {
    let num_fragments = core_trace_contexts.len();
    let total_allocated_rows = num_fragments * fragment_size;

    let mut core_trace_data = Felt::zero_vec(total_allocated_rows * CORE_STORAGE_WIDTH);

    // Save the first stack top for initialization
    let first_stack_top = if let Some(first_context) = core_trace_contexts.first() {
        first_context.state.stack.stack_top.to_vec()
    } else {
        vec![ZERO; MIN_STACK_DEPTH]
    };

    let writers: Vec<RowMajorTraceWriter<'_, Felt>> = core_trace_data
        .chunks_exact_mut(fragment_size * CORE_STORAGE_WIDTH)
        .map(|chunk| {
            RowMajorTraceWriter::with_stride(chunk, CORE_STORAGE_WIDTH, CORE_STORAGE_WIDTH)
        })
        .collect();

    // Build the core trace fragments in parallel
    let fragment_results: Result<Vec<_>, ExecutionError> = core_trace_contexts
        .into_par_iter()
        .zip(writers.into_par_iter())
        .map(|(trace_state, writer)| {
            let (mut processor, mut tracer, mut continuation_stack, mut current_forest) =
                split_trace_fragment_context(
                    trace_state,
                    writer,
                    fragment_size,
                    mast_forest_store,
                    max_stack_depth,
                )?;

            processor.execute(
                &mut continuation_stack,
                &mut current_forest,
                &kernel,
                &mut tracer,
            )?;

            tracer.into_final_state()
        })
        .collect();
    let fragment_results = fragment_results?;

    let mut stack_rows = Vec::new();
    let mut system_rows = Vec::new();
    let mut total_core_trace_rows = 0;

    for final_state in fragment_results {
        stack_rows.push(final_state.last_stack_cols);
        system_rows.push(final_state.last_system_cols);
        total_core_trace_rows += final_state.num_rows_written;
    }

    // Fix up stack and system rows
    fixup_stack_and_system_rows(
        &mut core_trace_data,
        fragment_size,
        &stack_rows,
        &system_rows,
        &first_stack_top,
    );

    // Run batch inversion on stack's H0 helper column, processing each fragment in parallel.
    // This must be done after fixup_stack_and_system_rows since that function overwrites the first
    // row of each fragment with non-inverted values.
    {
        let w = CORE_STORAGE_WIDTH;
        core_trace_data[..total_core_trace_rows * w]
            .par_chunks_mut(fragment_size * w)
            .for_each(|fragment_chunk| {
                let num_rows = fragment_chunk.len() / w;
                let mut h0_vals: Vec<Felt> = (0..num_rows)
                    .map(|r| {
                        let row: &CoreCols<Felt> = fragment_chunk[r * w..(r + 1) * w].borrow();
                        row.stack.h0
                    })
                    .collect();
                batch_inversion_allow_zeros(&mut h0_vals);
                for (r, &val) in h0_vals.iter().enumerate() {
                    let row: &mut CoreCols<Felt> = fragment_chunk[r * w..(r + 1) * w].borrow_mut();
                    row.stack.h0 = val;
                }
            });
    }

    // Truncate the core trace columns to the actual number of rows written.
    core_trace_data.truncate(total_core_trace_rows * CORE_STORAGE_WIDTH);

    push_halt_opcode_row(
        &mut core_trace_data,
        total_core_trace_rows,
        system_rows
            .last()
            .ok_or(ExecutionError::Internal("no trace fragments provided in the trace witness"))?,
        stack_rows
            .last()
            .ok_or(ExecutionError::Internal("no trace fragments provided in the trace witness"))?,
    );

    Ok(core_trace_data)
}

/// Initializing the first row of each fragment with the appropriate stack and system state.
///
/// This needs to be done as a separate pass after all fragments have been generated, because the
/// system and stack rows write the state at clk `i` to the row at index `i+1`. Hence, the state of
/// the last row of any given fragment cannot be written in parallel, since any given fragment
/// filler doesn't have access to the next fragment's first row.
fn fixup_stack_and_system_rows(
    core_trace_data: &mut [Felt],
    fragment_size: usize,
    stack_rows: &[StackCols<Felt>],
    system_rows: &[SystemCols<Felt>],
    first_stack_top: &[Felt],
) {
    const MIN_STACK_DEPTH_FELT: Felt = Felt::new_unchecked(MIN_STACK_DEPTH as u64);
    let w = CORE_STORAGE_WIDTH;

    {
        let row: &mut CoreCols<Felt> = core_trace_data[..w].borrow_mut();

        // Stack order in the trace is reversed vs `first_stack_top`.
        for (stack_col_idx, &value) in first_stack_top.iter().rev().enumerate() {
            row.stack.top[stack_col_idx] = value;
        }

        row.stack.b0 = MIN_STACK_DEPTH_FELT;
        row.stack.b1 = ZERO;
        row.stack.h0 = ZERO;
    }

    let total_rows = core_trace_data.len() / w;
    let num_fragments = total_rows / fragment_size;

    for frag_idx in 1..num_fragments {
        let row_idx = frag_idx * fragment_size;
        let row_start = row_idx * w;
        let row: &mut CoreCols<Felt> = core_trace_data[row_start..row_start + w].borrow_mut();
        row.system = system_rows[frag_idx - 1].clone();
        row.stack = stack_rows[frag_idx - 1].clone();
    }
}

/// Appends a HALT row (`num_rows_before` is the row count before append).
///
/// This ensures that the trace ends with at least one HALT operation, which is necessary to satisfy
/// the constraints.
fn push_halt_opcode_row(
    core_trace_data: &mut Vec<Felt>,
    num_rows_before: usize,
    last_system_state: &SystemCols<Felt>,
    last_stack_state: &StackCols<Felt>,
) {
    let w = CORE_STORAGE_WIDTH;
    let mut row_data = [ZERO; CORE_STORAGE_WIDTH];

    // Read the previous row's hasher state first half before we take a mutable borrow on
    // `row_data` (propagates the program hash into the HALT padding).
    let prev_hasher_state_first_half: [Felt; 4] = if num_rows_before > 0 {
        let last_row_start = (num_rows_before - 1) * w;
        let prev: &CoreCols<Felt> = core_trace_data[last_row_start..last_row_start + w].borrow();
        let hs = &prev.decoder.hasher_state;
        [hs[0], hs[1], hs[2], hs[3]]
    } else {
        [ZERO; 4]
    };

    {
        let row: &mut CoreCols<Felt> = row_data.as_mut_slice().borrow_mut();

        row.system = last_system_state.clone();
        row.stack = last_stack_state.clone();

        // Pad op_bits columns with HALT opcode bits
        let halt_opcode = opcodes::HALT;
        for bit_idx in 0..NUM_OP_BITS {
            row.decoder.op_bits[bit_idx] = Felt::from_u8((halt_opcode >> bit_idx) & 1);
        }

        // Pad hasher state columns (8 columns)
        // - First 4 columns: copy the last value (to propagate program hash)
        // - Remaining 4 columns: fill with ZEROs
        row.decoder.hasher_state[..4].copy_from_slice(&prev_hasher_state_first_half);

        // Pad op_bit_extra columns (2 columns)
        // - First column: do nothing (pre-filled with ZEROs, HALT doesn't use this)
        // - Second column: fill with ONEs (product of two most significant HALT bits, both are 1)
        row.decoder.extra[1] = ONE;
    }

    core_trace_data.extend_from_slice(&row_data);
}

/// Initializes the ranger checker from the recorded range checks during execution and returns it.
///
/// Note that the maximum number of rows that the range checker can produce is 2^16, which is less
/// than the maximum trace length (2^29). Hence, we can safely generate the entire range checker
/// trace and then pad it to the final trace length, without worrying about hitting memory limits.
fn initialize_range_checker(
    range_checker_replay: RangeCheckerReplay,
    chiplets: &Chiplets,
) -> RangeChecker {
    let mut range_checker = RangeChecker::new();

    // Add all range checks recorded during execution.
    for values in range_checker_replay {
        range_checker.add_range_checks(values.as_ref());
    }

    // Add all hasher- and memory-related range checks.
    chiplets.append_range_checks(&mut range_checker);

    range_checker
}

/// Replays recorded operations to populate chiplet traces. Results were already used during
/// execution; this pass only needs the trace-recording side effects.
///
/// The five chiplets are populated from disjoint replays, so they build in parallel. Their
/// non-hasher lengths are known from the replay metadata; checking those up front and giving the
/// hasher only the remaining rows preserves the hard cap before any builder materializes its
/// trace on the buffered path. A prebuilt (streamed) hasher was already built during execution
/// under the full `max_trace_len` budget and is instead validated against the remaining rows
/// after the fact.
fn initialize_chiplets(
    kernel: KernelDescriptor,
    core_trace_contexts: &[CoreTraceFragmentContext],
    memory_writes: MemoryWritesReplay,
    bitwise: BitwiseReplay,
    kernel_replay: KernelReplay,
    hasher_for_chiplet: HasherRequestReplay,
    prebuilt_hasher: Option<Hasher>,
    ace_replay: AceReplay,
    mast_forest_store: &[Arc<SparseMastForest>],
    max_trace_len: usize,
) -> Result<Chiplets, ExecutionError> {
    let non_hasher_trace_len = non_hasher_trace_len(
        &kernel,
        core_trace_contexts,
        &memory_writes,
        &bitwise,
        &ace_replay,
        max_trace_len,
    )?;
    let max_hasher_trace_len = max_trace_len
        .checked_sub(non_hasher_trace_len)
        .ok_or(ExecutionError::TraceLenExceeded(max_trace_len))?;

    if prebuilt_hasher
        .as_ref()
        .is_some_and(|hasher| hasher.trace_len() > max_hasher_trace_len)
    {
        return Err(ExecutionError::TraceLenExceeded(max_trace_len));
    }

    let (hasher, (bitwise, (memory, (ace, kernel_rom)))) = rayon::join(
        || match prebuilt_hasher {
            Some(hasher) => Ok(hasher),
            None => build_hasher_chiplet(
                hasher_for_chiplet.into_resolved_ops(mast_forest_store),
                max_hasher_trace_len,
            )
            .map_err(|err| match err {
                // The builder reports its internal remainder budget; surface the
                // configured cap instead, like every other rejection site.
                ExecutionError::TraceLenExceeded(_) => {
                    ExecutionError::TraceLenExceeded(max_trace_len)
                },
                other => other,
            }),
        },
        || {
            rayon::join(
                || build_bitwise_chiplet(bitwise, max_trace_len),
                || {
                    rayon::join(
                        || build_memory_chiplet(memory_writes, core_trace_contexts, max_trace_len),
                        || {
                            rayon::join(
                                || build_ace_chiplet(ace_replay, max_trace_len),
                                || build_kernel_rom_chiplet(kernel, kernel_replay, max_trace_len),
                            )
                        },
                    )
                },
            )
        },
    );

    let chiplets = Chiplets {
        hasher: hasher?,
        bitwise: bitwise?,
        memory: memory?,
        ace: ace?,
        kernel_rom: kernel_rom?,
    };
    debug_assert_eq!(
        non_hasher_trace_len,
        chiplets.trace_len() - chiplets.hasher.trace_len(),
        "chiplet preflight length differs from the materialized trace",
    );
    // Release-only insurance: in debug builds a preflight undercount trips the
    // assert above before this check can fire.
    if chiplets.trace_len() > max_trace_len {
        return Err(ExecutionError::TraceLenExceeded(max_trace_len));
    }
    Ok(chiplets)
}

fn non_hasher_trace_len(
    kernel: &KernelDescriptor,
    core_trace_contexts: &[CoreTraceFragmentContext],
    memory_writes: &MemoryWritesReplay,
    bitwise: &BitwiseReplay,
    ace: &AceReplay,
    max_trace_len: usize,
) -> Result<usize, ExecutionError> {
    let overflow = || ExecutionError::TraceLenExceeded(max_trace_len);
    let bitwise_len = bitwise.num_operations().checked_mul(OP_CYCLE_LEN).ok_or_else(overflow)?;
    let memory_reads_len = core_trace_contexts.iter().try_fold(0usize, |len, context| {
        len.checked_add(context.replay.memory_reads.num_accesses()?)
    });
    let memory_len = memory_writes
        .num_accesses()
        .and_then(|writes| memory_reads_len.and_then(|reads| writes.checked_add(reads)))
        .ok_or_else(overflow)?;
    let ace_len = ace.trace_len().ok_or_else(overflow)?;

    [1, kernel.proc_hashes().len(), bitwise_len, memory_len, ace_len]
        .into_iter()
        .try_fold(0usize, usize::checked_add)
        .filter(|&total| total <= max_trace_len)
        .ok_or_else(overflow)
}

/// Builds the hasher chiplet by replaying resolved requests in order.
///
/// The iterator abstracts over the two delivery modes: the buffered replay drained against the
/// finalized forest store, or a live channel fed by a concurrently executing processor (see
/// `FastProcessor::execute_and_build_trace_sync`).
pub(crate) fn build_hasher_chiplet<'a>(
    ops: impl IntoIterator<Item = Result<ResolvedHasherOp<'a>, ExecutionError>>,
    max_trace_len: usize,
) -> Result<Hasher, ExecutionError> {
    let mut hasher = Hasher::default();
    for hasher_op in ops {
        match hasher_op? {
            ResolvedHasherOp::Permute(input_state) => {
                let _ = hasher.permute(input_state);
            },
            ResolvedHasherOp::HashControlBlock((h1, h2, domain, expected_hash)) => {
                let _ = hasher.hash_control_block(h1, h2, domain, expected_hash);
            },
            ResolvedHasherOp::HashBasicBlock((batch_groups, expected_hash)) => match batch_groups {
                ResolvedBasicBlockGroups::Borrowed(op_batches) => {
                    let _ = hasher
                        .hash_basic_block(op_batches.iter().map(OpBatch::groups), expected_hash);
                },
                ResolvedBasicBlockGroups::Owned(batch_groups) => {
                    let _ = hasher.hash_basic_block(batch_groups.iter(), expected_hash);
                },
            },
            ResolvedHasherOp::BuildMerkleRoot((value, path, index)) => {
                let _ = hasher.build_merkle_root(value, &path, index);
            },
            ResolvedHasherOp::UpdateMerkleRoot((old_value, new_value, path, index)) => {
                hasher.update_merkle_root(old_value, new_value, &path, index);
            },
        }
        if hasher.trace_len() > max_trace_len {
            return Err(ExecutionError::TraceLenExceeded(max_trace_len));
        }
    }
    Ok(hasher)
}

/// Builds the bitwise chiplet by replaying recorded `u32and`/`u32xor` requests in order.
fn build_bitwise_chiplet(
    bitwise_replay: BitwiseReplay,
    max_trace_len: usize,
) -> Result<Bitwise, ExecutionError> {
    let mut bitwise = Bitwise::default();
    for (bitwise_op, a, b) in bitwise_replay {
        match bitwise_op {
            BitwiseOp::U32And => {
                bitwise.u32and(a, b).map_exec_err_no_ctx()?;
            },
            BitwiseOp::U32Xor => {
                bitwise.u32xor(a, b).map_exec_err_no_ctx()?;
            },
        }
        if bitwise.trace_len() > max_trace_len {
            return Err(ExecutionError::TraceLenExceeded(max_trace_len));
        }
    }
    Ok(bitwise)
}

/// Builds the memory chiplet by replaying recorded accesses merged in clock-cycle order.
fn build_memory_chiplet(
    memory_writes: MemoryWritesReplay,
    core_trace_contexts: &[CoreTraceFragmentContext],
    max_trace_len: usize,
) -> Result<Memory, ExecutionError> {
    enum MemoryAccess {
        ReadElement(Felt, ContextId, RowIndex),
        WriteElement(Felt, Felt, ContextId, RowIndex),
        ReadWord(Felt, ContextId, RowIndex),
        WriteWord(Felt, Word, ContextId, RowIndex),
    }

    impl MemoryAccess {
        fn clk(&self) -> RowIndex {
            match self {
                MemoryAccess::ReadElement(_, _, clk) => *clk,
                MemoryAccess::WriteElement(_, _, _, clk) => *clk,
                MemoryAccess::ReadWord(_, _, clk) => *clk,
                MemoryAccess::WriteWord(_, _, _, clk) => *clk,
            }
        }
    }

    let mut memory = Memory::default();

    // Note: care is taken to order all the accesses by clock cycle, since the memory chiplet
    // currently assumes that all memory accesses are issued in the same order as they appear in
    // the trace.
    let elements_written: Box<dyn Iterator<Item = MemoryAccess>> =
        Box::new(memory_writes.iter_elements_written().map(|(element, addr, ctx, clk)| {
            MemoryAccess::WriteElement(*addr, *element, *ctx, *clk)
        }));
    let words_written: Box<dyn Iterator<Item = MemoryAccess>> = Box::new(
        memory_writes
            .iter_words_written()
            .map(|(word, addr, ctx, clk)| MemoryAccess::WriteWord(*addr, *word, *ctx, *clk)),
    );
    let elements_read: Box<dyn Iterator<Item = MemoryAccess>> =
        Box::new(core_trace_contexts.iter().flat_map(|ctx| {
            ctx.replay
                .memory_reads
                .iter_read_elements()
                .map(|(_, addr, ctx, clk)| MemoryAccess::ReadElement(addr, ctx, clk))
        }));
    let words_read: Box<dyn Iterator<Item = MemoryAccess>> =
        Box::new(core_trace_contexts.iter().flat_map(|ctx| {
            ctx.replay
                .memory_reads
                .iter_read_words()
                .map(|(_, addr, ctx, clk)| MemoryAccess::ReadWord(addr, ctx, clk))
        }));

    [elements_written, words_written, elements_read, words_read]
        .into_iter()
        .kmerge_by(|a, b| a.clk() < b.clk())
        .try_for_each(|mem_access| {
            match mem_access {
                MemoryAccess::ReadElement(addr, ctx, clk) => memory
                    .read(ctx, addr, clk)
                    .map(|_| ())
                    .map_err(ExecutionError::MemoryErrorNoCtx)?,
                MemoryAccess::WriteElement(addr, element, ctx, clk) => memory
                    .write(ctx, addr, clk, element)
                    .map_err(ExecutionError::MemoryErrorNoCtx)?,
                MemoryAccess::ReadWord(addr, ctx, clk) => memory
                    .read_word(ctx, addr, clk)
                    .map(|_| ())
                    .map_err(ExecutionError::MemoryErrorNoCtx)?,
                MemoryAccess::WriteWord(addr, word, ctx, clk) => memory
                    .write_word(ctx, addr, clk, word)
                    .map_err(ExecutionError::MemoryErrorNoCtx)?,
            }
            if memory.trace_len() > max_trace_len {
                return Err(ExecutionError::TraceLenExceeded(max_trace_len));
            }
            Ok(())
        })?;

    Ok(memory)
}

/// Builds the ACE chiplet by replaying recorded circuit evaluations in order.
fn build_ace_chiplet(ace_replay: AceReplay, max_trace_len: usize) -> Result<Ace, ExecutionError> {
    let mut ace = Ace::default();
    for (clk, circuit_eval) in ace_replay.into_iter() {
        ace.add_circuit_evaluation(clk, circuit_eval);
        if ace.trace_len() > max_trace_len {
            return Err(ExecutionError::TraceLenExceeded(max_trace_len));
        }
    }
    Ok(ace)
}

/// Builds the kernel ROM chiplet by replaying recorded kernel procedure accesses in order.
fn build_kernel_rom_chiplet(
    kernel: KernelDescriptor,
    kernel_replay: KernelReplay,
    max_trace_len: usize,
) -> Result<KernelRom, ExecutionError> {
    let mut kernel_rom = KernelRom::new(kernel);
    for proc_hash in kernel_replay.into_iter() {
        kernel_rom.access_proc(proc_hash).map_exec_err_no_ctx()?;
        if kernel_rom.trace_len() > max_trace_len {
            return Err(ExecutionError::TraceLenExceeded(max_trace_len));
        }
    }
    Ok(kernel_rom)
}

/// Pads the core trace to `core_height` rows (HALT template, CLK incremented per row).
fn pad_core_row_major(core_trace_data: &mut Vec<Felt>, core_height: usize) {
    let w = CORE_STORAGE_WIDTH;
    let total_program_rows = core_trace_data.len() / w;
    assert!(total_program_rows <= core_height);
    assert!(total_program_rows > 0);

    let num_padding_rows = core_height - total_program_rows;
    if num_padding_rows == 0 {
        return;
    }
    let last_row_start = (total_program_rows - 1) * w;

    // Safety: per our documented safety guarantees, we know that `total_program_rows > 0`,
    // and row `total_program_rows - 1` is initialized.
    let (last_hasher_first_half, last_stack): ([Felt; 4], StackCols<Felt>) = {
        let last: &CoreCols<Felt> = core_trace_data[last_row_start..last_row_start + w].borrow();
        let hs = &last.decoder.hasher_state;
        let last_hasher: [Felt; 4] = [hs[0], hs[1], hs[2], hs[3]];
        (last_hasher, last.stack.clone())
    };

    let mut template_data = [ZERO; CORE_STORAGE_WIDTH];
    {
        let template: &mut CoreCols<Felt> = template_data.as_mut_slice().borrow_mut();

        // Decoder columns
        // ------------------------

        // Pad op_bits columns with HALT opcode bits
        let halt_opcode = opcodes::HALT;
        for i in 0..NUM_OP_BITS {
            template.decoder.op_bits[i] = Felt::from_u8((halt_opcode >> i) & 1);
        }
        // Pad hasher state columns (8 columns)
        // - First 4 columns: copy the last value (to propagate program hash)
        // - Remaining 4 columns: fill with ZEROs
        template.decoder.hasher_state[..4].copy_from_slice(&last_hasher_first_half);

        // Pad op_bit_extra columns (2 columns)
        // - First column: do nothing (filled with ZEROs, HALT doesn't use this)
        // - Second column: fill with ONEs (product of two most significant HALT bits, both are 1)
        template.decoder.extra[1] = ONE;

        // Stack columns
        // ------------------------

        // Pad stack columns with the last value in each column (analogous to Stack::into_trace())
        template.stack = last_stack;
    }

    // System columns
    // ------------------------

    // Pad CLK trace - fill with index values

    let pad_start = total_program_rows * w;
    core_trace_data.resize(pad_start + num_padding_rows * w, ZERO);
    core_trace_data[pad_start..]
        .par_chunks_mut(w)
        .enumerate()
        .for_each(|(idx, row_buf)| {
            row_buf.copy_from_slice(&template_data);
            let row: &mut CoreCols<Felt> = row_buf.borrow_mut();
            row.system.clk = Felt::from_u32((total_program_rows + idx) as u32);
        });
}

type SplitFragmentContext<'a> = (
    ReplayProcessor,
    CoreTraceGenerationTracer<'a>,
    ContinuationStack<Arc<SparseMastForest>>,
    Arc<SparseMastForest>,
);

/// Uses the provided `CoreTraceFragmentContext` to build and return a `ReplayProcessor` and
/// `CoreTraceGenerationTracer` that can be used to execute the fragment.
///
/// `mast_forest_store` provides the [`SparseMastForest`]s that the indices stored in the fragment
/// (the initial forest index and the `EnterForest` continuations) refer to.
///
/// # Errors
///
/// Returns [`ExecutionError::Internal`] if any [`MastForestId`] referenced by the fragment
/// (either `initial_mast_forest_id` or an `EnterForest` continuation) is out of range of
/// `mast_forest_store`. Because [`CoreTraceFragmentContext`] is attacker-controllable when fed in
/// from outside, we validate these indices rather than indexing-and-panicking.
fn split_trace_fragment_context<'a>(
    fragment_context: CoreTraceFragmentContext,
    writer: RowMajorTraceWriter<'a, Felt>,
    fragment_size: usize,
    mast_forest_store: &[Arc<SparseMastForest>],
    max_stack_depth: usize,
) -> Result<SplitFragmentContext<'a>, ExecutionError> {
    let CoreTraceFragmentContext {
        state: CoreTraceState { system, decoder, stack },
        replay:
            ExecutionReplay {
                block_stack: block_stack_replay,
                execution_context: execution_context_replay,
                stack_overflow: stack_overflow_replay,
                memory_reads: memory_reads_replay,
                advice: advice_replay,
                hasher: hasher_response_replay,
                block_address: block_address_replay,
                mast_forest_resolution: mast_forest_resolution_replay,
            },
        continuation,
        initial_mast_forest_id,
    } = fragment_context;

    let translated_continuation =
        translate_snapshot_continuation_stack(continuation, mast_forest_store)?;

    let initial_mast_forest =
        lookup_mast_forest(mast_forest_store, initial_mast_forest_id)?.clone();

    let processor = ReplayProcessor::new(
        system,
        stack,
        stack_overflow_replay,
        execution_context_replay,
        advice_replay,
        memory_reads_replay,
        hasher_response_replay,
        mast_forest_resolution_replay,
        mast_forest_store.to_vec(),
        max_stack_depth,
        fragment_size.into(),
    );
    let tracer =
        CoreTraceGenerationTracer::new(writer, decoder, block_address_replay, block_stack_replay);

    Ok((processor, tracer, translated_continuation, initial_mast_forest))
}

/// Translates a snapshotted `ContinuationStack<MastForestId>` into one carrying actual
/// [`Arc<SparseMastForest>`] handles, ready to drive `execute_impl`.
///
/// Returns [`ExecutionError::Internal`] if any `EnterForest` continuation carries a
/// [`MastForestId`] that is out of range of `mast_forest_store`.
fn translate_snapshot_continuation_stack(
    snapshot: ContinuationStack<MastForestId>,
    mast_forest_store: &[Arc<SparseMastForest>],
) -> Result<ContinuationStack<Arc<SparseMastForest>>, ExecutionError> {
    let mut out: ContinuationStack<Arc<SparseMastForest>> = ContinuationStack::default();
    for cont in snapshot.into_inner() {
        let translated = match cont {
            Continuation::EnterForest {
                forest: id,
                package_debug_info,
                inline_context_depth,
            } => Continuation::EnterForest {
                forest: lookup_mast_forest(mast_forest_store, id)?.clone(),
                package_debug_info,
                inline_context_depth,
            },
            Continuation::StartNode(id) => Continuation::StartNode(id),
            Continuation::FinishJoin(id) => Continuation::FinishJoin(id),
            Continuation::FinishSplit(id) => Continuation::FinishSplit(id),
            Continuation::FinishLoop(node_id) => Continuation::FinishLoop(node_id),
            Continuation::FinishCall(id) => Continuation::FinishCall(id),
            Continuation::FinishDyn(id) => Continuation::FinishDyn(id),
            Continuation::ResumeBasicBlock { node_id, batch_index, op_idx_in_batch } => {
                Continuation::ResumeBasicBlock { node_id, batch_index, op_idx_in_batch }
            },
            Continuation::Respan { node_id, batch_index } => {
                Continuation::Respan { node_id, batch_index }
            },
            Continuation::FinishBasicBlock(id) => Continuation::FinishBasicBlock(id),
        };
        out.push_continuation(translated);
    }
    Ok(out)
}

/// Looks up `id` in `mast_forest_store`, returning [`ExecutionError::Internal`] if it is out of
/// range.
pub(super) fn lookup_mast_forest(
    mast_forest_store: &[Arc<SparseMastForest>],
    id: MastForestId,
) -> Result<&Arc<SparseMastForest>, ExecutionError> {
    mast_forest_store
        .get(id.to_usize())
        .ok_or(ExecutionError::Internal("MastForestId out of range of mast_forest_store"))
}