proofman-common 1.3.1-alpha

use std::collections::HashSet;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
use crate::{ProofmanResult, ProofmanError};

/// Lifecycle of an instance's witness within one proof. Single answer to "does this still need
/// computing?": the dispatch-time bool it replaces could disagree with the trace it stood for.
#[repr(u8)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum WitnessState {
    /// Nothing has happened to it in this proof.
    Absent = 0,
    /// On the witness channel, waiting for a hook to pick it up.
    Queued = 1,
    /// A witness hook owns it right now.
    Running = 2,
    /// Its witness exists.
    Done = 3,
    /// Was `Done`; its trace was reclaimed, so it must be recomputed before it can be proved.
    Evicted = 4,
}

impl WitnessState {
    fn from_u8(v: u8) -> Self {
        match v {
            0 => Self::Absent,
            1 => Self::Queued,
            2 => Self::Running,
            3 => Self::Done,
            _ => Self::Evicted,
        }
    }
}

/// One instance's witness lifecycle. Every transition goes through `move_to`, so the legal ones are
/// the only ones expressible and each is a single atomic.
#[derive(Debug, Default)]
pub struct WitnessSlot(AtomicU8);

impl WitnessSlot {
    pub fn get(&self) -> WitnessState {
        WitnessState::from_u8(self.0.load(Ordering::SeqCst))
    }

    fn move_to(&self, to: WitnessState, from: &[WitnessState]) -> bool {
        let mut cur = self.0.load(Ordering::SeqCst);
        loop {
            if !from.contains(&WitnessState::from_u8(cur)) {
                return false;
            }
            match self.0.compare_exchange_weak(cur, to as u8, Ordering::SeqCst, Ordering::SeqCst) {
                Ok(_) => return true,
                Err(actual) => cur = actual,
            }
        }
    }

    /// False means it is queued, owned, or already computed: do not enqueue it again. `Evicted` is
    /// excluded on purpose -- it has had its witness in this phase, and a blanket re-announce would
    /// put it back on the channel to be recomputed for nobody. `rearm` is the explicit way back.
    pub fn try_queue(&self) -> bool {
        self.move_to(WitnessState::Queued, &[WitnessState::Absent])
    }

    /// The one gate every witness hook goes through. `Evicted` is not a source: its owner may still
    /// be running, and a second hook taking the slot would make that owner's `release` land on the
    /// second one's `Running`. `rearm` is the way back, and every recompute path already calls it.
    pub fn try_acquire(&self) -> bool {
        self.move_to(WitnessState::Running, &[WitnessState::Absent, WitnessState::Queued])
    }

    /// End of a hook. Producing nothing returns to `Absent` because the queue slot is spent; only
    /// from `Running`, so a trace reclaimed mid-hook is not overwritten with `Done`.
    pub fn release(&self, produced: bool) {
        let to = if produced { WitnessState::Done } else { WitnessState::Absent };
        self.move_to(to, &[WitnessState::Running]);
    }

    /// The trace is gone. Covers `Running` so a buffer freed under a live hook cannot end as `Done`.
    pub fn mark_evicted(&self) {
        self.move_to(WitnessState::Evicted, &[WitnessState::Done, WitnessState::Running]);
    }

    /// Back to square one: this instance is wanted again and must be prepared, queued and computed
    /// from scratch. Used by the on-demand queries and by the phase boundary.
    pub fn rearm(&self) -> bool {
        self.move_to(WitnessState::Absent, &[WitnessState::Evicted, WitnessState::Queued])
    }
}

/// Witness-lifecycle diagnostics. Pure observation: nothing reads these to make a decision.
#[derive(Debug, Default)]
pub struct WitnessStats {
    /// Dispatches refused because the instance was already queued or computed.
    pub duplicate_sends: AtomicUsize,
    /// Traces reclaimed from computed instances; each is a witness that must be redone.
    pub evictions: AtomicUsize,
    /// Witnesses computed over a trace that already held one.
    pub recomputed_over_trace: AtomicUsize,
}

impl WitnessStats {
    pub fn snapshot(&self) -> (usize, usize, usize) {
        let load = |c: &AtomicUsize| c.load(Ordering::Relaxed);
        (load(&self.duplicate_sends), load(&self.evictions), load(&self.recomputed_over_trace))
    }
}

#[cfg(test)]
mod witness_slot_tests {
    use super::*;
    use std::sync::Arc;

    #[test]
    fn a_second_queue_of_the_same_instance_is_refused() {
        let s = WitnessSlot::default();
        assert!(s.try_queue());
        assert!(!s.try_queue(), "the duplicate must not reach the channel");
    }

    #[test]
    fn a_queued_instance_is_still_acquirable_and_only_once() {
        let s = WitnessSlot::default();
        assert!(s.try_queue());
        assert!(s.try_acquire());
        assert!(!s.try_acquire());
        s.release(true);
        assert_eq!(s.get(), WitnessState::Done);
        assert!(!s.try_acquire(), "a produced witness must not be computed again");
    }

    #[test]
    fn producing_nothing_leaves_the_instance_computable_and_queueable() {
        let s = WitnessSlot::default();
        assert!(s.try_queue());
        assert!(s.try_acquire());
        s.release(false);
        assert_eq!(s.get(), WitnessState::Absent);
        assert!(s.try_queue());
        assert!(s.try_acquire());
    }

    #[test]
    fn reclaiming_a_trace_makes_the_instance_computable_again_once_rearmed() {
        let s = WitnessSlot::default();
        assert!(s.try_acquire());
        s.release(true);
        s.mark_evicted();
        assert_eq!(s.get(), WitnessState::Evicted);
        assert!(!s.try_acquire(), "its owner may still be running");
        assert!(s.rearm());
        assert!(s.try_acquire());
    }

    /// A reclaim mid-hook must not let a second hook in: the first owner's `release` would then
    /// land on the second one's `Running` and leave that one's result untracked.
    #[test]
    fn an_evicted_owner_is_not_displaced_by_a_second_hook() {
        let s = WitnessSlot::default();
        assert!(s.try_acquire());
        s.mark_evicted();
        assert!(!s.try_acquire(), "the evicted owner has not finished");
        s.release(true);
        assert_eq!(s.get(), WitnessState::Evicted, "the late release must not publish Done");
    }

    #[test]
    fn an_evicted_instance_is_not_re_queued_by_a_blanket_announce() {
        let s = WitnessSlot::default();
        assert!(s.try_acquire());
        s.release(true);
        s.mark_evicted();
        assert!(!s.try_queue(), "a blanket re-announce must not schedule it again");
        assert!(s.rearm());
        assert_eq!(s.get(), WitnessState::Absent);
        assert!(s.try_queue());
    }

    #[test]
    fn rearm_leaves_a_computed_instance_alone() {
        let s = WitnessSlot::default();
        assert!(s.try_acquire());
        s.release(true);
        assert!(!s.rearm());
        assert_eq!(s.get(), WitnessState::Done);
    }

    #[test]
    fn a_reclaim_mid_hook_beats_the_late_release() {
        let s = WitnessSlot::default();
        assert!(s.try_acquire());
        s.mark_evicted();
        s.release(true);
        assert_eq!(s.get(), WitnessState::Evicted);
        assert!(s.rearm());
        assert!(s.try_acquire());
    }

    #[test]
    fn reclaiming_an_instance_that_never_ran_is_a_no_op() {
        let s = WitnessSlot::default();
        s.mark_evicted();
        assert_eq!(s.get(), WitnessState::Absent);
        assert!(s.try_queue());
        s.mark_evicted();
        assert_eq!(s.get(), WitnessState::Queued);
    }

    /// Both gates must elect exactly one winner under contention.
    fn race(f: fn(&WitnessSlot) -> bool) -> usize {
        let s = Arc::new(WitnessSlot::default());
        let winners = Arc::new(AtomicUsize::new(0));
        let handles: Vec<_> = (0..16)
            .map(|_| {
                let (s, winners) = (s.clone(), winners.clone());
                std::thread::spawn(move || {
                    if f(&s) {
                        winners.fetch_add(1, Ordering::SeqCst);
                    }
                })
            })
            .collect();
        handles.into_iter().for_each(|h| h.join().unwrap());
        winners.load(Ordering::SeqCst)
    }

    #[test]
    fn concurrent_gates_elect_exactly_one_winner() {
        assert_eq!(race(WitnessSlot::try_acquire), 1);
        assert_eq!(race(WitnessSlot::try_queue), 1);
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct InstanceChunks {
    pub chunks: Vec<usize>,
    pub slow: bool,
}

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct InstanceInfo {
    pub airgroup_id: usize,
    pub air_id: usize,
    pub table: bool,
    pub shared: bool,
    pub n_chunks: usize,
    pub weight: u64,            // Cost of the basic proof plus the recursion chain it triggers
    pub compressor_weight: u64, // Cost of the compressor proof it triggers, 0 if the air has none
}

impl InstanceInfo {
    pub fn new(
        airgroup_id: usize,
        air_id: usize,
        table: bool,
        shared: bool,
        weight: u64,
        compressor_weight: u64,
    ) -> Self {
        Self { airgroup_id, air_id, table, shared, n_chunks: 0, weight, compressor_weight }
    }

    /// Total cost this instance puts on its owner: basic, recursion chain and compressor
    #[inline]
    pub fn total_weight(&self) -> u64 {
        self.weight + self.compressor_weight
    }

    /// Whether this instance triggers a compressor proof, i.e. its air has a compressor and the
    /// compressor setups are loaded (they are not without aggregation, and then there is nothing
    /// to spread across partitions either)
    #[inline]
    pub fn has_compressor(&self) -> bool {
        self.compressor_weight > 0
    }
}

/// Context for distributed computing with two-level hierarchy:
/// 1. Distributed Prover MANAGER distributes PARTITIONS across WORKERS
/// 2. Instances are distributed across PARTITIONS as they are created
/// 3. Instances are distributed across local PROCESSES within each WORKER
#[derive(Default)]
pub struct DistributionCtx {
    // Worker-level
    pub n_partitions: usize,       // Total number of partitions in the system
    pub partition_mask: Vec<bool>, // Which partitions are assigned to this worker

    // Process-level
    pub n_processes: usize, // Number of processes used by this worker
    pub process_id: usize,  // ID of current process within this worker

    // DYNAMIC PARAMETERS

    // Instances
    pub n_instances: usize,                               // Total number of instances
    pub instances: Vec<InstanceInfo>,                     // Instances info
    pub instances_chunks: Vec<InstanceChunks>,            // Chunks info per instance
    pub witness_states: Vec<WitnessSlot>,                 // Witness lifecycle per instance
    pub n_tables: usize,                                  // Number of table instances
    pub aux_tables: Vec<InstanceInfo>,                    // Table instances info (lately appended to instances)
    pub aux_table_map: Vec<i32>,                          // Map from aux tables to original instances
    pub table_assignment: HashMap<(usize, usize), usize>, // (airgroup_id, air_id) -> reference gid the table is forced onto

    // Worker-level distribution
    pub partition_set: bool,                      // Whether the partition assignation is done
    pub instance_partition: Vec<i32>, // Which partition each instance belongs to (>=0 assigned, -1 unassigned, -2 appended table)
    pub assigned_table_instances: HashSet<usize>, // Global indices of assigned tables (forced onto a reference instance's node)
    pub worker_instances: Vec<usize>,             // Indexes of instances assigned to this worker
    pub partition_count: Vec<u32>,                // #instances in each partition (does not include tables)
    pub partition_weight: Vec<u64>,               // Weight per partition, basic + compressor (does not include tables)
    pub partition_compressor_count: Vec<u32>,     // #compressor instances assigned to each partition

    // Process-level distribution
    pub instance_process: Vec<(i32, usize)>, // For each instance: (process_id or -1 if other worker, local_idx)
    pub process_instances: Vec<usize>,       // Indexes of instances assigned to current process
    pub skipped_process_instances: Vec<usize>, // Indexes of instances assigned to current process but skipped for some reason
    pub process_count: Vec<usize>,             // #instances assigned to each process
    pub process_weight: Vec<u64>,              // Weight per process, basic + recursion + compressor
    pub process_compressor_count: Vec<u32>,    // #compressor instances assigned to each process

    pub worker_index: i32, // Index of the current worker

    // Control
    pub assignation_done: bool, // Whether the instance assignation is done
}

impl std::fmt::Debug for DistributionCtx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut dbg = f.debug_struct("DistributionCtx");

        // STATIC PARAMETERS
        dbg.field("=== STATIC PARAMS ===", &"");
        dbg.field("n_partitions", &self.n_partitions)
            .field("partition_mask", &self.partition_mask)
            .field("n_processes", &self.n_processes)
            .field("process_id", &self.process_id);

        // DYNAMIC PARAMETERS
        dbg.field("=== DYNAMIC PARAMS ===", &"");
        dbg.field("n_instances", &self.n_instances)
            .field("instances", &self.instances)
            .field("n_tables", &self.n_tables)
            .field("tables", &self.aux_tables)
            .field("instance_partition", &self.instance_partition)
            .field("worker_instances", &self.worker_instances)
            .field("partition_count", &self.partition_count)
            .field("partition_weight", &self.partition_weight)
            .field("partition_compressor_count", &self.partition_compressor_count)
            .field("instance_process", &self.instance_process)
            .field("process_instances", &self.process_instances)
            .field("process_count", &self.process_count)
            .field("process_weight", &self.process_weight)
            .field("process_compressor_count", &self.process_compressor_count)
            .field("assignation_done", &self.assignation_done);
        dbg.finish()
    }
}

impl DistributionCtx {
    pub fn new() -> Self {
        DistributionCtx {
            n_partitions: 0,
            partition_mask: Vec::new(),
            n_processes: 0,
            process_id: 0,
            n_instances: 0,
            instances: Vec::new(),
            witness_states: Vec::new(),
            instances_chunks: Vec::new(),
            n_tables: 0,
            aux_tables: Vec::new(),
            aux_table_map: Vec::new(),
            table_assignment: HashMap::new(),
            instance_partition: Vec::new(),
            assigned_table_instances: HashSet::new(),
            worker_instances: Vec::new(),
            partition_count: Vec::new(),
            partition_weight: Vec::new(),
            partition_compressor_count: Vec::new(),
            instance_process: Vec::new(),
            process_instances: Vec::new(),
            skipped_process_instances: Vec::new(),
            process_count: Vec::new(),
            process_weight: Vec::new(),
            process_compressor_count: Vec::new(),
            worker_index: -1,
            assignation_done: false,
            partition_set: false,
        }
    }

    /// Configure the static partitioning parameters
    /// - `n_partitions`: Total number of partitions in the distributed system
    /// - `partition_ids`: Which partition IDs are assigned to this worker
    /// - `balance`: Whether to balance load across partitions or use round-robin
    pub fn setup_partitions(&mut self, n_partitions: usize, partition_ids: Vec<u32>) -> ProofmanResult<()> {
        self.n_partitions = n_partitions;
        self.partition_mask = vec![false; n_partitions];

        for id in &partition_ids {
            if *id < n_partitions as u32 {
                self.partition_mask[*id as usize] = true;
            } else {
                return Err(ProofmanError::InvalidConfiguration(format!(
                    "Partition ID {} exceeds total partitions {}",
                    id, n_partitions
                )));
            }
        }

        self.partition_count = vec![0; n_partitions];
        self.partition_weight = vec![0; n_partitions];
        self.partition_compressor_count = vec![0; n_partitions];
        self.partition_set = true;
        Ok(())
    }

    /// Configure the static processes parameters
    /// - `n_processes`: Number of processes available to this worker
    /// - `process_id`: The rank/ID of the current process (must be < n_processes)
    pub fn setup_processes(&mut self, n_processes: usize, process_id: usize) -> ProofmanResult<()> {
        if process_id >= n_processes {
            return Err(ProofmanError::InvalidConfiguration(format!(
                "Process rank {} exceeds total processes {}",
                process_id, n_processes
            )));
        }
        self.n_processes = n_processes;
        self.process_id = process_id;
        self.process_count = vec![0; n_processes];
        self.process_weight = vec![0; n_processes];
        self.process_compressor_count = vec![0; n_processes];
        Ok(())
    }

    pub fn setup_worker_index(&mut self, worker_index: usize) {
        self.worker_index = worker_index as i32;
    }

    /// Reset all DYNAMIC parameters for a new proof
    /// This clears all instance-specific data while preserving the STATIC configuration
    pub fn reset_instances(&mut self) {
        // Global instance management
        self.n_instances = 0;
        self.instances.clear();
        self.instances_chunks.clear();
        self.witness_states.clear();
        self.n_tables = 0;
        self.aux_tables.clear();
        self.aux_table_map.clear();
        self.table_assignment.clear();
        self.assigned_table_instances.clear();

        // Worker-level
        self.instance_partition.clear();
        self.worker_instances.clear();
        self.partition_count.fill(0);
        self.partition_weight.fill(0);
        self.partition_compressor_count.fill(0);

        // Process-level
        self.instance_process.clear();
        self.process_instances.clear();
        self.skipped_process_instances.clear();
        self.process_count.fill(0);
        self.process_weight.fill(0);
        self.process_compressor_count.fill(0);

        //control
        self.assignation_done = false;
        self.partition_set = false;
    }

    /// Verify that the static configuration has been properly set up
    #[inline]
    pub fn validate_static_config(&self) -> Result<(), ProofmanError> {
        // Check partition configuration
        if self.n_partitions == 0 {
            return Err(ProofmanError::InvalidConfiguration(
                "Partition configuration not set. Call setup_partitions() first.".to_string(),
            ));
        }
        if self.partition_mask.len() != self.n_partitions {
            return Err(ProofmanError::InvalidConfiguration(
                "Partition mask size mismatch with n_partitions".to_string(),
            ));
        }

        // Check process configuration
        if self.n_processes == 0 {
            return Err(ProofmanError::InvalidConfiguration(
                "Process configuration not set. Call setup_processes() first.".to_string(),
            ));
        }
        if self.process_id >= self.n_processes {
            return Err(ProofmanError::InvalidConfiguration(format!(
                "Invalid process rank {} >= {}",
                self.process_id, self.n_processes
            )));
        }

        Ok(())
    }

    #[inline]
    pub fn is_setup_partition_init(&self) -> bool {
        self.partition_set
    }

    pub fn skip_instance(&mut self, instance_id: usize) {
        if !self.skipped_process_instances.contains(&instance_id) {
            self.skipped_process_instances.push(instance_id);
        }
    }

    pub fn is_skipped_instance(&self, instance_id: usize) -> bool {
        self.skipped_process_instances.contains(&instance_id)
    }

    /// Check if the current process is the owner of a given instance
    #[inline]
    pub fn is_my_process_instance(&self, instance_id: usize) -> ProofmanResult<bool> {
        if instance_id >= self.instance_process.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instance_process.len()
            )));
        }
        Ok(self.instance_process[instance_id].0 == self.process_id as i32)
    }

    #[inline]
    pub fn get_process_owner_instance(&self, instance_id: usize) -> ProofmanResult<i32> {
        if instance_id >= self.instance_process.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instance_process.len()
            )));
        }
        let owner = self.instance_process[instance_id].0;
        if owner == -1 {
            return Err(ProofmanError::InvalidAssignation(format!(
                "Instance {} is not owned by any process",
                instance_id
            )));
        }
        Ok(owner)
    }

    /// Get the airgroup and air ID for a given instance
    /// Returns (airgroup_id, air_id)
    #[inline]
    pub fn get_instance_info(&self, instance_id: usize) -> ProofmanResult<(usize, usize)> {
        if instance_id >= self.instances.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instances.len()
            )));
        }
        Ok((self.instances[instance_id].airgroup_id, self.instances[instance_id].air_id))
    }

    /// Get the airgroup and air ID of a given table
    /// Returns (airgroup_id, air_id)
    #[inline]
    pub fn get_table_info(&self, table_idx: usize) -> ProofmanResult<(usize, usize)> {
        if self.assignation_done {
            let instance_id = self.aux_table_map[table_idx] as usize;
            self.get_instance_info(instance_id)
        } else {
            if table_idx >= self.aux_tables.len() {
                return Err(ProofmanError::OutOfBounds(format!(
                    "Table index {} out of bounds (max: {})",
                    table_idx,
                    self.aux_tables.len()
                )));
            }
            Ok((self.aux_tables[table_idx].airgroup_id, self.aux_tables[table_idx].air_id))
        }
    }

    pub fn get_table_instance_idx(&self, table_idx: usize) -> ProofmanResult<usize> {
        if self.assignation_done {
            Ok(self.aux_table_map[table_idx] as usize)
        } else {
            Err(ProofmanError::InvalidAssignation("Table instances not yet assigned".into()))
        }
    }

    /// Get the local index of the instance within its owner process
    #[inline]
    pub fn get_instance_local_idx(&self, instance_id: usize) -> ProofmanResult<usize> {
        if instance_id >= self.instance_process.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instance_process.len()
            )));
        }
        Ok(self.instance_process[instance_id].1)
    }

    /// Get the number of Minimum Trace chunks to be processes for a given global instance
    #[inline]
    pub fn get_instance_chunks(&self, instance_id: usize) -> ProofmanResult<usize> {
        if instance_id >= self.instances.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instances.len()
            )));
        }
        Ok(self.instances_chunks[instance_id].chunks.len())
    }

    /// Set the number of chunks for a given instance (these may be used for balancing purposes)
    pub fn set_n_chunks(&mut self, instance_id: usize, n_chunks: usize) -> ProofmanResult<()> {
        if instance_id >= self.instances.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instances.len()
            )));
        }
        let instance_info = &mut self.instances[instance_id];
        instance_info.n_chunks = n_chunks;
        Ok(())
    }

    /// Check if the current worker is the owner of a given instance
    #[inline]
    pub fn is_my_worker_instance(&self, instance_id: usize) -> ProofmanResult<bool> {
        if instance_id >= self.instance_process.len() {
            return Err(ProofmanError::OutOfBounds(format!(
                "Instance index {} out of bounds (max: {})",
                instance_id,
                self.instance_process.len()
            )));
        }
        Ok(self.instance_process[instance_id].0 >= 0)
    }

    #[inline]
    pub fn find_air_instance_id(&self, instance_id: usize) -> ProofmanResult<usize> {
        let mut air_instance_id = 0;
        let (airgroup_id, air_id) = self.get_instance_info(instance_id)?;
        for idx in 0..instance_id {
            let (instance_airgroup_id, instance_air_id) = self.get_instance_info(idx)?;
            if (instance_airgroup_id, instance_air_id) == (airgroup_id, air_id) {
                air_instance_id += 1;
            }
        }
        Ok(air_instance_id)
    }

    /// Find an instance with the given airgroup_id and air_id among the current process's instances
    /// Returns (found, local_index)
    #[inline]
    pub fn find_process_instance(&self, airgroup_id: usize, air_id: usize) -> ProofmanResult<(bool, usize)> {
        let mut matches = self
            .process_instances
            .iter()
            .enumerate()
            .filter(|&(_pos, &id)| {
                let inst = &self.instances[id];
                inst.airgroup_id == airgroup_id && inst.air_id == air_id
            })
            .map(|(pos, _)| pos);

        match (matches.next(), matches.next()) {
            (None, _) => Ok((false, 0)),
            (Some(pos), None) => Ok((true, pos)),
            (Some(_), Some(_)) => Err(ProofmanError::InvalidAssignation(format!(
                "Multiple instances found for airgroup_id: {airgroup_id}, air_id: {air_id}"
            ))),
        }
    }

    #[inline]
    pub fn find_instance_id(&self, airgroup_id: usize, air_id: usize) -> ProofmanResult<(bool, usize)> {
        let mut matches = self
            .instances
            .iter()
            .enumerate()
            .filter(|&(_gid, inst)| inst.airgroup_id == airgroup_id && inst.air_id == air_id)
            .map(|(gid, _)| gid);

        match (matches.next(), matches.next()) {
            (None, _) => Ok((false, 0)),
            (Some(gid), None) => Ok((true, gid)),
            (Some(_), Some(_)) => Err(ProofmanError::InvalidAssignation(format!(
                "Multiple instances found for airgroup_id: {airgroup_id}, air_id: {air_id}"
            ))),
        }
    }

    /// Find a table with the given airgroup_id and air_id among the current process's tables
    /// Returns (found, local_index)
    #[inline]
    pub fn find_process_table(&self, airgroup_id: usize, air_id: usize) -> ProofmanResult<(bool, usize)> {
        if self.assignation_done {
            self.find_process_instance(airgroup_id, air_id)
        } else {
            let mut matches = self
                .aux_tables
                .iter()
                .enumerate()
                .filter(|&(_pos, &info)| info.airgroup_id == airgroup_id && info.air_id == air_id)
                .map(|(pos, _)| pos);

            match (matches.next(), matches.next()) {
                (None, _) => Ok((false, 0)),
                (Some(pos), None) => Ok((true, pos)),
                (Some(_), Some(_)) => Err(ProofmanError::InvalidAssignation(format!(
                    "Multiple tables found for airgroup_id: {airgroup_id}, air_id: {air_id}"
                ))),
            }
        }
    }

    #[inline]
    pub fn find_by_air_instance_id(&self, airgroup_id: usize, air_id: usize, air_instance_id: usize) -> Option<usize> {
        let mut count = 0;
        for (instance_idx, instance) in self.instances.iter().enumerate() {
            let (inst_airgroup_id, inst_air_id) = (instance.airgroup_id, instance.air_id);
            if airgroup_id == inst_airgroup_id && air_id == inst_air_id {
                if count == air_instance_id {
                    return Some(instance_idx);
                }
                count += 1;
            }
        }
        None
    }

    /// Partition with the least accumulated cost, ties broken by #instances. Compressor count is
    /// not a criterion of its own: that cost is already in the weight.
    #[inline]
    fn least_loaded_partition(&self) -> usize {
        let mut best_idx = 0;
        let mut best_key = (u64::MAX, u32::MAX);
        for (i, &weight) in self.partition_weight.iter().enumerate() {
            let key = (weight, self.partition_count[i]);
            if key < best_key {
                best_key = key;
                best_idx = i;
            }
        }
        best_idx
    }

    /// Same criterion as `least_loaded_partition`, over the processes of this worker
    #[inline]
    fn least_loaded_process(&self, counts: &[usize]) -> usize {
        let mut best_idx = 0;
        let mut best_key = (u64::MAX, usize::MAX);
        for (i, &weight) in self.process_weight.iter().enumerate() {
            let key = (weight, counts[i]);
            if key < best_key {
                best_key = key;
                best_idx = i;
            }
        }
        best_idx
    }

    /// add an instance and assign it immediately to the least loaded partition/process
    /// the instance added is not a table
    #[inline]
    pub fn add_instance(
        &mut self,
        airgroup_id: usize,
        air_id: usize,
        weight: u64,
        compressor_weight: u64,
    ) -> ProofmanResult<usize> {
        if self.assignation_done {
            return Err(ProofmanError::InvalidAssignation("Instances already assigned".to_string()));
        }
        self.validate_static_config().expect("Static configuration invalid or incomplete");
        let gid: usize = self.instances.len();
        let instance = InstanceInfo::new(airgroup_id, air_id, false, false, weight, compressor_weight);
        let (total_weight, has_compressor) = (instance.total_weight(), instance.has_compressor());
        self.instances.push(instance);
        self.instances_chunks.push(InstanceChunks { chunks: vec![], slow: false });
        self.witness_states.push(WitnessSlot::default());
        self.n_instances += 1;
        // Placed as created, without knowing the instances still to come: greedy least loaded.
        // Round-robin on the gid handed out the heaviest airs (Main, Keccakf) blindly, leaving
        // only the instances of assign_instances() to compensate for it.
        let partition_id = self.least_loaded_partition() as u32;
        self.instance_partition.push(partition_id as i32);
        self.partition_count[partition_id as usize] += 1;
        self.partition_weight[partition_id as usize] += total_weight;
        if has_compressor {
            self.partition_compressor_count[partition_id as usize] += 1;
        }
        let mut local_idx = 0;
        let mut owner = -1;
        if self.partition_mask[partition_id as usize] {
            self.worker_instances.push(gid);
            let process_id = self.least_loaded_process(&self.process_count);
            owner = process_id as i32;
            local_idx = self.process_count[process_id];
            self.process_count[process_id] += 1;
            self.process_weight[process_id] += total_weight;
            if has_compressor {
                self.process_compressor_count[process_id] += 1;
            }
            if process_id == self.process_id {
                self.process_instances.push(gid);
            }
        }
        self.instance_process.push((owner, local_idx));
        Ok(gid)
    }

    pub fn is_first_process(&self) -> bool {
        self.partition_mask[0] && self.process_id == 0
    }

    /// add an instance without assigning it to any partition/process
    /// It will be assigned later by assign_instances()
    /// the instance added is not a table
    #[inline]
    pub fn add_instance_no_assign(
        &mut self,
        airgroup_id: usize,
        air_id: usize,
        weight: u64,
        compressor_weight: u64,
    ) -> ProofmanResult<usize> {
        if self.assignation_done {
            return Err(ProofmanError::InvalidAssignation("Instances already assigned".to_string()));
        }
        self.validate_static_config().expect("Static configuration invalid or incomplete");
        self.instances.push(InstanceInfo::new(airgroup_id, air_id, false, false, weight, compressor_weight));
        self.instances_chunks.push(InstanceChunks { chunks: vec![], slow: false });
        self.witness_states.push(WitnessSlot::default());
        self.instance_partition.push(-1);
        self.instance_process.push((-1, 0_usize));
        self.n_instances += 1;
        Ok(self.n_instances - 1)
    }

    /// Add local table instances
    pub fn add_table(&mut self, airgroup_id: usize, air_id: usize, weight: u64) -> ProofmanResult<usize> {
        if self.assignation_done {
            return Err(ProofmanError::InvalidAssignation("Instances already assigned".to_string()));
        }
        self.validate_static_config().expect("Static configuration invalid or incomplete");
        let lid = self.aux_tables.len();
        self.aux_tables.push(InstanceInfo::new(airgroup_id, air_id, true, true, weight, 0));
        self.aux_table_map.push(-1);
        self.n_tables += 1;
        Ok(lid)
    }

    pub fn set_chunks(&mut self, global_idx: usize, chunks: Vec<usize>, slow: bool) {
        let instance_info = &mut self.instances_chunks[global_idx];
        instance_info.chunks = chunks;
        instance_info.slow = slow;
    }

    pub fn add_table_all(&mut self, airgroup_id: usize, air_id: usize, weight: u64) -> ProofmanResult<usize> {
        if self.assignation_done {
            return Err(ProofmanError::InvalidAssignation("Instances already assigned".to_string()));
        }
        self.validate_static_config().expect("Static configuration invalid or incomplete");
        let lid = self.aux_tables.len();
        self.aux_tables.push(InstanceInfo::new(airgroup_id, air_id, true, false, weight, 0));
        self.aux_table_map.push(-1);
        self.n_tables += 1;
        Ok(lid)
    }

    /// Assign instances to partitions and processes
    pub fn assign_instances(&mut self) -> ProofmanResult<()> {
        if self.assignation_done {
            return Err(ProofmanError::InvalidAssignation("Instances already assigned".to_string()));
        }
        //assign instances
        self.validate_static_config().expect("Static configuration invalid or incomplete");

        // Sort the unassigned instances by proof weight
        let mut unassigned_instances = Vec::new();
        for (gid, &partition_id) in self.instance_partition.iter().enumerate() {
            if partition_id == -1 {
                unassigned_instances.push((gid, self.instances[gid].total_weight()));
            }
        }

        // Sort the unassigned instances by proof weight
        unassigned_instances.sort_by_key(|b| std::cmp::Reverse(b.1));

        // Distribute the unassigned instances to the process with minimum weight each time
        // cost: O(n^2) may be optimized if needed

        let mut instances_assigned_partition = vec![HashMap::<(usize, usize), usize>::new(); self.n_partitions];
        let mut instances_assigned_process = vec![HashMap::<(usize, usize), usize>::new(); self.n_processes];

        let mut local_process_count = self.process_count.clone();
        for (gid, _) in &unassigned_instances {
            let (airgroup_id, air_id) = self.get_instance_info(*gid)?;
            let has_compressor = self.instances[*gid].has_compressor();

            // Select target partition: least loaded first (see least_loaded_partition)
            let min_weight_idx = self.least_loaded_partition();
            if has_compressor {
                self.partition_compressor_count[min_weight_idx] += 1;
            }

            *instances_assigned_partition[min_weight_idx].entry((airgroup_id, air_id)).or_insert(0) += 1;
            self.partition_count[min_weight_idx] += 1;
            self.partition_weight[min_weight_idx] += self.instances[*gid].total_weight();
            if self.partition_mask[min_weight_idx] {
                // Select target process with the same criterion as the partition above
                let min_weight_process_idx = self.least_loaded_process(&local_process_count);
                if has_compressor {
                    self.process_compressor_count[min_weight_process_idx] += 1;
                }

                local_process_count[min_weight_process_idx] += 1;
                self.process_weight[min_weight_process_idx] += self.instances[*gid].total_weight();
                instances_assigned_process[min_weight_process_idx]
                    .entry((airgroup_id, air_id))
                    .and_modify(|c| *c += 1)
                    .or_insert(1);
            }
        }

        unassigned_instances.sort_by_key(|&(idx, weight)| (self.instances_chunks[idx].slow, std::cmp::Reverse(weight)));

        let partitions_chunks: &mut Vec<HashSet<usize>> = &mut (0..self.n_partitions).map(|_| HashSet::new()).collect();
        let process_chunks: &mut Vec<HashSet<usize>> = &mut (0..self.n_processes).map(|_| HashSet::new()).collect();

        for (gid, _) in &unassigned_instances {
            let chunks = &self.instances_chunks[*gid].chunks;
            let (airgroup_id, air_id) = self.get_instance_info(*gid)?;
            let mut min_chunks = usize::MAX;
            let mut min_chunks_idx = 0;
            for partition_id in 0..self.n_partitions {
                if instances_assigned_partition[partition_id].get(&(airgroup_id, air_id)).unwrap_or(&0) > &0 {
                    let mut new_chunks_added = 0;
                    for chunk in chunks {
                        if !partitions_chunks[partition_id].contains(chunk) {
                            new_chunks_added += 1;
                        }
                    }
                    if new_chunks_added < min_chunks {
                        min_chunks = new_chunks_added;
                        min_chunks_idx = partition_id;
                    }
                }
            }

            if let Some(c) = instances_assigned_partition[min_chunks_idx].get_mut(&(airgroup_id, air_id)) {
                *c -= 1;
            }

            self.instance_partition[*gid] = min_chunks_idx as i32;

            for chunk in chunks {
                partitions_chunks[min_chunks_idx].insert(*chunk);
            }

            if self.partition_mask[min_chunks_idx] {
                self.worker_instances.push(*gid);
                let mut min_chunks = usize::MAX;
                let mut min_process_id = 0;
                for process_id in 0..self.n_processes {
                    if instances_assigned_process[process_id].get(&(airgroup_id, air_id)).unwrap_or(&0) > &0 {
                        let mut new_chunks_added = 0;
                        for chunk in chunks {
                            if !process_chunks[process_id].contains(chunk) {
                                new_chunks_added += 1;
                            }
                        }
                        if new_chunks_added < min_chunks {
                            min_chunks = new_chunks_added;
                            min_process_id = process_id;
                        }
                    }
                }

                for chunk in chunks {
                    process_chunks[min_process_id].insert(*chunk);
                }

                if min_process_id == self.process_id {
                    self.process_instances.push(*gid);
                }

                self.instance_process[*gid].0 = min_process_id as i32;
                self.instance_process[*gid].1 = self.process_count[min_process_id];
                self.process_count[min_process_id] += 1;
                if let Some(c) = instances_assigned_process[min_process_id].get_mut(&(airgroup_id, air_id)) {
                    *c -= 1;
                }
            }
        }

        // Add tables that
        self.n_tables = 0;
        for (table_idx, table) in self.aux_tables.iter().enumerate() {
            if let Some(&ref_gid) = self.table_assignment.get(&(table.airgroup_id, table.air_id)) {
                // Assigned table: single instance forced onto ref_gid's node.
                if ref_gid >= self.instance_partition.len() {
                    return Err(ProofmanError::InvalidAssignation(format!(
                        "Table assignment reference gid {ref_gid} out of bounds"
                    )));
                }
                let ref_partition = self.instance_partition[ref_gid];
                if ref_partition < 0 {
                    return Err(ProofmanError::InvalidAssignation(format!(
                        "Table assignment reference gid {ref_gid} is not assigned to a partition"
                    )));
                }
                // -1 when ref_partition is another worker's: every worker keeps the gid and the
                // table mapping, only the owner materializes the instance.
                let ref_process = self.instance_process[ref_gid].0;
                let gid = self.instances.len();
                self.instances.push(*table);
                self.witness_states.push(WitnessSlot::default());
                self.instances_chunks.push(InstanceChunks { chunks: vec![], slow: false });
                self.n_instances += 1;
                self.n_tables += 1;
                self.instance_partition.push(ref_partition); // REAL partition, not -2
                if ref_process >= 0 {
                    let ref_process = ref_process as usize;
                    self.worker_instances.push(gid);
                    let lid = self.process_count[ref_process];
                    self.process_count[ref_process] += 1;
                    self.process_weight[ref_process] += table.weight;
                    if ref_process == self.process_id {
                        self.process_instances.push(gid);
                    }
                    self.instance_process.push((ref_process as i32, lid));
                } else {
                    self.instance_process.push((-1, 0_usize));
                }
                self.aux_table_map[table_idx] = gid as i32;
                self.assigned_table_instances.insert(gid);
            } else if table.shared {
                let process_id = self.least_loaded_process(&self.process_count);
                let gid = self.instances.len();
                self.instances.push(*table);
                self.witness_states.push(WitnessSlot::default());
                self.instances_chunks.push(InstanceChunks { chunks: vec![], slow: false });
                self.n_instances += 1;
                self.n_tables += 1;
                self.instance_partition.push(-2); // Mark as table
                self.worker_instances.push(gid);
                let lid = self.process_count[process_id];
                self.process_count[process_id] += 1;
                self.process_weight[process_id] += table.weight;
                if process_id == self.process_id {
                    self.process_instances.push(gid);
                }
                self.aux_table_map[table_idx] = gid as i32;
                self.instance_process.push((process_id as i32, lid));
            } else {
                for rank in 0..self.n_processes {
                    let gid = self.instances.len();
                    self.instances.push(InstanceInfo::new(
                        table.airgroup_id,
                        table.air_id,
                        true,
                        false,
                        table.weight,
                        0,
                    ));
                    self.instances_chunks.push(InstanceChunks { chunks: vec![], slow: false });
                    self.witness_states.push(WitnessSlot::default());
                    self.n_instances += 1;
                    self.n_tables += 1;
                    self.instance_partition.push(-2); // Mark as table
                    self.worker_instances.push(gid);
                    let lid = self.process_count[rank];
                    self.process_count[rank] += 1;
                    self.process_weight[rank] += table.weight;
                    if rank == self.process_id {
                        self.process_instances.push(gid);
                        self.aux_table_map[table_idx] = gid as i32;
                    }
                    self.instance_process.push((rank as i32, lid));
                }
            }
        }
        for key in self.table_assignment.keys() {
            let matched = (0..self.aux_tables.len())
                .any(|i| self.aux_tables[i].airgroup_id == key.0 && self.aux_tables[i].air_id == key.1);
            if !matched {
                return Err(ProofmanError::InvalidAssignation(format!(
                    "Table assignment for airgroup_id: {}, air_id: {} has no registered table",
                    key.0, key.1
                )));
            }
        }
        self.aux_tables.clear();
        self.assignation_done = true;

        Ok(())
    }

    /// Record that the table for (airgroup_id, air_id) must be assigned to the
    /// SAME worker (partition) and SAME process as instance `gid`. Recording only;
    /// the actual placement happens in assign_instances(). Must be called before
    /// assignation is done. Reference validity is checked in assign_instances().
    pub fn assign_table_to(&mut self, airgroup_id: usize, air_id: usize, gid: usize) -> ProofmanResult<()> {
        if self.assignation_done {
            return Err(ProofmanError::InvalidAssignation("Instances already assigned".to_string()));
        }
        match self.table_assignment.entry((airgroup_id, air_id)) {
            Entry::Occupied(_) => {
                return Err(ProofmanError::InvalidAssignation(format!(
                    "Table assignment already set for airgroup_id: {airgroup_id}, air_id: {air_id}"
                )))
            }
            Entry::Vacant(slot) => {
                slot.insert(gid);
            }
        }
        Ok(())
    }

    /// True if instance_id is an assigned table (single owner, no cross-rank
    /// multiplicity reduction). Valid post-assignment.
    pub fn is_assigned_table(&self, instance_id: usize) -> ProofmanResult<bool> {
        Ok(self.assigned_table_instances.contains(&instance_id))
    }

    ///  Load balance info for partitions
    ///  Does not include tables
    pub fn load_balance_info_partition(&self) -> (f64, u64, u64, f64) {
        let mut average_partition_weight = 0.0;
        let mut max_partition_weight = 0;
        let mut min_partition_weight = u64::MAX;
        for i in 0..self.n_partitions {
            average_partition_weight += self.partition_weight[i] as f64;
            if self.partition_weight[i] > max_partition_weight {
                max_partition_weight = self.partition_weight[i];
            }
            if self.partition_weight[i] < min_partition_weight {
                min_partition_weight = self.partition_weight[i];
            }
        }
        average_partition_weight /= self.n_partitions as f64;
        let max_deviation = max_partition_weight as f64 / average_partition_weight;
        (average_partition_weight, max_partition_weight, min_partition_weight, max_deviation)
    }

    /// Load balance info for processes
    pub fn load_balance_info_process(&self) -> (f64, u64, u64, f64) {
        let mut average_process_weight = 0.0;
        let mut max_process_weight = 0;
        let mut min_process_weight = u64::MAX;

        for i in 0..self.n_processes {
            average_process_weight += self.process_weight[i] as f64;
            if self.process_weight[i] > max_process_weight {
                max_process_weight = self.process_weight[i];
            }
            if self.process_weight[i] < min_process_weight {
                min_process_weight = self.process_weight[i];
            }
        }
        average_process_weight /= self.n_processes as f64;
        let max_deviation = max_process_weight as f64 / average_process_weight;
        (average_process_weight, max_process_weight, min_process_weight, max_deviation)
    }
}

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

    // A compressor air (weight 1000 + compressor 100) and a heavier plain one (5000)
    const HEAVY_PLAIN: (usize, u64, u64) = (0, 5000, 0);
    const COMPRESSOR: (usize, u64, u64) = (1, 1000, 100);
    const LIGHT_PLAIN: (usize, u64, u64) = (2, 100, 0);

    fn ctx(n_partitions: usize) -> DistributionCtx {
        let mut dctx = DistributionCtx::new();
        dctx.setup_processes(1, 0).unwrap();
        dctx.setup_partitions(n_partitions, (0..n_partitions as u32).collect()).unwrap();
        dctx
    }

    /// Were compressor count a criterion of its own, the second COMPRESSOR would land on the
    /// overloaded partition 0 (6100 vs 1100) instead of partition 1 (5000 vs 2200).
    #[test]
    fn compressor_air_follows_cost_not_compressor_count() {
        let mut dctx = ctx(2);
        for (air_id, weight, compressor_weight) in [HEAVY_PLAIN, COMPRESSOR, COMPRESSOR] {
            dctx.add_instance_no_assign(0, air_id, weight, compressor_weight).unwrap();
        }
        dctx.assign_instances().unwrap();

        assert_eq!(dctx.instance_partition, vec![0, 1, 1]);
        assert_eq!(dctx.partition_weight, vec![5000, 2200]);
        assert_eq!(dctx.partition_compressor_count, vec![0, 2]);
    }

    /// Instances assigned as they are created go to the least loaded partition, not to `gid %
    /// n_partitions`: round-robin would put the second LIGHT_PLAIN back on the heavy partition 0.
    #[test]
    fn immediate_assignation_balances_cost_instead_of_round_robin() {
        let mut dctx = ctx(2);
        for (air_id, weight, compressor_weight) in [HEAVY_PLAIN, LIGHT_PLAIN, LIGHT_PLAIN] {
            dctx.add_instance(0, air_id, weight, compressor_weight).unwrap();
        }

        assert_eq!(dctx.instance_partition, vec![0, 1, 1]);
        assert_eq!(dctx.partition_weight, vec![5000, 200]);
    }

    /// The first instance added still lands on partition 0 (zisk relies on it for the Rom instance)
    #[test]
    fn first_instance_goes_to_partition_zero() {
        let mut dctx = ctx(4);
        dctx.add_instance(0, LIGHT_PLAIN.0, LIGHT_PLAIN.1, LIGHT_PLAIN.2).unwrap();

        assert_eq!(dctx.instance_partition[0], 0);
        assert_eq!(dctx.instance_process[0], (0, 0));
    }

    /// The compressor proof is part of the load a partition is given
    #[test]
    fn compressor_weight_counts_towards_the_partition_load() {
        let mut dctx = ctx(1);
        dctx.add_instance(0, COMPRESSOR.0, COMPRESSOR.1, COMPRESSOR.2).unwrap();

        assert_eq!(dctx.partition_weight, vec![1100]);
        assert_eq!(dctx.process_weight, vec![1100]);
        assert_eq!(dctx.partition_compressor_count, vec![1]);
    }

    fn dctx_1w_2p() -> DistributionCtx {
        // 1 partition (this worker owns it), 2 processes, we are process 0
        let mut dctx = DistributionCtx::new();
        dctx.setup_partitions(1, vec![0]).expect("setup_partitions");
        dctx.setup_processes(2, 0).expect("setup_processes");
        dctx
    }

    #[test]
    fn assign_table_to_records_intent() {
        let mut dctx = dctx_1w_2p();
        // add a regular instance to act as the reference (immediate path sets coords)
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance");
        // register the table normally
        dctx.add_table(7, 1, 50).expect("add_table");
        // record the assignment intent
        dctx.assign_table_to(7, 1, rom).expect("assign_table_to");
        assert_eq!(dctx.table_assignment.get(&(7, 1)), Some(&rom));
    }

    #[test]
    fn assign_table_to_rejects_after_assignation_done() {
        let mut dctx = dctx_1w_2p();
        dctx.assignation_done = true;
        let err = dctx.assign_table_to(7, 1, 0);
        assert!(matches!(err, Err(ProofmanError::InvalidAssignation(_))));
    }

    #[test]
    fn assign_table_to_rejects_duplicate_key() {
        let mut dctx = dctx_1w_2p();
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance");
        dctx.add_table(7, 1, 50).expect("add_table");
        dctx.assign_table_to(7, 1, rom).expect("first assign_table_to");
        let err = dctx.assign_table_to(7, 1, rom);
        assert!(matches!(err, Err(ProofmanError::InvalidAssignation(_))));
    }

    #[test]
    fn is_assigned_table_false_for_unknown_id() {
        let dctx = dctx_1w_2p();
        assert!(!dctx.is_assigned_table(999).expect("is_assigned_table"));
    }

    #[test]
    fn assigned_table_copies_reference_coords_and_weight() {
        let mut dctx = dctx_1w_2p();
        // Two regular instances so process weights differ; ROM is the reference.
        let _other = dctx.add_instance(7, 0, 100, 0).expect("add_instance other");
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        let rom_partition = dctx.instance_partition[rom];
        let rom_process = dctx.instance_process[rom].0;
        let weight_before = dctx.process_weight[rom_process as usize];

        dctx.add_table(7, 1, 50).expect("add_table");
        dctx.assign_table_to(7, 1, rom).expect("assign_table_to");
        dctx.assign_instances().expect("assign_instances");

        // The table instance is the last instance pushed.
        let table_gid = dctx.get_table_instance_idx(0).expect("table idx");
        assert_eq!(dctx.instance_partition[table_gid], rom_partition, "real partition, not -2");
        assert_eq!(dctx.instance_process[table_gid].0, rom_process, "same process as ROM");
        assert!(dctx.is_assigned_table(table_gid).expect("is_assigned_table"));
        assert_eq!(dctx.witness_states.len(), dctx.instances.len(), "witness slots stay aligned");
        assert_eq!(
            dctx.process_weight[rom_process as usize],
            weight_before + 50,
            "table weight lands on ROM's process"
        );
    }

    #[test]
    fn assigned_table_on_foreign_worker_is_listed_but_not_materialized() {
        // 2 partitions, this worker owns only partition 1, so the reference lands on partition 0.
        let mut dctx = DistributionCtx::new();
        dctx.setup_partitions(2, vec![1]).expect("setup_partitions");
        dctx.setup_processes(2, 0).expect("setup_processes");

        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        assert_eq!(dctx.instance_process[rom].0, -1, "reference is not owned by this worker");
        dctx.add_table(7, 1, 50).expect("add_table");
        dctx.assign_table_to(7, 1, rom).expect("assign_table_to");
        let counts_before = dctx.process_count.clone();
        let weights_before = dctx.process_weight.clone();
        dctx.assign_instances().expect("must not abort on a foreign reference");

        // Still in the global list (same gids on every worker), but not materialized here.
        let table_gid = dctx.get_table_instance_idx(0).expect("table idx");
        assert_eq!(dctx.instance_partition[table_gid], 0, "follows the reference partition");
        assert_eq!(dctx.instance_process[table_gid], (-1, 0), "not owned by any local process");
        assert!(!dctx.worker_instances.contains(&table_gid));
        assert!(!dctx.process_instances.contains(&table_gid));
        assert_eq!(dctx.process_count, counts_before);
        assert_eq!(dctx.process_weight, weights_before);
    }

    #[test]
    fn duplicate_assign_table_to_keeps_the_original_gid() {
        let mut dctx = dctx_1w_2p();
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        let other = dctx.add_instance(7, 0, 100, 0).expect("add_instance other");
        dctx.add_table(7, 1, 50).expect("add_table");
        dctx.assign_table_to(7, 1, rom).expect("first assign_table_to");
        assert!(matches!(dctx.assign_table_to(7, 1, other), Err(ProofmanError::InvalidAssignation(_))));
        assert_eq!(dctx.table_assignment.get(&(7, 1)), Some(&rom), "rejected assignment must not overwrite");
    }

    #[test]
    fn assigned_table_overrides_all_ranks_to_single_instance() {
        let mut dctx = dctx_1w_2p();
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        let instances_before = dctx.instances.len();
        // Register as ALL-RANKS (shared=false) but then assign it.
        dctx.add_table_all(7, 1, 50).expect("add_table_all");
        dctx.assign_table_to(7, 1, rom).expect("assign_table_to");
        dctx.assign_instances().expect("assign_instances");
        // All-ranks would have added n_processes (2) instances; assignment collapses to 1.
        assert_eq!(dctx.instances.len(), instances_before + 1, "single instance, not replicated");
    }

    #[test]
    fn assigned_table_with_out_of_range_reference_errors() {
        let mut dctx = dctx_1w_2p();
        let _rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        dctx.add_table(7, 1, 50).expect("add_table");
        // reference a gid that does not exist
        dctx.assign_table_to(7, 1, 999).expect("assign_table_to");
        let err = dctx.assign_instances();
        assert!(matches!(err, Err(ProofmanError::InvalidAssignation(_))));
    }

    #[test]
    fn assigned_table_with_unknown_key_errors() {
        let mut dctx = dctx_1w_2p();
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        // assignment for a table that was never registered in aux_tables
        dctx.assign_table_to(7, 99, rom).expect("assign_table_to");
        let err = dctx.assign_instances();
        assert!(matches!(err, Err(ProofmanError::InvalidAssignation(_))));
    }

    #[test]
    fn ordinary_shared_table_is_not_assigned() {
        let mut dctx = dctx_1w_2p();
        let _rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance");
        dctx.add_table(7, 1, 50).expect("add_table");
        dctx.assign_instances().expect("assign_instances");
        let table_gid = dctx.get_table_instance_idx(0).expect("table idx");
        assert!(!dctx.is_assigned_table(table_gid).expect("is_assigned_table"));
    }

    #[test]
    fn reset_clears_assignment_state_across_cycles() {
        let mut dctx = dctx_1w_2p();
        // Cycle 1: assign a table to a reference, then run assignment.
        let rom = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom");
        dctx.add_table(7, 1, 50).expect("add_table");
        dctx.assign_table_to(7, 1, rom).expect("assign_table_to cycle1");
        dctx.assign_instances().expect("assign cycle1");
        let table_gid_c1 = dctx.get_table_instance_idx(0).expect("table idx c1");
        assert!(dctx.is_assigned_table(table_gid_c1).expect("assigned c1"));

        // Reset for the next proof cycle.
        dctx.reset_instances();

        // table_assignment must be empty so the same key can be re-registered
        // without a duplicate-key error.
        assert!(dctx.table_assignment.is_empty(), "table_assignment not cleared on reset");
        // assigned_table_instances must be empty so stale gids don't misclassify.
        assert!(dctx.assigned_table_instances.is_empty(), "assigned_table_instances not cleared on reset");

        // Cycle 2: re-registering the same (ag, air) assignment must succeed.
        let rom2 = dctx.add_instance(7, 0, 100, 0).expect("add_instance rom cycle2");
        dctx.add_table(7, 1, 50).expect("add_table cycle2");
        dctx.assign_table_to(7, 1, rom2).expect("assign_table_to cycle2 must not be duplicate");
        dctx.assign_instances().expect("assign cycle2");
        let table_gid_c2 = dctx.get_table_instance_idx(0).expect("table idx c2");
        assert!(dctx.is_assigned_table(table_gid_c2).expect("assigned c2"));
    }
}