qzkp 0.1.0

//! BB84 quantum key distribution protocol simulation.
//!
//! Implements the core BB84 protocol: Alice encodes random bits in random bases,
//! Bob measures in random bases, they sift to matching bases, and estimate the
//! error rate on a sample to detect eavesdropping.

use crate::quantum::QuantumCircuit;
use rand::Rng;

/// BB84 quantum key distribution simulator.
///
/// # Example
///
/// ```
/// use qzkp::bb84::BB84QKD;
///
/// let bb84 = BB84QKD::new(100, 0.02);
/// let (key_alice, key_bob, error_rate) = bb84.generate_key(&mut rand::thread_rng());
/// assert!(error_rate < 0.11);
/// ```
pub struct BB84QKD {
    /// Number of qubits transmitted.
    pub n_bits: usize,
    /// Channel noise level (depolarizing error probability per qubit).
    pub noise_level: f64,
}

impl BB84QKD {
    /// Create a new BB84 instance.
    ///
    /// * `n_bits` -- number of qubits to transmit
    /// * `noise_level` -- depolarizing noise probability per qubit
    pub fn new(n_bits: usize, noise_level: f64) -> Self {
        BB84QKD {
            n_bits,
            noise_level,
        }
    }

    fn encode_bit(&self, bit: u8, basis: u8) -> QuantumCircuit {
        let mut qc = QuantumCircuit::new(1, self.noise_level);
        if bit == 1 {
            qc.x(0);
        }
        if basis == 1 {
            qc.h(0);
        }
        qc
    }

    fn measure_bit(&self, qc: &mut QuantumCircuit, basis: u8, rng: &mut impl Rng) -> u8 {
        if basis == 1 {
            qc.h(0);
        }
        qc.measure(0, rng)
    }

    /// Run the BB84 protocol and return `(alice_key, bob_key, error_rate)`.
    ///
    /// The returned keys are the sifted keys (after basis reconciliation and
    /// error sampling). The error rate is estimated on half the sifted bits.
    pub fn generate_key(&self, rng: &mut impl Rng) -> (Vec<u8>, Vec<u8>, f64) {
        let alice_bits: Vec<u8> = (0..self.n_bits).map(|_| rng.gen_range(0..2)).collect();
        let alice_bases: Vec<u8> = (0..self.n_bits).map(|_| rng.gen_range(0..2)).collect();
        let bob_bases: Vec<u8> = (0..self.n_bits).map(|_| rng.gen_range(0..2)).collect();

        let mut bob_results = Vec::new();
        for i in 0..self.n_bits {
            let mut qc = self.encode_bit(alice_bits[i], alice_bases[i]);
            qc.apply_noise(rng);
            let result = self.measure_bit(&mut qc, bob_bases[i], rng);
            bob_results.push(result);
        }

        let mut sifted_key_alice = Vec::new();
        let mut sifted_key_bob = Vec::new();
        for i in 0..self.n_bits {
            if alice_bases[i] == bob_bases[i] {
                sifted_key_alice.push(alice_bits[i]);
                sifted_key_bob.push(bob_results[i]);
            }
        }

        let sample_size = sifted_key_alice.len() / 2;
        let error_count = (0..sample_size)
            .filter(|&i| sifted_key_alice[i] != sifted_key_bob[i])
            .count();

        let error_rate = if sample_size > 0 {
            error_count as f64 / sample_size as f64
        } else {
            0.0
        };

        let final_key_alice = sifted_key_alice[sample_size..].to_vec();
        let final_key_bob = sifted_key_bob[sample_size..].to_vec();

        (final_key_alice, final_key_bob, error_rate)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::rngs::StdRng;
    use rand::SeedableRng;

    fn seeded_rng() -> StdRng {
        StdRng::seed_from_u64(123)
    }

    #[test]
    fn low_noise_yields_matching_keys() {
        let bb84 = BB84QKD::new(200, 0.0);
        let mut rng = seeded_rng();
        let (key_a, key_b, error_rate) = bb84.generate_key(&mut rng);
        assert_eq!(key_a, key_b, "zero-noise keys should match exactly");
        assert!((error_rate - 0.0).abs() < 1e-10);
    }

    #[test]
    fn keys_have_reasonable_length() {
        let bb84 = BB84QKD::new(100, 0.02);
        let mut rng = seeded_rng();
        let (key_a, _key_b, _error_rate) = bb84.generate_key(&mut rng);
        // ~50% basis match, then half used for sampling => ~25% of n_bits
        assert!(key_a.len() >= 10, "sifted key too short: {}", key_a.len());
        assert!(
            key_a.len() <= 60,
            "sifted key unexpectedly long: {}",
            key_a.len()
        );
    }

    #[test]
    fn low_noise_error_rate_below_threshold() {
        let bb84 = BB84QKD::new(200, 0.02);
        let mut rng = seeded_rng();
        let (_ka, _kb, error_rate) = bb84.generate_key(&mut rng);
        assert!(
            error_rate < 0.11,
            "error rate {} too high for low noise",
            error_rate
        );
    }

    #[test]
    fn high_noise_degrades_keys() {
        let bb84 = BB84QKD::new(200, 0.3);
        let mut rng = seeded_rng();
        let (key_a, key_b, error_rate) = bb84.generate_key(&mut rng);
        // With 30% noise, many bits should differ
        let mismatches = key_a.iter().zip(&key_b).filter(|(a, b)| a != b).count();
        assert!(
            mismatches > 0 || error_rate > 0.05,
            "high noise should cause errors"
        );
    }
}