prism-q 0.35.0

//! Shot collections and probability-weighted shot sampling.

use std::collections::HashMap;

use crate::hash::FxHashMap;

use rand::RngExt;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;

use super::Probabilities;
use super::compiled;
use super::metadata::{ResolvedBackend, RunMetadata};

/// Marker on a result no route has stamped yet; `every_terminal_stamps_metadata`
/// asserts none reaches a caller.
pub(crate) const UNSTAMPED: &str = "unstamped";

/// Result of a multi-shot simulation run.
#[derive(Debug, Clone)]
pub struct ShotsResult {
    /// `shots[i][j]` is the j-th classical bit from the i-th shot.
    pub shots: Vec<Vec<bool>>,
    pub(crate) num_classical_bits: usize,
    /// Which engine ran, whether the answer is exact, and where the state lived.
    pub metadata: RunMetadata,
    /// `leaked[i][q]` is `true` when qubit `q` leaked at any point of shot `i`,
    /// the erasure herald of a model carrying
    /// [`NoiseChannel::Leakage`](crate::NoiseChannel::Leakage); `None` for every
    /// other run.
    pub leaked: Option<Vec<Vec<bool>>>,
}

impl ShotsResult {
    /// Build with placeholder provenance; the entry point that picked the
    /// engine overwrites it through [`ShotsResult::with_metadata`], which
    /// `every_terminal_stamps_metadata` pins.
    pub(crate) fn from_shots(shots: Vec<Vec<bool>>, num_classical_bits: usize) -> Self {
        Self {
            shots,
            num_classical_bits,
            metadata: RunMetadata::exact(ResolvedBackend::Other(UNSTAMPED)),
            leaked: None,
        }
    }

    pub(crate) fn with_leaked(mut self, leaked: Vec<Vec<bool>>) -> Self {
        self.leaked = Some(leaked);
        self
    }

    pub(crate) fn with_metadata(mut self, metadata: RunMetadata) -> Self {
        let shots = self.shots.len();
        self.metadata = metadata.with_shots(shots);
        self
    }

    /// Keys are packed `Vec<u64>` where bit `i` of word `i/64` corresponds
    /// to classical bit `i`. Use [`bitstring`] to format keys for display.
    pub fn counts(&self) -> HashMap<Vec<u64>, u64> {
        let m_words = self.num_classical_bits.div_ceil(64).max(1);
        let mut counts: HashMap<Vec<u64>, u64> = HashMap::new();
        let mut key = vec![0u64; m_words];
        for shot in &self.shots {
            key.fill(0);
            for (word, bits) in key.iter_mut().zip(shot.chunks(64)) {
                *word = bits
                    .iter()
                    .enumerate()
                    .fold(0, |acc, (i, &b)| acc | (u64::from(b) << i));
            }
            match counts.get_mut(key.as_slice()) {
                Some(count) => *count += 1,
                None => {
                    counts.insert(key.clone(), 1);
                }
            }
        }
        counts
    }

    pub fn num_shots(&self) -> usize {
        self.shots.len()
    }

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

#[cfg(test)]
impl ShotsResult {
    pub(crate) fn marginal(&self, bit: usize) -> f64 {
        self.shots.iter().filter(|s| s[bit]).count() as f64 / self.shots.len() as f64
    }

    pub(crate) fn coherent_fraction(&self) -> f64 {
        self.shots
            .iter()
            .filter(|s| s.iter().all(|&b| b == s[0]))
            .count() as f64
            / self.shots.len() as f64
    }
}

impl std::fmt::Display for ShotsResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let counts = self.counts();
        let mut entries: Vec<_> = counts.into_iter().collect();
        entries.sort_by_key(|e| std::cmp::Reverse(e.1));
        for (bits, count) in &entries {
            let bs = bitstring(bits, self.num_classical_bits);
            writeln!(f, "{bs}: {count}")?;
        }
        Ok(())
    }
}

/// Format a packed `Vec<u64>` key (from [`ShotsResult::counts`]) as a binary string.
///
/// Bit 0 of the first word corresponds to classical bit 0 (leftmost character).
pub fn bitstring(key: &[u64], num_bits: usize) -> String {
    let mut s = String::with_capacity(num_bits);
    for i in 0..num_bits {
        let word = i / 64;
        let bit = i % 64;
        if word < key.len() && (key[word] >> bit) & 1 == 1 {
            s.push('1');
        } else {
            s.push('0');
        }
    }
    s
}

pub(crate) fn build_cdf(probs: &[f64]) -> Vec<f64> {
    let mut cdf = Vec::with_capacity(probs.len());
    let mut acc = 0.0;
    for &p in probs {
        acc += p;
        cdf.push(acc);
    }
    if let Some(last) = cdf.last_mut() {
        *last = 1.0;
    }
    cdf
}

/// Shots drawn per ChaCha8 substream by the block samplers. One stream per
/// block rather than per shot: cipher setup and the first keystream block cost
/// more than a whole draw when the per-shot work is short.
pub(crate) const SHOTS_PER_STREAM: usize = 256;

/// Shot count below which a block sampler stays on the calling thread.
pub(crate) const MIN_SHOTS_FOR_PAR: usize = 32;

/// Draws a call must make before the product and factored samplers split their
/// blocks across workers: 2000 factored draws (1000 shots over two sub-states)
/// ran 18% slower forked than serial.
const MIN_DRAWS_FOR_PAR: usize = 1 << 14;

/// Whether `num_shots` shots of `draws_per_shot` draws each carry enough work
/// to fill their blocks in parallel. The blocks draw the same values either way.
pub(crate) fn draws_split_across_workers(num_shots: usize, draws_per_shot: usize) -> bool {
    num_shots >= MIN_SHOTS_FOR_PAR && num_shots.saturating_mul(draws_per_shot) >= MIN_DRAWS_FOR_PAR
}

/// Fill `rows`, `per_shot` entries to a shot, one block of [`SHOTS_PER_STREAM`]
/// shots at a time. Block `b` draws from stream `b + 2` of the ChaCha8 generator
/// seeded with `seed`, so the partition, and with it every draw, is a function
/// of the shot count alone: `parallel` moves the schedule, never a value.
pub(crate) fn sample_in_shot_blocks<T: Send>(
    rows: &mut [T],
    per_shot: usize,
    seed: u64,
    parallel: bool,
    sample_block: impl Fn(&mut ChaCha8Rng, &mut [T]) + Sync,
) {
    let block_len = SHOTS_PER_STREAM * per_shot;
    let run = |block: usize, rows: &mut [T]| {
        let mut rng = ChaCha8Rng::seed_from_u64(seed);
        rng.set_stream(block as u64 + 2);
        sample_block(&mut rng, rows);
    };

    #[cfg(feature = "parallel")]
    if parallel {
        use rayon::prelude::*;
        rows.par_chunks_mut(block_len)
            .enumerate()
            .for_each(|(block, rows)| run(block, rows));
        return;
    }
    #[cfg(not(feature = "parallel"))]
    let _ = parallel;

    for (block, rows) in rows.chunks_mut(block_len).enumerate() {
        run(block, rows);
    }
}

pub(crate) fn sample_from_cdf(cdf: &[f64], r: f64) -> usize {
    match cdf.binary_search_by(|p| p.partial_cmp(&r).unwrap_or(std::cmp::Ordering::Equal)) {
        Ok(i) => i,
        Err(i) => i.min(cdf.len() - 1),
    }
}

pub(crate) fn sample_shots(
    probs: &Probabilities,
    meas_map: &[(usize, usize)],
    num_classical_bits: usize,
    num_shots: usize,
    seed: u64,
) -> Vec<Vec<bool>> {
    if meas_map.is_empty() {
        return vec![vec![false; num_classical_bits]; num_shots];
    }

    let mut shots = vec![vec![false; num_classical_bits]; num_shots];

    match probs {
        Probabilities::Dense(v) => {
            let mut rng = ChaCha8Rng::seed_from_u64(seed);
            let cdf = build_cdf(v);
            for shot in &mut shots {
                let r: f64 = rng.random();
                let state_idx = sample_from_cdf(&cdf, r);
                for &(qubit, cbit) in meas_map {
                    shot[cbit] = (state_idx >> qubit) & 1 == 1;
                }
            }
        }
        // Same blocks, streams and per-shot draw order as the factored
        // backend's native sampler, so the two routes agree shot for shot.
        Probabilities::Factored { blocks, .. } => {
            let block_cdfs: Vec<Vec<f64>> = blocks.iter().map(|b| build_cdf(&b.probs)).collect();
            let parallel = draws_split_across_workers(num_shots, blocks.len());
            sample_in_shot_blocks(&mut shots, 1, seed, parallel, |rng, block_shots| {
                for shot in block_shots {
                    let mut global_idx = 0usize;
                    for (block, cdf) in blocks.iter().zip(block_cdfs.iter()) {
                        let r: f64 = rng.random();
                        let local_idx = sample_from_cdf(cdf, r);
                        let mut m = block.mask;
                        let mut bit = 0;
                        while m != 0 {
                            let pos = m.trailing_zeros() as usize;
                            if local_idx & (1 << bit) != 0 {
                                global_idx |= 1 << pos;
                            }
                            bit += 1;
                            m &= m.wrapping_sub(1);
                        }
                    }
                    for &(qubit, cbit) in meas_map {
                        shot[cbit] = (global_idx >> qubit) & 1 == 1;
                    }
                }
            });
        }
    }

    shots
}

/// Project packed per-qubit outcomes from a native backend sampler onto
/// classical bits. Later measurements of the same classical bit win, matching
/// [`sample_shots`].
pub(crate) fn shots_from_basis_samples(
    samples: &crate::backend::BasisSamples,
    meas_map: &[(usize, usize)],
    num_classical_bits: usize,
) -> Vec<Vec<bool>> {
    let dense_identity_map = meas_map.len() == num_classical_bits
        && meas_map
            .iter()
            .enumerate()
            .all(|(idx, &(qubit, classical_bit))| idx == qubit && idx == classical_bit);
    if dense_identity_map {
        return samples.to_shots(num_classical_bits);
    }

    let mut shots = vec![vec![false; num_classical_bits]; samples.num_shots()];
    for (index, shot) in shots.iter_mut().enumerate() {
        for &(qubit, cbit) in meas_map {
            shot[cbit] = samples.bit(index, qubit);
        }
    }
    shots
}

/// Re-key counts drawn in measurement-record order onto the classical
/// register, merging records that land on the same classical outcome. A later
/// measurement of a classical bit overwrites an earlier one.
pub(super) fn record_counts_to_classical_bits(
    counts: HashMap<Vec<u64>, u64>,
    meas_map: &[(usize, usize)],
    num_classical_bits: usize,
) -> HashMap<Vec<u64>, u64> {
    let words = num_classical_bits.div_ceil(64).max(1);
    let in_order = meas_map.len() <= num_classical_bits
        && meas_map
            .iter()
            .enumerate()
            .all(|(idx, &(_, classical_bit))| idx == classical_bit);
    if in_order && counts.keys().all(|key| key.len() == words) {
        return counts;
    }

    let mut out: HashMap<Vec<u64>, u64> = HashMap::with_capacity(counts.len());
    for (record, count) in counts {
        let key = if in_order {
            let mut key = record;
            key.resize(words, 0);
            key
        } else {
            let mut key = vec![0u64; words];
            for (idx, &(_, classical_bit)) in meas_map.iter().enumerate() {
                if classical_bit >= num_classical_bits {
                    continue;
                }
                let bit = record
                    .get(idx / 64)
                    .is_some_and(|w| (w >> (idx % 64)) & 1 == 1);
                let mask = 1u64 << (classical_bit % 64);
                if bit {
                    key[classical_bit / 64] |= mask;
                } else {
                    key[classical_bit / 64] &= !mask;
                }
            }
            key
        };
        *out.entry(key).or_insert(0) += count;
    }
    out
}

/// Histogram of native samples keyed as [`ShotsResult::counts`] keys them,
/// equal to counting [`shots_from_basis_samples`]. Rows are counted with
/// unmeasured qubits masked off, so only distinct outcomes pay the projection
/// onto classical bits.
pub(crate) fn counts_from_basis_samples(
    samples: &crate::backend::BasisSamples,
    meas_map: &[(usize, usize)],
    num_classical_bits: usize,
) -> HashMap<Vec<u64>, u64> {
    let words_per_shot = samples.words_per_shot();
    let mut mask = vec![0u64; words_per_shot];
    for &(qubit, _) in meas_map {
        mask[qubit / 64] |= 1u64 << (qubit % 64);
    }
    let m_words = num_classical_bits.div_ceil(64).max(1);
    let in_place = meas_map.iter().all(|&(qubit, cbit)| qubit == cbit);
    let classical_key = |row: &[u64]| -> Vec<u64> {
        let mut key = vec![0u64; m_words];
        if in_place {
            let live = m_words.min(row.len());
            key[..live].copy_from_slice(&row[..live]);
        } else {
            for &(qubit, cbit) in meas_map {
                let bit = (row[qubit / 64] >> (qubit % 64)) & 1;
                let word = &mut key[cbit / 64];
                *word = (*word & !(1u64 << (cbit % 64))) | (bit << (cbit % 64));
            }
        }
        key
    };

    let mut counts = HashMap::new();
    if words_per_shot == 1 {
        let mask = mask[0];
        let mut rows: FxHashMap<u64, u64> = FxHashMap::default();
        for &word in samples.words() {
            *rows.entry(word & mask).or_insert(0) += 1;
        }
        for (row, n) in rows {
            *counts.entry(classical_key(&[row])).or_insert(0) += n;
        }
    } else {
        let mut rows: FxHashMap<Vec<u64>, u64> = FxHashMap::default();
        let mut scratch = vec![0u64; words_per_shot];
        for row in samples.words().chunks_exact(words_per_shot) {
            for ((dst, &word), &m) in scratch.iter_mut().zip(row).zip(&mask) {
                *dst = word & m;
            }
            match rows.get_mut(scratch.as_slice()) {
                Some(n) => *n += 1,
                None => {
                    rows.insert(scratch.clone(), 1);
                }
            }
        }
        for (row, n) in rows {
            *counts.entry(classical_key(&row)).or_insert(0) += n;
        }
    }
    counts
}

pub(super) fn packed_shots_to_classical_bits(
    packed: &compiled::PackedShots,
    meas_map: &[(usize, usize)],
    num_classical_bits: usize,
) -> Vec<Vec<bool>> {
    let dense_identity_map = meas_map.len() == num_classical_bits
        && meas_map
            .iter()
            .enumerate()
            .all(|(idx, &(_, classical_bit))| idx == classical_bit);
    if dense_identity_map {
        return packed.to_shots();
    }

    let mut shots = vec![vec![false; num_classical_bits]; packed.num_shots()];
    let mut seen = vec![false; num_classical_bits];
    let unique_classical_bits = meas_map.iter().all(|&(_, classical_bit)| {
        if classical_bit >= num_classical_bits {
            return true;
        }
        if seen[classical_bit] {
            false
        } else {
            seen[classical_bit] = true;
            true
        }
    });

    if unique_classical_bits {
        match packed.layout() {
            compiled::ShotLayout::ShotMajor => {
                let m_words = packed.m_words();
                let data = packed.raw_data();
                for (shot_idx, shot) in shots.iter_mut().enumerate() {
                    let row = &data[shot_idx * m_words..(shot_idx + 1) * m_words];
                    for (measurement, &(_, classical_bit)) in meas_map.iter().enumerate() {
                        if classical_bit >= num_classical_bits {
                            continue;
                        }
                        let word = row[measurement / 64];
                        shot[classical_bit] = (word >> (measurement % 64)) & 1 != 0;
                    }
                }
            }
            compiled::ShotLayout::MeasMajor => {
                let s_words = packed.s_words();
                let data = packed.raw_data();
                let tail = packed.num_shots() % 64;
                let last_mask = if tail == 0 {
                    u64::MAX
                } else {
                    (1u64 << tail) - 1
                };
                for (measurement, &(_, classical_bit)) in meas_map.iter().enumerate() {
                    if classical_bit >= num_classical_bits {
                        continue;
                    }
                    let row = &data[measurement * s_words..(measurement + 1) * s_words];
                    for (sw, mut bits) in row.iter().copied().enumerate() {
                        if sw + 1 == s_words {
                            bits &= last_mask;
                        }
                        while bits != 0 {
                            let shot = sw * 64 + bits.trailing_zeros() as usize;
                            shots[shot][classical_bit] = true;
                            bits &= bits - 1;
                        }
                    }
                }
            }
        }
        return shots;
    }

    for (measurement, &(_, classical_bit)) in meas_map.iter().enumerate() {
        if classical_bit >= num_classical_bits {
            continue;
        }
        for (shot_idx, shot) in shots.iter_mut().enumerate() {
            shot[classical_bit] = packed.get_bit(shot_idx, measurement);
        }
    }
    shots
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sim::compiled::PackedShots;
    use crate::sim::probability::{FactoredBlock, Probabilities};

    fn basis_samples_fixture(num_shots: usize, num_qubits: usize) -> crate::backend::BasisSamples {
        let mut samples = crate::backend::BasisSamples::new(num_shots, num_qubits);
        for shot in 0..num_shots {
            for qubit in 0..num_qubits {
                if (shot * 31 + qubit * 7) % 3 == 0 {
                    samples.set(shot, qubit);
                }
            }
        }
        samples
    }

    fn shots_per_bit(
        samples: &crate::backend::BasisSamples,
        meas_map: &[(usize, usize)],
        num_classical_bits: usize,
    ) -> Vec<Vec<bool>> {
        let mut shots = vec![vec![false; num_classical_bits]; samples.num_shots()];
        for (index, shot) in shots.iter_mut().enumerate() {
            for &(qubit, cbit) in meas_map {
                shot[cbit] = samples.bit(index, qubit);
            }
        }
        shots
    }

    #[test]
    fn basis_sample_fast_path_matches_per_bit_expansion() {
        for num_qubits in [4usize, 64, 70, 129] {
            let samples = basis_samples_fixture(17, num_qubits);
            let meas_map: Vec<(usize, usize)> = (0..num_qubits).map(|q| (q, q)).collect();
            assert_eq!(
                shots_from_basis_samples(&samples, &meas_map, num_qubits),
                shots_per_bit(&samples, &meas_map, num_qubits)
            );
        }
    }

    #[test]
    fn basis_sample_fast_path_ignores_qubits_above_the_classical_register() {
        let samples = basis_samples_fixture(9, 70);
        let meas_map: Vec<(usize, usize)> = (0..5).map(|q| (q, q)).collect();
        assert_eq!(
            shots_from_basis_samples(&samples, &meas_map, 5),
            shots_per_bit(&samples, &meas_map, 5)
        );
    }

    // A classical register wider than the qubit register: the straddling word
    // lies past the last word the samples hold, so nothing may be masked off.
    #[test]
    fn basis_sample_unpack_keeps_every_bit_of_a_narrow_register() {
        let samples = basis_samples_fixture(5, 70);
        let unpacked = samples.to_shots(130);
        for (shot, row) in unpacked.iter().enumerate() {
            for (qubit, &got) in row[..70].iter().enumerate() {
                assert_eq!(got, samples.bit(shot, qubit), "shot {shot} qubit {qubit}");
            }
            assert!(row[70..].iter().all(|b| !b));
        }
    }

    #[test]
    fn basis_sample_permuted_map_takes_the_general_path() {
        let samples = basis_samples_fixture(6, 8);
        let meas_map = [(3, 0), (0, 1), (7, 2)];
        assert_eq!(
            shots_from_basis_samples(&samples, &meas_map, 3),
            shots_per_bit(&samples, &meas_map, 3)
        );
    }

    #[test]
    fn basis_sample_counts_match_the_expanded_histogram() {
        let cases = [
            (8, 8, (0..8).map(|q| (q, q)).collect::<Vec<_>>()),
            (8, 3, vec![(3, 0), (0, 1), (7, 2)]),
            (8, 3, vec![(1, 0), (2, 0), (5, 2), (6, 0)]),
            (8, 20, vec![(1, 7), (4, 19)]),
            (8, 130, vec![(0, 0), (3, 64), (5, 129)]),
            (130, 130, (0..130).map(|q| (q, q)).collect()),
            (
                130,
                70,
                vec![(129, 0), (64, 69), (2, 64), (100, 64), (63, 63)],
            ),
        ];
        for (num_qubits, num_classical_bits, meas_map) in cases {
            let samples = basis_samples_fixture(40, num_qubits);
            let expanded = ShotsResult::from_shots(
                shots_per_bit(&samples, &meas_map, num_classical_bits),
                num_classical_bits,
            )
            .counts();
            assert_eq!(
                counts_from_basis_samples(&samples, &meas_map, num_classical_bits),
                expanded,
                "{num_qubits} qubits, {num_classical_bits} bits, map {meas_map:?}"
            );
        }
        assert!(counts_from_basis_samples(&basis_samples_fixture(0, 8), &[(0, 0)], 1).is_empty());
    }

    #[test]
    fn build_cdf_normalizes_last_to_one() {
        let cdf = build_cdf(&[0.2, 0.3, 0.4999]);
        assert_eq!(cdf.len(), 3);
        assert!((cdf[0] - 0.2).abs() < 1e-12);
        assert!((cdf[1] - 0.5).abs() < 1e-12);
        assert!((cdf[2] - 1.0).abs() < 1e-12);
    }

    #[test]
    fn build_cdf_empty_and_single() {
        let empty = build_cdf(&[]);
        assert!(empty.is_empty());
        let single = build_cdf(&[0.42]);
        assert_eq!(single, vec![1.0]);
    }

    #[test]
    fn sample_from_cdf_bounds() {
        let cdf = [0.25, 0.5, 0.75, 1.0];
        assert_eq!(sample_from_cdf(&cdf, 0.0), 0);
        assert_eq!(sample_from_cdf(&cdf, 0.3), 1);
        assert_eq!(sample_from_cdf(&cdf, 0.99), 3);
        assert_eq!(sample_from_cdf(&cdf, 1.0), 3);
    }

    #[test]
    fn bitstring_packs_bits_lsb_first() {
        let bits = vec![0b1011u64];
        let s = bitstring(&bits, 4);
        assert_eq!(s, "1101");
    }

    #[test]
    fn bitstring_short_key_pads_zero() {
        let s = bitstring(&[], 3);
        assert_eq!(s, "000");
    }

    #[test]
    fn shots_result_counts_and_display() {
        let result = ShotsResult::from_shots(
            vec![vec![true, false], vec![true, false], vec![false, true]],
            2,
        );
        assert_eq!(result.num_shots(), 3);
        assert_eq!(result.num_classical_bits(), 2);
        let counts = result.counts();
        assert_eq!(counts.len(), 2);
        let s = format!("{}", result);
        assert!(s.contains("10: 2"));
        assert!(s.contains("01: 1"));
    }

    #[test]
    fn shots_result_counts_match_per_bit_keys_across_words() {
        let num_bits: usize = 130;
        let mut state = 42u64;
        let mut distinct = Vec::new();
        for _ in 0..7 {
            let shot: Vec<bool> = (0..num_bits)
                .map(|_| {
                    state ^= state << 13;
                    state ^= state >> 7;
                    state ^= state << 17;
                    state & 1 == 1
                })
                .collect();
            distinct.push(shot);
        }
        let shots: Vec<Vec<bool>> = (0..200).map(|i| distinct[(i * i) % 7].clone()).collect();

        let mut expected: HashMap<Vec<u64>, u64> = HashMap::new();
        for shot in &shots {
            let mut key = vec![0u64; num_bits.div_ceil(64)];
            for (i, &b) in shot.iter().enumerate() {
                key[i / 64] |= u64::from(b) << (i % 64);
            }
            *expected.entry(key).or_insert(0) += 1;
        }

        let counts = ShotsResult::from_shots(shots, num_bits).counts();
        assert_eq!(counts, expected);
    }

    #[test]
    fn sample_shots_empty_meas_map_returns_all_false() {
        let probs = Probabilities::Dense(vec![1.0]);
        let shots = sample_shots(&probs, &[], 3, 4, 42);
        assert_eq!(shots.len(), 4);
        for shot in shots {
            assert_eq!(shot, vec![false, false, false]);
        }
    }

    #[test]
    fn sample_shots_dense_deterministic() {
        let probs = Probabilities::Dense(vec![0.0, 1.0]);
        let shots = sample_shots(&probs, &[(0, 0)], 1, 5, 42);
        for shot in shots {
            assert_eq!(shot, vec![true]);
        }
    }

    #[test]
    fn sample_shots_factored_reconstructs_global_index() {
        let probs = Probabilities::Factored {
            blocks: vec![
                FactoredBlock {
                    probs: vec![0.0, 1.0],
                    mask: 0b001,
                },
                FactoredBlock {
                    probs: vec![0.0, 0.0, 0.0, 1.0],
                    mask: 0b110,
                },
            ],
            total_qubits: 3,
        };
        let shots = sample_shots(&probs, &[(0, 0), (1, 1), (2, 2)], 3, 10, 42);
        for shot in shots {
            assert_eq!(shot, vec![true, true, true]);
        }
    }

    #[test]
    fn packed_shots_identity_fast_path() {
        let mut data = vec![0u64; 4];
        data[0] = 0b101;
        data[1] = 0b010;
        data[2] = 0b111;
        data[3] = 0b000;
        let packed = PackedShots::from_shot_major(data, 4, 3);
        let meas_map = [(0, 0), (0, 1), (0, 2)];
        let shots = packed_shots_to_classical_bits(&packed, &meas_map, 3);
        assert_eq!(shots.len(), 4);
        assert_eq!(shots[0], vec![true, false, true]);
        assert_eq!(shots[1], vec![false, true, false]);
        assert_eq!(shots[2], vec![true, true, true]);
        assert_eq!(shots[3], vec![false, false, false]);
    }

    #[test]
    fn packed_shots_non_identity_mapping() {
        let mut data = vec![0u64; 2];
        data[0] = 0b011;
        data[1] = 0b100;
        let packed = PackedShots::from_shot_major(data, 2, 3);
        let meas_map = [(0, 2), (0, 1), (0, 0)];
        let shots = packed_shots_to_classical_bits(&packed, &meas_map, 3);
        assert_eq!(shots[0], vec![false, true, true]);
        assert_eq!(shots[1], vec![true, false, false]);
    }

    #[test]
    fn packed_shots_out_of_range_classical_bit_skipped() {
        let data = vec![0b111u64];
        let packed = PackedShots::from_shot_major(data, 1, 3);
        let meas_map = [(0, 0), (0, 5), (0, 1)];
        let shots = packed_shots_to_classical_bits(&packed, &meas_map, 2);
        assert_eq!(shots[0], vec![true, true]);
    }
}