helpers_codec/common/

sampling.rs

// MIT License
//
// Copyright (c) 2026 Raja Lehtihet & Wael El Oraiby
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! Example sampling helpers backed by `rand_core::CryptoRng`.
//!
//! These routines provide simple building blocks for experimentation:
//! - uniform sampling in `Z_q`,
//! - centered-binomial noise sampling,
//! - discrete Gaussian-like sampling via Box-Muller,
//! - explicit rejection sampling helpers.

use core::f64::consts::TAU;

use rand_core::CryptoRng;

use nc_polynomial::{PolynomialError, RingContext, RingElem};

/// Errors returned by example sampling helpers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SamplingError {
    /// Rejection sampling requires a positive upper bound.
    InvalidUpperBound,
    /// Centered binomial parameter must be positive.
    InvalidEta,
    /// Gaussian sigma must be finite and positive.
    InvalidSigma,
    /// Ring element construction or ring arithmetic failed.
    Polynomial(PolynomialError),
}

impl core::fmt::Display for SamplingError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::InvalidUpperBound => write!(f, "invalid upper bound, expected > 0"),
            Self::InvalidEta => write!(f, "invalid eta, expected eta >= 1"),
            Self::InvalidSigma => write!(f, "invalid sigma, expected finite sigma > 0"),
            Self::Polynomial(err) => write!(f, "polynomial operation failed: {err}"),
        }
    }
}

impl std::error::Error for SamplingError {}

impl From<PolynomialError> for SamplingError {
    fn from(value: PolynomialError) -> Self {
        Self::Polynomial(value)
    }
}

/// Samples a uniform value from `[0, upper_bound)` using rejection sampling.
pub fn sample_rejection_u64_below_example<R: CryptoRng>(
    upper_bound: u64,
    rng: &mut R,
) -> Result<u64, SamplingError> {
    if upper_bound == 0 {
        return Err(SamplingError::InvalidUpperBound);
    }

    // Rejection-sampling setup:
    // - `range` is the cardinality of `u64` outputs.
    // - `zone` is the largest multiple of `upper_bound` inside `range`.
    // Any random value below `zone` can be mapped with `% upper_bound`
    // without modulo bias.
    let range = (u64::MAX as u128) + 1;
    let ub = upper_bound as u128;
    let zone = range - (range % ub);

    loop {
        // Draw and accept only values in the unbiased zone.
        let candidate = rng.next_u64() as u128;
        if candidate < zone {
            return Ok((candidate % ub) as u64);
        }
    }
}

/// Samples one polynomial with coefficients uniformly random in `Z_q`.
pub fn sample_uniform_poly_example<R: CryptoRng>(
    ctx: &RingContext,
    rng: &mut R,
) -> Result<RingElem, SamplingError> {
    let mut coeffs = vec![0_u64; ctx.max_degree() + 1];
    for coeff in &mut coeffs {
        // Independent uniform sampling for each coefficient.
        *coeff = sample_rejection_u64_below_example(ctx.modulus(), rng)?;
    }
    // `ctx.element(...)` canonicalizes modulo the ring polynomial.
    Ok(ctx.element(&coeffs)?)
}

/// Samples one polynomial with centered-binomial noise of parameter `eta`.
///
/// Each coefficient is generated as `sum(a_i) - sum(b_i)` over `eta` Bernoulli bits,
/// then mapped into `Z_q`.
pub fn sample_cbd_poly_example<R: CryptoRng>(
    ctx: &RingContext,
    eta: u8,
    rng: &mut R,
) -> Result<RingElem, SamplingError> {
    if eta == 0 {
        return Err(SamplingError::InvalidEta);
    }

    let mut coeffs = vec![0_u64; ctx.max_degree() + 1];
    for coeff in &mut coeffs {
        // CBD(eta): difference between two eta-bit Hamming weights.
        let mut a = 0_i64;
        let mut b = 0_i64;
        for _ in 0..eta {
            a += (rng.next_u32() & 1) as i64;
            b += (rng.next_u32() & 1) as i64;
        }
        // Map signed noise into `[0, q)` representation.
        *coeff = wrap_signed_to_modulus(a - b, ctx.modulus());
    }

    Ok(ctx.element(&coeffs)?)
}

/// Samples one polynomial with approximate discrete Gaussian coefficients.
///
/// This example helper uses a Box-Muller normal sampler and rounds to nearest integer.
pub fn sample_discrete_gaussian_poly_example<R: CryptoRng>(
    ctx: &RingContext,
    sigma: f64,
    rng: &mut R,
) -> Result<RingElem, SamplingError> {
    if !(sigma.is_finite() && sigma > 0.0) {
        return Err(SamplingError::InvalidSigma);
    }

    let mut coeffs = vec![0_u64; ctx.max_degree() + 1];
    for coeff in &mut coeffs {
        // Sample one integer noise value and wrap it into `Z_q`.
        let sampled = sample_box_muller_integer(rng, sigma);
        *coeff = wrap_signed_to_modulus(sampled, ctx.modulus());
    }

    Ok(ctx.element(&coeffs)?)
}

fn sample_box_muller_integer<R: CryptoRng>(rng: &mut R, sigma: f64) -> i64 {
    // Box-Muller transform:
    // z0 = sqrt(-2 ln u1) * cos(2 pi u2), where u1/u2 are uniform(0,1).
    // Then scale by sigma and round to nearest integer.
    let mut u1 = sample_unit_open_closed(rng);
    while u1 <= 0.0 {
        u1 = sample_unit_open_closed(rng);
    }
    let u2 = sample_unit_open_closed(rng);

    let radius = (-2.0 * u1.ln()).sqrt();
    let z0 = radius * (TAU * u2).cos();
    (z0 * sigma).round() as i64
}

fn sample_unit_open_closed<R: CryptoRng>(rng: &mut R) -> f64 {
    // Convert 53 random bits into a floating-point fraction in (0, 1).
    const SCALE: f64 = 1.0 / ((1_u64 << 53) as f64);
    let raw = rng.next_u64() >> 11;
    let mut out = (raw as f64) * SCALE;
    // Clamp away exact endpoints to keep logarithm and cosine stable.
    if out <= 0.0 {
        out = f64::from_bits(1);
    }
    if out >= 1.0 {
        out = 1.0 - f64::EPSILON;
    }
    out
}

fn wrap_signed_to_modulus(value: i64, modulus: u64) -> u64 {
    // Normalize signed integer into canonical residue class `[0, q)`.
    let m = modulus as i128;
    let reduced = ((value as i128 % m) + m) % m;
    reduced as u64
}

#[cfg(test)]
mod tests {
    use super::{
        SamplingError, sample_cbd_poly_example, sample_discrete_gaussian_poly_example,
        sample_rejection_u64_below_example, sample_uniform_poly_example,
    };
    use core::convert::Infallible;
    use nc_polynomial::RingContext;
    use rand_core::{TryCryptoRng, TryRng};

    #[derive(Debug, Clone)]
    struct TestRng {
        state: u64,
    }

    impl TestRng {
        fn new(seed: u64) -> Self {
            Self { state: seed }
        }

        fn step(&mut self) -> u64 {
            let mut x = self.state;
            x ^= x >> 12;
            x ^= x << 25;
            x ^= x >> 27;
            self.state = x;
            x.wrapping_mul(0x2545_F491_4F6C_DD1D)
        }
    }

    impl TryRng for TestRng {
        type Error = Infallible;

        fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
            Ok(self.step() as u32)
        }

        fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
            Ok(self.step())
        }

        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
            let mut offset = 0;
            while offset < dest.len() {
                let word = self.step().to_le_bytes();
                let chunk = (dest.len() - offset).min(8);
                dest[offset..offset + chunk].copy_from_slice(&word[..chunk]);
                offset += chunk;
            }
            Ok(())
        }
    }

    impl TryCryptoRng for TestRng {}

    fn ctx() -> RingContext {
        RingContext::from_parts(8, 998_244_353, &[1, 0, 0, 0, 0, 0, 0, 0, 1], 3)
            .expect("context should build")
    }

    #[test]
    fn rejection_sampling_respects_bounds() {
        let mut rng = TestRng::new(7);
        for bound in [1_u64, 2, 3, 17, 1_000_000, 998_244_353] {
            for _ in 0..200 {
                // Every output must stay in `[0, bound)`.
                let sampled = sample_rejection_u64_below_example(bound, &mut rng)
                    .expect("sampling should work");
                assert!(sampled < bound);
            }
        }
    }

    #[test]
    fn rejection_sampling_rejects_zero_bound() {
        let mut rng = TestRng::new(11);
        let err = sample_rejection_u64_below_example(0, &mut rng).expect_err("expected error");
        assert_eq!(err, SamplingError::InvalidUpperBound);
    }

    #[test]
    fn uniform_poly_sampling_outputs_valid_ring_element() {
        let ctx = ctx();
        let mut rng = TestRng::new(13);
        let sampled = sample_uniform_poly_example(&ctx, &mut rng).expect("sampling should work");

        // Shape and domain checks.
        assert_eq!(sampled.coefficients().len(), ctx.max_degree() + 1);
        assert!(sampled.coefficients().iter().all(|&c| c < ctx.modulus()));
    }

    #[test]
    fn cbd_poly_sampling_outputs_small_noise_mod_q() {
        let ctx = ctx();
        let mut rng = TestRng::new(17);
        let sampled = sample_cbd_poly_example(&ctx, 3, &mut rng).expect("sampling should work");

        // Interpret coefficient residues as signed values near zero and check bounds.
        let eta = 3_i64;
        for &coeff in sampled.coefficients() {
            let signed = if coeff > ctx.modulus() / 2 {
                coeff as i64 - ctx.modulus() as i64
            } else {
                coeff as i64
            };
            assert!(signed >= -eta && signed <= eta);
        }
    }

    #[test]
    fn cbd_sampling_rejects_eta_zero() {
        let ctx = ctx();
        let mut rng = TestRng::new(19);
        let err = sample_cbd_poly_example(&ctx, 0, &mut rng).expect_err("expected error");
        assert_eq!(err, SamplingError::InvalidEta);
    }

    #[test]
    fn gaussian_sampling_accepts_positive_sigma() {
        let ctx = ctx();
        let mut rng = TestRng::new(23);
        let sampled = sample_discrete_gaussian_poly_example(&ctx, 2.5, &mut rng)
            .expect("sampling should work");
        assert_eq!(sampled.coefficients().len(), ctx.max_degree() + 1);
    }

    #[test]
    fn gaussian_sampling_rejects_invalid_sigma() {
        let ctx = ctx();
        let mut rng = TestRng::new(29);
        let err =
            sample_discrete_gaussian_poly_example(&ctx, 0.0, &mut rng).expect_err("expected error");
        assert_eq!(err, SamplingError::InvalidSigma);
    }
}