zisk-executor 1.1.0-alpha

//! The `ZiskExecutor` module serves as the core orchestrator for executing the ZisK ROM program
//! and generating witness computations.
//!
//! ## Executor Workflow
//! The execution is divided into distinct, sequential phases:
//!
//! 1. **Minimal Traces**: Rapidly process the ROM to collect minimal traces with minimal overhead.
//! 2. **Counting**: Creates the metrics required for the secondary state machine instances.
//! 3. **Planning**: Strategically plan the execution of instances to optimize resource usage.
//! 4. **Instance Creation**: Creates the AIR instances for the main and secondary state machines.
//! 5. **Witness Computation**: Compute the witnesses for all AIR instances, leveraging parallelism
//!    for efficiency.
//!
//! By structuring these phases, the `ZiskExecutor` ensures high-performance execution while
//! maintaining clarity and modularity in the computation process.

use crate::{
    ports::ProofRegistry, witness::WitnessContext, AirClassifier, AsmResources, EmulatorAsm,
    ExecutionPhase, ExecutionState, InstanceAssigner, NoopProofRegistry, PlanPhase,
    ProofmanAdapter, StaticSMBundle, WitnessPhase,
};
use proofman_common::{lease_pool, BufferPool, ProofCtx, ProofmanError, ProofmanResult, SetupCtx};
use proofman_fields::PrimeField64;
use proofman_util::{timer_start_info, timer_stop_and_log_info};
use proofman_witness::{WitnessComponent, WitnessManager};

use std::{
    sync::{Arc, RwLock},
    time::Instant,
};
use zisk_common::{
    io::ZiskStdin, stats_begin, stats_end, AirInstanceCount, BusDeviceMetrics, ChunkId, EmuTrace,
    ExecutorStatsHandle, Plan, ZiskExecutorSummary, ZiskExecutorTime,
};
use zisk_core::{ZiskRom, CHUNK_SIZE};
use zisk_sm_main::{MainPlanner, MainSM};

use crate::error::{ExecutorError, ExecutorResult, RwLockExt};

/// `(chunk_id, metrics)` pair — the per-chunk device-metrics output
/// produced by counter-phase processing.
pub(crate) type DeviceMetricsByChunk = (ChunkId, Box<dyn BusDeviceMetrics>);

/// One entry in the standalone plan summary — counts of planned instances
/// per AIR. No proving-key / setup data; just shape from the planner.
pub struct PlanSummaryEntry {
    /// AIR group id.
    pub airgroup_id: usize,
    /// AIR id within the group.
    pub air_id: usize,
    /// Display name for this AIR (e.g. "Main", "Mem", "Keccakf"). "Unknown" for unregistered ids.
    pub name: &'static str,
    /// Number of instances planned for this AIR.
    pub count: usize,
}

/// The maximum number of steps to execute in the emulator or assembly runner.
pub(crate) const MAX_NUM_STEPS: u64 = 1 << 36;

/// Appends to the progressive minimal-trace store every chunk it has not seen
/// yet, up to and including `idx`. `traces` is the emulator's own buffer, so
/// `traces[idx]` is its last element.
///
/// In-order delivery is an invariant of the ASM reader; a chunk that does not
/// extend the store contiguously is rejected rather than silently skipped.
fn publish_chunks(
    store: &mut Vec<Arc<EmuTrace>>,
    traces: &[Arc<EmuTrace>],
    idx: usize,
) -> ExecutorResult<()> {
    if store.len() > idx {
        return Err(ExecutorError::ChunkOutOfOrder { got: idx, expected: store.len() });
    }
    store.extend_from_slice(&traces[store.len()..=idx]);
    Ok(())
}

/// The `ZiskExecutor` struct orchestrates the execution of the ZisK ROM program, managing state
/// machines, planning, and witness computation.
pub struct ZiskExecutor<F: PrimeField64> {
    /// Shared execution state.
    state: ExecutionState<F>,
    /// Phase-1 Execution. Runs the chosen emulator and produces an `ExecutionOutput`.
    execution: ExecutionPhase,
    /// Phase-2 Plan (pure planning, no bundle).
    plan: PlanPhase<F>,
    /// Phase-3 Witness computation. `None` on the standalone path
    /// (executor constructed without `WitnessManager` / `Std`).
    witness: Option<WitnessPhase<F>>,
}

impl<F: PrimeField64> ZiskExecutor<F> {
    /// Creates a new instance of the `ZiskExecutor` with default state machines.
    ///
    /// This function initializes the executor with a default set of state machines.
    ///
    /// # Arguments
    ///
    /// * `wcm` - Witness manager for managing witness data.
    /// * `verbose_mode` - Verbose mode for logging.
    /// * `shared_tables` - Whether to use shared tables for execution.
    /// * `with_asm_emulator` - Whether the executor supports the ASM backend at runtime.
    /// * `packed` - Whether to use packed representation for witness computation. For Main
    ///   this selects the compact indexed row (+ instruction table).
    pub fn new(
        wcm: &WitnessManager<F>,
        verbose_mode: proofman_common::VerboseMode,
        shared_tables: bool,
        with_asm_emulator: bool,
        packed: bool,
    ) -> ExecutorResult<Arc<Self>> {
        let rank_info = wcm.get_rank_info();
        proofman_common::initialize_logger(verbose_mode, Some(&rank_info));

        let std = pil2_std_lib::Std::new(wcm.get_pctx(), wcm.get_sctx(), shared_tables)?;
        proofman::register_std(wcm, &std);

        let precompiles = crate::Precompiles::all(std.clone());
        let sm_bundle = Arc::new(StaticSMBundle::new(std, precompiles));

        let executor = Arc::new(Self {
            state: ExecutionState::new(),
            execution: ExecutionPhase::new(CHUNK_SIZE, with_asm_emulator),
            plan: PlanPhase::new(CHUNK_SIZE),
            witness: Some(WitnessPhase::new(CHUNK_SIZE, sm_bundle)),
        });
        executor.set_packed(packed);

        wcm.register_component(executor.clone());
        wcm.set_witness_initialized();

        Ok(executor)
    }

    /// Constructs a standalone executor — no `WitnessManager`, no `Std`,
    /// no `StaticSMBundle`, no `WitnessPhase`. Only the emulate + plan
    /// path is wired up; calls to witness-mode-only public methods (e.g.
    /// `calculate_witness`) will panic.
    pub fn new_standalone(
        verbose_mode: proofman_common::VerboseMode,
        with_asm_emulator: bool,
    ) -> ExecutorResult<Arc<Self>> {
        proofman_common::initialize_logger(verbose_mode, None);
        Ok(Arc::new(Self {
            state: ExecutionState::new(),
            execution: ExecutionPhase::new(CHUNK_SIZE, with_asm_emulator),
            plan: PlanPhase::new(CHUNK_SIZE),
            witness: None,
        }))
    }

    /// Standalone execution entry point: emulate + count + plan. Returns
    /// the executor summary, the program's captured `(index, value)`
    /// public-output pairs, and a per-AIR plan summary. `cost_per_type`
    /// on the returned summary is `Default::default()` since cost
    /// computation is skipped (no `SetupCtx`).
    #[allow(clippy::type_complexity)]
    pub fn execute_standalone(
        &self,
        zisk_rom: Arc<ZiskRom>,
        stdin: ZiskStdin,
        use_hints: bool,
    ) -> ExecutorResult<(ZiskExecutorSummary, Vec<(u64, u32)>, Vec<PlanSummaryEntry>)> {
        self.state.set_rom(zisk_rom, use_hints);
        self.state.set_stdin(stdin);
        let registry = NoopProofRegistry::default();
        let global_ids = RwLock::new(Vec::new());
        self.execute_inner(&registry, None, &global_ids)?;

        let mut plan: Vec<PlanSummaryEntry> = registry
            .take_instance_counts()
            .into_iter()
            .map(|((airgroup_id, air_id), count)| PlanSummaryEntry {
                airgroup_id,
                air_id,
                name: AirClassifier::name(airgroup_id, air_id),
                count,
            })
            .collect();
        plan.sort_by_key(|e| (e.airgroup_id, e.air_id));

        Ok((self.state.get_execution_result(), registry.take_pub_outs(), plan))
    }

    /// Sets the ZisK ROM (ELF) for execution.
    ///
    /// This method allows changing the ROM between executions without
    /// recreating the executor, making the executor more reusable.
    ///
    /// # Arguments
    /// * `zisk_rom` - The ZisK ROM to execute.
    pub fn set_rom(&self, zisk_rom: Arc<ZiskRom>, use_hints: bool) -> ExecutorResult<()> {
        self.state.set_rom(zisk_rom.clone(), use_hints);
        if let Some(witness) = self.witness.as_ref() {
            witness.set_rom(zisk_rom)?;
        }
        Ok(())
    }

    /// Sets whether to use packed representation for witness computation.
    pub fn set_packed(&self, packed: bool) {
        if let Some(witness) = self.witness.as_ref() {
            witness.set_packed(packed);
        }
    }

    /// Whether the Main trace is built in the compact indexed form (i.e. packed).
    pub fn is_packed(&self) -> bool {
        self.witness.as_ref().map(|w| w.is_packed()).unwrap_or(false)
    }

    /// Sets the standard input for execution.
    pub fn set_stdin(&self, stdin: ZiskStdin) -> ExecutorResult<()> {
        self.state.set_stdin(stdin);
        Ok(())
    }

    /// Sets ASM resources for execution (only applicable for ASM emulator).
    pub fn set_asm_resources(&self, asm_resources: Arc<AsmResources>) -> ExecutorResult<()> {
        self.execution.set_asm_resources(asm_resources)
    }

    /// Clears any previously-installed ASM resources. No-op when the
    /// executor was built with the Rust emulator backend.
    pub fn clear_asm_resources(&self) -> ExecutorResult<()> {
        self.execution.clear_asm_resources();
        Ok(())
    }

    /// Returns a reference to the ASM emulator if ASM execution is active.
    pub fn asm_emulator(&self) -> Option<&EmulatorAsm> {
        self.execution.asm_emulator()
    }

    /// Gets the execution result and stats.
    #[allow(clippy::type_complexity)]
    pub fn get_execution_result(&self) -> (ZiskExecutorSummary, ExecutorStatsHandle) {
        (self.state.get_execution_result(), self.state.get_stats())
    }

    /// Stores statistics to persistent storage.
    pub fn store_stats(&self) {
        self.state.stats.store_stats();
    }

    /// Inner implementation of [`WitnessComponent::execute`].
    ///
    /// Returns [`ExecutorResult`] so the body can use `?` freely; the
    /// trait-method wrapper maps any error to `ProofmanError::InvalidSetup`
    /// once at the FFI seam.
    fn execute_inner(
        &self,
        registry: &dyn ProofRegistry,
        proofman_extras: Option<&ProofmanAdapter<'_, F>>,
        global_ids: &RwLock<Vec<usize>>,
    ) -> ExecutorResult<()> {
        let start_total = Instant::now();
        self.state.reset();
        if let Some(witness) = self.witness.as_ref() {
            witness.reset()?;
        }

        stats_begin!(self.state.stats, 0, _exec_scope, "EXECUTE", 0);
        self.state.stats.set_start_time(Instant::now());

        let is_asm_emulator = self.execution.is_asm_execution();

        // Reserve proofman's unified GPU buffer for MO count-and-plan
        // (no-op on CPU / standalone).
        if is_asm_emulator {
            if let Some(extras) = proofman_extras {
                extras.acquire_gpu_buffer();
            }
        }

        // ────────────────────────────────────────────────────────────
        // Phase 1.1: Emulate (+ incremental Main advancement)
        // ────────────────────────────────────────────────────────────
        // ROM instance is assigned BEFORE the run so the global-id sequence
        // (ROM, Main segments in order, secondary) is identical to the old
        // batch path while Main segments are now assigned mid-emulation.
        InstanceAssigner::assign_rom_instance(registry)?;

        // Incremental Main-instance advancement. The ASM MT reader calls this on
        // this thread, in chunk order: as soon as a chunk completes a Main
        // instance — or ends the execution — that instance's chunks are published
        // to `state.min_traces` and the instance is planned, assigned and marked
        // witness-ready, so main witness computation (and its GPU streaming-slot
        // commit) overlaps the rest of the emulation instead of waiting for the
        // run to finish. Global-id order is unchanged: ROM (assigned above), then
        // Main segments in order, then secondary.
        //
        // No-op in standalone mode and never called by the Rust emulator; both
        // keep the post-run batch path below.
        let num_within = MainPlanner::traces_per_segment(self.plan.chunk_size())?;
        let on_chunk =
            |idx: usize, traces: &[Arc<EmuTrace>], is_last: bool| -> ExecutorResult<()> {
                let Some(witness) = self.witness.as_ref() else { return Ok(()) };

                // This chunk neither completes a Main instance nor ends the execution.
                let Some(segment) = MainPlanner::segment_completed_by(idx, num_within, is_last)
                else {
                    return Ok(());
                };

                {
                    let mut guard = self.state.min_traces.write_or_poison("min_traces")?;
                    publish_chunks(guard.get_or_insert_with(Vec::new), traces, idx)?;
                }

                let plan = MainPlanner::plan_segment(segment, is_last);
                let assignments =
                    InstanceAssigner::assign_main_instances(registry, global_ids, vec![plan])?;
                witness.populate_main_instances(registry, &self.state, assignments)
            };
        let chunk_hook = &on_chunk;

        timer_start_info!(COMPUTE_MINIMAL_TRACE);
        let start_partial = Instant::now();

        let zisk_rom = self.state.get_rom()?;
        let stdin = self.state.get_stdin();
        let output = self.execution.run::<F>(
            &zisk_rom,
            &stdin,
            registry.is_first_process(),
            self.state.use_hints.load(std::sync::atomic::Ordering::SeqCst),
            &self.state.stats,
            &_exec_scope,
            chunk_hook,
        )?;

        let execution_duration = start_partial.elapsed();
        timer_stop_and_log_info!(COMPUTE_MINIMAL_TRACE);

        // ────────────────────────────────────────────────────────────
        // Phase 1.2: Plan + assign main, then populate main (witness only)
        // ────────────────────────────────────────────────────────────
        let steps = output.steps;

        let crate::ExecutionOutput { min_traces, mut counters, pub_outs, mut backend, .. } = output;
        let num_chunks = min_traces.len();

        // The hook published every chunk it saw, so on the ASM path the store is
        // already complete and a read lock is enough (an exclusive lock here would
        // stall main witnesses that are already computing). The Rust emulator
        // never calls the hook, so its store is still empty and gets the whole
        // vector at once.
        let published =
            self.state.min_traces.read_or_poison("min_traces")?.as_ref().map_or(0, Vec::len);
        if published != num_chunks {
            *self.state.min_traces.write_or_poison("min_traces")? = Some(min_traces);
        }

        // ASM + witness: the hook already released every Main instance during the
        // run (the runner only returns `Ok` after delivering the final chunk).
        let main_instances_count = if is_asm_emulator && self.witness.is_some() {
            num_chunks.div_ceil(num_within)
        } else {
            // Rust emulator / standalone: plan and release every segment now.
            let main_plans = self.plan.run_main(num_chunks, &self.state.stats, &_exec_scope)?;
            let main_assignments =
                InstanceAssigner::assign_main_instances(registry, global_ids, main_plans)?;
            let count = main_assignments.len();
            if let Some(witness) = self.witness.as_ref() {
                witness.populate_main_instances(registry, &self.state, main_assignments)?;
            }
            count
        };

        // ────────────────────────────────────────────────────────────
        // Phase 1.3: Plan secondary, await async, configure + populate (witness only)
        // ────────────────────────────────────────────────────────────
        let secn_artifacts = self.plan.run_secondary(
            &mut counters,
            num_chunks,
            is_asm_emulator,
            &mut backend,
            &self.state.stats,
            &_exec_scope,
        )?;

        // MO runner joined in `run_secondary`; release the buffer back to proofman.
        // Earlier error paths skip the release on purpose: the MO thread may
        // still be using the buffer.
        if is_asm_emulator {
            if let Some(extras) = proofman_extras {
                if let Some(used) = secn_artifacts.gpu_mops_used_bytes {
                    extras.pctx().report_first_gpu_buffer_usage(used);
                }
                extras.release_gpu_buffer();
            }
        }

        timer_start_info!(WAIT_ASM_RH);
        if let Some(rh_data) = backend.await_rom_histogram()? {
            if let Some(witness) = self.witness.as_ref() {
                witness.set_rh_data(rh_data)?;
            }
        }
        timer_stop_and_log_info!(WAIT_ASM_RH);

        stats_begin!(self.state.stats, &_exec_scope, _config_scope, "CONFIGURE_INSTANCES", 0);

        if let (Some(witness), Some(extras)) = (self.witness.as_ref(), proofman_extras) {
            witness.configure_sm_instances(extras.pctx(), &secn_artifacts.secn_planning);
        }

        let mut secn_plans: Vec<Plan> =
            secn_artifacts.secn_planning.into_values().flatten().collect();
        InstanceAssigner::assign_secn_instances(registry, global_ids, &mut secn_plans)?;
        let secn_global_ids: Vec<usize> = secn_plans
            .iter()
            .map(|plan| {
                plan.global_id
                    .ok_or(crate::error::ExecutorError::SecnPlanMissing { phase: "assignment" })
            })
            .collect::<ExecutorResult<Vec<_>>>()?;

        registry.write_pub_outs(&pub_outs.0);

        if let Some(witness) = self.witness.as_ref() {
            witness.populate_secn_instances(&self.state, secn_plans)?;
            witness.configure_checkpoints(registry, &self.state, &secn_global_ids)?;
        }

        stats_end!(self.state.stats, &_config_scope);

        // Reset hints stream and input shmem after the ASM
        // backend-specific await calls have drained the runners.
        self.execution.reset()?;

        // ────────────────────────────────────────────────────────────
        // Phase 1.4: Cost accumulation (witness only — needs sctx)
        // ────────────────────────────────────────────────────────────
        let cost_per_type = match proofman_extras {
            Some(extras) => extras.compute_costs(&self.state, main_instances_count)?,
            None => Default::default(),
        };

        stats_end!(self.state.stats, &_exec_scope);

        let zisk_execution_time = ZiskExecutorTime {
            execution_duration: execution_duration.as_millis() as u64,
            count_and_plan_duration: secn_artifacts.count_and_plan_duration.as_millis() as u64,
            count_and_plan_mo_duration: secn_artifacts.count_and_plan_mo_duration.as_millis()
                as u64,
            total_duration: start_total.elapsed().as_millis() as u64,
            asm_execution_duration: self.execution.get_asm_execution_info()?,
        };
        let mut execution_result =
            ZiskExecutorSummary::new(steps, zisk_execution_time, cost_per_type);
        // Per-AIR instance plan, captured from the registry's planning counts. Only the
        // full (proofman) path exposes this via the summary; the standalone path returns
        // its own (named) plan directly, so skip the work when there's no `SetupCtx`.
        if proofman_extras.is_some() {
            execution_result.plan = registry
                .instance_counts()
                .into_iter()
                .map(|((airgroup_id, air_id), count)| AirInstanceCount {
                    airgroup_id,
                    air_id,
                    count: count as u64,
                })
                .collect();
        }

        // Store the execution result
        self.state.set_execution_result(execution_result);

        Ok(())
    }

    fn witness_or_panic(&self) -> &WitnessPhase<F> {
        self.witness.as_ref().expect("witness phase missing on a witness-mode entry point")
    }

    /// Inner implementation of [`WitnessComponent::calculate_witness`].
    fn calculate_witness_inner(
        &self,
        stage: u32,
        pctx: Arc<ProofCtx<F>>,
        sctx: Arc<SetupCtx<F>>,
        global_ids: &[usize],
        n_cores: usize,
        buffer_pool: &dyn BufferPool<F>,
    ) -> ExecutorResult<()> {
        if stage != 1 {
            return Ok(());
        }

        stats_begin!(self.state.stats, 0, _witness_scope, "CALCULATE_WITNESS", 0);

        let pool = lease_pool(n_cores);
        let adapter = ProofmanAdapter::new(&pctx, &sctx);
        let is_asm_emulator = self.execution.is_asm_execution();
        let witness = self.witness_or_panic();
        pool.install(|| -> ExecutorResult<()> {
            let ctx = WitnessContext::new(
                &pctx,
                &sctx,
                &self.state,
                buffer_pool,
                &_witness_scope,
                &adapter,
                is_asm_emulator,
            );
            for &global_id in global_ids {
                witness.dispatch(&ctx, global_id)?;
            }
            Ok(())
        })?;

        stats_end!(self.state.stats, &_witness_scope);

        Ok(())
    }

    /// Inner implementation of [`WitnessComponent::pre_calculate_witness`].
    fn pre_calculate_witness_inner(
        &self,
        stage: u32,
        pctx: Arc<ProofCtx<F>>,
        sctx: Arc<SetupCtx<F>>,
        global_ids: &[usize],
        n_cores: usize,
        _buffer_pool: &dyn BufferPool<F>,
    ) -> ExecutorResult<()> {
        stats_begin!(self.state.stats, 0, _pre_scope, "PRE_CALCULATE_WITNESS", 0);

        if stage != 1 {
            return Ok(());
        }

        let pool = lease_pool(n_cores);
        let adapter = ProofmanAdapter::new(&pctx, &sctx);
        let is_asm_emulator = self.execution.is_asm_execution();
        let witness = self.witness_or_panic();

        pool.install(|| {
            witness.pre_calculate(&pctx, &adapter, &self.state, global_ids, is_asm_emulator)
        })?;

        stats_end!(self.state.stats, &_pre_scope);
        Ok(())
    }
}

impl<F: PrimeField64> WitnessComponent<F> for ZiskExecutor<F> {
    /// Executes the ZisK ROM program and calculate the plans for main and secondary state machines.
    fn execute(
        &self,
        pctx: Arc<ProofCtx<F>>,
        sctx: Arc<SetupCtx<F>>,
        global_ids: &RwLock<Vec<usize>>,
    ) -> ProofmanResult<()> {
        let adapter = ProofmanAdapter::new(&pctx, &sctx);
        self.execute_inner(&adapter, Some(&adapter), global_ids)
            .map_err(|e| ProofmanError::InvalidSetup(format!("{e:#}")))
    }

    /// Computes the witness for the main and secondary state machines.
    fn calculate_witness(
        &self,
        stage: u32,
        pctx: Arc<ProofCtx<F>>,
        sctx: Arc<SetupCtx<F>>,
        global_ids: &[usize],
        n_cores: usize,
        buffer_pool: &dyn BufferPool<F>,
    ) -> ProofmanResult<()> {
        self.calculate_witness_inner(stage, pctx, sctx, global_ids, n_cores, buffer_pool)
            .map_err(|e| ProofmanError::InvalidSetup(format!("{e:#}")))
    }

    fn pre_calculate_witness(
        &self,
        stage: u32,
        pctx: Arc<ProofCtx<F>>,
        sctx: Arc<SetupCtx<F>>,
        global_ids: &[usize],
        n_cores: usize,
        buffer_pool: &dyn BufferPool<F>,
    ) -> ProofmanResult<()> {
        self.pre_calculate_witness_inner(stage, pctx, sctx, global_ids, n_cores, buffer_pool)
            .map_err(|e| ProofmanError::InvalidSetup(format!("{e:#}")))
    }

    /// Debugs the main and secondary state machines.
    fn debug(
        &self,
        pctx: Arc<ProofCtx<F>>,
        sctx: Arc<SetupCtx<F>>,
        global_ids: &[usize],
    ) -> ProofmanResult<()> {
        for &global_id in global_ids {
            let (_airgroup_id, air_id) = pctx.dctx_get_instance_info(global_id)?;

            if AirClassifier::is_main(air_id) {
                MainSM::debug(&pctx, &sctx);
            } else {
                let secn_instances =
                    self.state.instance_set.secn_instances.read().map_err(|e| {
                        ProofmanError::InvalidSetup(format!("secn_instances lock poisoned: {e}"))
                    })?;
                let secn_instance = secn_instances.get(&global_id).ok_or_else(|| {
                    ProofmanError::InvalidSetup(format!(
                        "Instance not found for global_id {global_id}"
                    ))
                })?;

                secn_instance.debug(&pctx, &sctx);
            }
        }
        Ok(())
    }
}

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

    /// `n` distinct chunks — `steps` is used only as an identity marker here.
    fn chunks(n: usize) -> Vec<Arc<EmuTrace>> {
        (0..n).map(|i| Arc::new(EmuTrace { steps: i as u64 + 1, ..EmuTrace::default() })).collect()
    }

    #[test]
    fn publish_chunks_fills_an_empty_store_up_to_idx() {
        // First release of a 4-chunk segment: the store catches up in one call.
        let traces = chunks(4);
        let mut store = Vec::new();
        publish_chunks(&mut store, &traces, 3).expect("contiguous");
        assert_eq!(store.len(), 4);
        assert!(
            store.iter().zip(&traces).all(|(a, b)| Arc::ptr_eq(a, b)),
            "chunks are published in stream order, not cloned or reordered"
        );
    }

    #[test]
    fn publish_chunks_appends_only_the_unseen_tail() {
        // Segment 1 of a num_within = 2 run: chunks 0..=1 are already published.
        let traces = chunks(4);
        let mut store: Vec<Arc<EmuTrace>> = traces[..2].to_vec();
        publish_chunks(&mut store, &traces, 3).expect("contiguous");
        assert_eq!(store.len(), 4, "only chunks 2 and 3 were appended");
        assert!(Arc::ptr_eq(&store[2], &traces[2]));
        assert!(Arc::ptr_eq(&store[3], &traces[3]));
    }

    #[test]
    fn publish_chunks_is_idempotent_per_chunk_index() {
        // Re-delivering an already-published chunk must not duplicate it.
        let traces = chunks(2);
        let mut store = Vec::new();
        publish_chunks(&mut store, &traces, 1).expect("contiguous");
        let err = publish_chunks(&mut store, &traces, 1).expect_err("already published");
        assert!(matches!(err, ExecutorError::ChunkOutOfOrder { got: 1, expected: 2 }));
        assert_eq!(store.len(), 2, "store is left untouched on rejection");
    }

    #[test]
    fn publish_chunks_rejects_a_gap() {
        // A store at 0 asked to publish chunk 1 would skip chunk 0 — reject instead
        // of leaving a hole the Main witness would read as the wrong chunk.
        let traces = chunks(3);
        let mut store: Vec<Arc<EmuTrace>> = traces[..2].to_vec();
        let err = publish_chunks(&mut store, &traces, 0).expect_err("goes backwards");
        assert!(matches!(err, ExecutorError::ChunkOutOfOrder { got: 0, expected: 2 }));
    }

    /// Replays the hook's decisions over a whole run: for each streamed chunk,
    /// publish + release when `segment_completed_by` says so. Returns the
    /// released `(segment, is_last_segment)` pairs and the final store length.
    fn replay(num_chunks: usize, num_within: usize) -> (Vec<(usize, bool)>, usize) {
        let traces = chunks(num_chunks);
        let mut store = Vec::new();
        let mut released = Vec::new();

        for idx in 0..num_chunks {
            let is_last = idx == num_chunks - 1;
            if let Some(segment) = MainPlanner::segment_completed_by(idx, num_within, is_last) {
                publish_chunks(&mut store, &traces, idx).expect("in-order stream");
                released.push((segment, is_last));
            }
        }
        (released, store.len())
    }

    #[test]
    fn replay_releases_every_segment_exactly_once_and_publishes_every_chunk() {
        let num_within = MainPlanner::traces_per_segment(CHUNK_SIZE).expect("valid chunk size");
        assert_eq!(num_within, MainTrace::<()>::NUM_ROWS / CHUNK_SIZE as usize);

        for num_chunks in 1..=(2 * num_within + 1) {
            let (released, published) = replay(num_chunks, num_within);

            let expected: Vec<usize> = (0..num_chunks.div_ceil(num_within)).collect();
            assert_eq!(
                released.iter().map(|(s, _)| *s).collect::<Vec<_>>(),
                expected,
                "segments released in order, once each ({num_chunks} chunks)"
            );
            assert!(
                released.iter().rev().skip(1).all(|(_, is_last)| !is_last),
                "only the final release is flagged as the last segment"
            );
            assert_eq!(
                released.last().map(|(_, is_last)| *is_last),
                Some(true),
                "the final segment is always released, partial or not"
            );
            assert_eq!(
                published, num_chunks,
                "every chunk reaches the store before its segment is released"
            );
        }
    }

    #[test]
    fn replay_publishes_a_segments_chunks_before_releasing_it() {
        // The Main witness for segment `s` reads store[s * num_within ..], so the
        // store must already cover that range at release time.
        let num_within = 4;
        let num_chunks = 10;
        let traces = chunks(num_chunks);
        let mut store = Vec::new();

        for idx in 0..num_chunks {
            let is_last = idx == num_chunks - 1;
            if let Some(segment) = MainPlanner::segment_completed_by(idx, num_within, is_last) {
                publish_chunks(&mut store, &traces, idx).expect("in-order stream");
                let start = segment * num_within;
                let end = if is_last { store.len() } else { start + num_within };
                assert!(store.len() >= end, "segment {segment} released with a short store");
                assert!(start < store.len(), "segment {segment} released empty");
            }
        }
    }
}