wai-quantum 0.3.30

//! Stabilizer (CHP) backend — the O(n²) Clifford tier (`wai.quantum.stabilizer`).
//!
//! The cheapest-sufficient cascade's fast tier: a Clifford circuit — however many
//! qubits — is simulated in `O(n²)` per gate and sampled exactly by the
//! Aaronson–Gottesman stabilizer tableau (the CHP algorithm), far past the dense
//! simulator's `2ⁿ` wall. Where the dense sink dies at ~26 qubits, this runs a
//! thousand-qubit Clifford circuit and draws measurement shots from it.
//!
//! Determinism: the tableau is exact integer bit-arithmetic (no floats at all), and
//! measurement randomness comes from the pinned `splitmix64` — so a `(circuit,
//! seed, shots)` histogram is **bit-identical on every machine**, the same contract
//! the dense simulator makes. The Clifford subset it accepts is
//! `H, S, S†, X, Y, Z, P(1)=Z, P(2)=S, CX, CZ, CP(1)` and `SWAP`; anything else
//! (a `T`, a Toffoli) is rejected so the caller falls back to a universal backend.

use crate::quantum::{BaseGate, Circuit, Gate};

fn splitmix64(s: &mut u64) -> u64 {
    *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *s;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// A gate outside the Clifford subset this backend simulates.
#[derive(Debug)]
pub struct NotClifford(pub String);

// ===========================================================================
// Tableau  (2n+1 rows: n destabilizers, n stabilizers, 1 scratch)
// ===========================================================================

/// A stabilizer tableau on `n` qubits. Rows `0..n` are destabilizers, `n..2n`
/// stabilizers, `2n` is scratch. Each row is a signed Pauli stored as bit-packed
/// `x`/`z` words plus a phase `r ∈ {0,1}`.
#[derive(Clone)]
pub struct Tableau {
    n: usize,
    words: usize,
    xs: Vec<u64>,
    zs: Vec<u64>,
    r: Vec<u8>,
}

impl Tableau {
    /// The identity tableau: destabilizer `i = X_i`, stabilizer `i = Z_i`.
    pub fn identity(n: usize) -> Tableau {
        let words = n.div_ceil(64).max(1);
        let rows = 2 * n + 1;
        let mut t = Tableau { n, words, xs: vec![0; rows * words], zs: vec![0; rows * words], r: vec![0; rows] };
        for i in 0..n {
            t.set_x(i, i, 1); // destabilizer i = X_i
            t.set_z(n + i, i, 1); // stabilizer i = Z_i
        }
        t
    }

    #[inline]
    fn xget(&self, row: usize, j: usize) -> u64 {
        (self.xs[row * self.words + (j >> 6)] >> (j & 63)) & 1
    }
    #[inline]
    fn zget(&self, row: usize, j: usize) -> u64 {
        (self.zs[row * self.words + (j >> 6)] >> (j & 63)) & 1
    }
    #[inline]
    fn set_x(&mut self, row: usize, j: usize, v: u64) {
        let w = row * self.words + (j >> 6);
        let b = 1u64 << (j & 63);
        if v != 0 {
            self.xs[w] |= b;
        } else {
            self.xs[w] &= !b;
        }
    }
    #[inline]
    fn set_z(&mut self, row: usize, j: usize, v: u64) {
        let w = row * self.words + (j >> 6);
        let b = 1u64 << (j & 63);
        if v != 0 {
            self.zs[w] |= b;
        } else {
            self.zs[w] &= !b;
        }
    }
    #[inline]
    fn flip_x(&mut self, row: usize, j: usize) {
        self.xs[row * self.words + (j >> 6)] ^= 1u64 << (j & 63);
    }
    #[inline]
    fn flip_z(&mut self, row: usize, j: usize) {
        self.zs[row * self.words + (j >> 6)] ^= 1u64 << (j & 63);
    }

    // --- Clifford primitives (operate over the 2n state rows) ---
    fn cnot(&mut self, a: usize, b: usize) {
        for i in 0..2 * self.n {
            let xa = self.xget(i, a);
            let za = self.zget(i, a);
            let xb = self.xget(i, b);
            let zb = self.zget(i, b);
            self.r[i] ^= (xa & zb & (xb ^ za ^ 1)) as u8;
            if xa != 0 {
                self.flip_x(i, b);
            }
            if zb != 0 {
                self.flip_z(i, a);
            }
        }
    }
    fn hadamard(&mut self, a: usize) {
        for i in 0..2 * self.n {
            let xa = self.xget(i, a);
            let za = self.zget(i, a);
            self.r[i] ^= (xa & za) as u8;
            // swap x,z at (i,a)
            self.set_x(i, a, za);
            self.set_z(i, a, xa);
        }
    }
    fn phase(&mut self, a: usize) {
        for i in 0..2 * self.n {
            let xa = self.xget(i, a);
            let za = self.zget(i, a);
            self.r[i] ^= (xa & za) as u8;
            if xa != 0 {
                self.flip_z(i, a);
            }
            let _ = za;
        }
    }
    fn x(&mut self, a: usize) {
        for i in 0..2 * self.n {
            self.r[i] ^= self.zget(i, a) as u8;
        }
    }
    fn z(&mut self, a: usize) {
        for i in 0..2 * self.n {
            self.r[i] ^= self.xget(i, a) as u8;
        }
    }
    fn y(&mut self, a: usize) {
        for i in 0..2 * self.n {
            self.r[i] ^= (self.xget(i, a) ^ self.zget(i, a)) as u8;
        }
    }
    fn cz(&mut self, a: usize, b: usize) {
        self.hadamard(b);
        self.cnot(a, b);
        self.hadamard(b);
    }

    /// Apply one circuit gate; error on a non-Clifford gate.
    fn apply(&mut self, g: &Gate) -> Result<(), NotClifford> {
        let t = g.target as usize;
        match (g.base, g.controls.as_slice()) {
            (BaseGate::I, []) => {}
            (BaseGate::H, []) => self.hadamard(t),
            (BaseGate::S, []) => self.phase(t),
            (BaseGate::Sdg, []) => {
                self.phase(t);
                self.phase(t);
                self.phase(t);
            }
            (BaseGate::X, []) => self.x(t),
            (BaseGate::Y, []) => self.y(t),
            (BaseGate::Z, []) => self.z(t),
            (BaseGate::P, []) if g.param == 1 => self.z(t),
            (BaseGate::P, []) if g.param == 2 => self.phase(t),
            (BaseGate::X, [c]) => self.cnot(*c as usize, t),
            (BaseGate::Z, [c]) => self.cz(*c as usize, t),
            (BaseGate::P, [c]) if g.param == 1 => self.cz(*c as usize, t),
            (BaseGate::Y, [c]) => {
                // CY = Sdg(t) · CX(c,t) · S(t)
                self.phase(t);
                self.cnot(*c as usize, t);
                self.phase(t);
                self.phase(t);
                self.phase(t);
            }
            (base, ctrls) => {
                return Err(NotClifford(format!("{base:?} with controls {ctrls:?} param {}", g.param)));
            }
        }
        Ok(())
    }

    // --- measurement (CHP) ---
    fn g_exp(x1: u64, z1: u64, x2: u64, z2: u64) -> i32 {
        let (x2, z2) = (x2 as i32, z2 as i32);
        match (x1, z1) {
            (0, 0) => 0,
            (1, 1) => z2 - x2,
            (1, 0) => z2 * (2 * x2 - 1),
            _ => x2 * (1 - 2 * z2),
        }
    }
    fn rowsum(&mut self, h: usize, i: usize) {
        let mut s = 2 * self.r[h] as i32 + 2 * self.r[i] as i32;
        for j in 0..self.n {
            s += Self::g_exp(self.xget(i, j), self.zget(i, j), self.xget(h, j), self.zget(h, j));
        }
        self.r[h] = if s.rem_euclid(4) == 2 { 1 } else { 0 };
        for w in 0..self.words {
            self.xs[h * self.words + w] ^= self.xs[i * self.words + w];
            self.zs[h * self.words + w] ^= self.zs[i * self.words + w];
        }
    }
    fn copy_row(&mut self, dst: usize, src: usize) {
        for w in 0..self.words {
            self.xs[dst * self.words + w] = self.xs[src * self.words + w];
            self.zs[dst * self.words + w] = self.zs[src * self.words + w];
        }
        self.r[dst] = self.r[src];
    }
    fn zero_row(&mut self, row: usize) {
        for w in 0..self.words {
            self.xs[row * self.words + w] = 0;
            self.zs[row * self.words + w] = 0;
        }
        self.r[row] = 0;
    }

    /// Measure qubit `a` in the computational basis, collapsing the state. Returns
    /// the outcome bit. `rng` supplies randomness for indeterminate outcomes.
    pub fn measure(&mut self, a: usize, rng: &mut u64) -> u8 {
        let n = self.n;
        // is there a stabilizer that anticommutes with Z_a? (x_{p,a}=1, p in n..2n)
        let mut p = None;
        for row in n..2 * n {
            if self.xget(row, a) == 1 {
                p = Some(row);
                break;
            }
        }
        if let Some(p) = p {
            // random outcome
            for i in 0..2 * n {
                if i != p && self.xget(i, a) == 1 {
                    self.rowsum(i, p);
                }
            }
            self.copy_row(p - n, p); // destabilizer <- old stabilizer
            self.zero_row(p);
            self.set_z(p, a, 1);
            let bit = (splitmix64(rng) & 1) as u8;
            self.r[p] = bit;
            bit
        } else {
            // determined outcome: accumulate into scratch row 2n
            let scratch = 2 * n;
            self.zero_row(scratch);
            for i in 0..n {
                if self.xget(i, a) == 1 {
                    self.rowsum(scratch, i + n);
                }
            }
            self.r[scratch]
        }
    }
}

// ===========================================================================
// Runner + sampler
// ===========================================================================

/// Run a Clifford `circuit` on the stabilizer tableau. Errors on a non-Clifford
/// gate (the caller should fall back to a universal backend).
pub fn run(circuit: &Circuit) -> Result<Tableau, NotClifford> {
    let mut t = Tableau::identity(circuit.n_qubits as usize);
    for g in &circuit.ops {
        t.apply(g)?;
    }
    Ok(t)
}

/// One shot: measure every qubit of a fresh copy, packing the outcome into a `u64`
/// (`n ≤ 64`). Bit `i` = qubit `i`, matching the dense simulator's index order.
fn shot(base: &Tableau, rng: &mut u64) -> u64 {
    let mut t = base.clone();
    let mut out = 0u64;
    for q in 0..t.n.min(64) {
        if t.measure(q, rng) == 1 {
            out |= 1u64 << q;
        }
    }
    out
}

/// Sample `shots` measurements of a Clifford circuit as `(outcome, count)` pairs
/// (outcomes as `u64`, `n ≤ 64`), sorted by count descending. Deterministic in
/// `(seed, shots)`.
pub fn sample(circuit: &Circuit, seed: u64, shots: u64) -> Result<Vec<(u64, u64)>, NotClifford> {
    let base = run(circuit)?;
    let mut rng = seed ^ 0x5DEE_CE66_D3A9_1B2C;
    let mut counts: std::collections::HashMap<u64, u64> = std::collections::HashMap::new();
    for _ in 0..shots {
        *counts.entry(shot(&base, &mut rng)).or_insert(0) += 1;
    }
    let mut v: Vec<(u64, u64)> = counts.into_iter().collect();
    v.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
    Ok(v)
}

/// The number of qubits the stabilizer path can handle for a given circuit — the
/// cascade's scaling story (`O(n²)` per gate, no `2ⁿ` state vector).
pub fn max_practical_qubits() -> usize {
    4096
}

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

    fn ghz(n: u8) -> Circuit {
        let mut c = Circuit::new(n);
        c.h(0);
        for i in 1..n {
            c.cx(i - 1, i);
        }
        c
    }

    #[test]
    fn zero_state_is_deterministic_zero() {
        let c = Circuit::new(4);
        let s = sample(&c, 1, 200).unwrap();
        assert_eq!(s.len(), 1);
        assert_eq!(s[0].0, 0);
        assert_eq!(s[0].1, 200);
    }

    #[test]
    fn ghz_samples_all_or_nothing() {
        let s = sample(&ghz(5), 42, 4000).unwrap();
        // only |00000> and |11111> appear, each ~half
        assert_eq!(s.len(), 2, "GHZ has exactly two outcomes");
        let outcomes: Vec<u64> = s.iter().map(|(o, _)| *o).collect();
        assert!(outcomes.contains(&0) && outcomes.contains(&0b11111));
        for (_, cnt) in &s {
            assert!((*cnt as i64 - 2000).abs() < 250, "GHZ halves should be ~even: {cnt}");
        }
    }

    #[test]
    fn matches_dense_support_on_a_clifford_circuit() {
        // a non-trivial Clifford circuit: every sampled outcome must have nonzero
        // dense probability, and the count of distinct outcomes must match.
        let mut c = Circuit::new(4);
        c.h(0);
        c.h(1);
        c.cx(0, 2);
        c.cx(1, 3);
        c.s(2);
        c.cz(0, 1);
        let stab = sample(&c, 7, 8000).unwrap();
        let dense = c.simulate().unwrap();
        let probs = dense.prob_weights();
        let total: i128 = probs.iter().sum();
        // support match
        for (o, _) in &stab {
            assert!(probs[*o as usize] > 0, "stabilizer outcome {o} has zero dense probability");
        }
        let dense_support = probs.iter().filter(|&&p| p > 0).count();
        assert_eq!(stab.len(), dense_support, "distinct-outcome count must match dense");
        // frequencies roughly uniform over the support (Clifford ⇒ flat)
        let expect = 8000.0 / dense_support as f64;
        for (_, cnt) in &stab {
            assert!((*cnt as f64 - expect).abs() < expect * 0.35, "freq off: {cnt} vs {expect}");
        }
        let _ = total;
    }

    #[test]
    fn scales_past_the_dense_wall() {
        // 200-qubit GHZ — 2^200 amplitudes is impossible densely; here it's instant,
        // and every shot must be all-0 or all-1 (perfect parity correlation).
        let n = 200u8;
        let c = ghz(n);
        let base = run(&c).unwrap();
        let mut rng = 123u64;
        for _ in 0..50 {
            let mut t = base.clone();
            let first = t.measure(0, &mut rng);
            for q in 1..n as usize {
                assert_eq!(t.measure(q, &mut rng), first, "GHZ parity broken at large n");
            }
        }
    }

    #[test]
    fn deterministic_in_seed() {
        let a = sample(&ghz(6), 99, 500).unwrap();
        let b = sample(&ghz(6), 99, 500).unwrap();
        assert_eq!(a, b, "same seed must give the same histogram");
    }

    #[test]
    fn rejects_non_clifford() {
        let mut c = Circuit::new(2);
        c.h(0);
        c.t(0); // non-Clifford
        assert!(run(&c).is_err());
        // Toffoli (2 controls) also rejected
        let mut c2 = Circuit::new(3);
        c2.ccx(0, 1, 2);
        assert!(run(&c2).is_err());
    }
}