Skip to main content

dcrypt_algorithms/poly/
sampling.rs

1//! sampling.rs - Cryptographic sampling algorithms
2
3use super::params::Modulus;
4use super::polynomial::Polynomial;
5use crate::error::{Error, Result};
6use dcrypt_internal::random::{try_fill_bytes_zeroing_on_error, CryptoRng, RngCore};
7use dcrypt_internal::zeroing::Zeroizing;
8
9/// Trait for sampling polynomials uniformly at random
10pub trait UniformSampler<M: Modulus> {
11    /// Samples a polynomial with coefficients uniformly random in [0, Q-1]
12    fn sample_uniform<R: RngCore + CryptoRng>(rng: &mut R) -> Result<Polynomial<M>>;
13}
14
15/// Trait for sampling polynomials from a Centered Binomial Distribution (CBD)
16pub trait CbdSampler<M: Modulus> {
17    /// Samples a polynomial with coefficients from CBD(eta)
18    fn sample_cbd<R: RngCore + CryptoRng>(rng: &mut R, eta: u8) -> Result<Polynomial<M>>;
19}
20
21/// Default implementation of cryptographic samplers
22pub struct DefaultSamplers;
23
24impl<M: Modulus> UniformSampler<M> for DefaultSamplers {
25    fn sample_uniform<R: RngCore + CryptoRng>(rng: &mut R) -> Result<Polynomial<M>> {
26        let mut poly = Polynomial::<M>::zero();
27        let q = M::Q;
28
29        // Handle different modulus sizes
30        if q <= (1 << 16) {
31            // For small moduli, use rejection sampling with u16
32            sample_uniform_small::<M, R>(rng, &mut poly)?;
33        } else if q <= (1 << 24) {
34            // For medium moduli, use rejection sampling with u32
35            sample_uniform_medium::<M, R>(rng, &mut poly)?;
36        } else {
37            // For large moduli up to 2^31
38            sample_uniform_large::<M, R>(rng, &mut poly)?;
39        }
40
41        Ok(poly)
42    }
43}
44
45/// Rejection sampling for small moduli (Q <= 2^16)
46fn sample_uniform_small<M: Modulus, R: RngCore + CryptoRng>(
47    rng: &mut R,
48    poly: &mut Polynomial<M>,
49) -> Result<()> {
50    let q = M::Q;
51    let n = M::N;
52
53    // Find the largest multiple of q that fits in u16
54    let threshold = ((1u32 << 16) / q) * q;
55
56    for i in 0..n {
57        loop {
58            let mut bytes = Zeroizing::new([0u8; 2]);
59            try_fill_bytes_zeroing_on_error(rng, &mut bytes[..])?;
60            let sample = u32::from(bytes[0]) | (u32::from(bytes[1]) << 8);
61
62            // Rejection sampling for uniform distribution
63            if sample < threshold {
64                poly.coeffs[i] = sample % q;
65                break;
66            }
67        }
68    }
69
70    Ok(())
71}
72
73/// Rejection sampling for medium moduli (2^16 < Q <= 2^24)
74fn sample_uniform_medium<M: Modulus, R: RngCore + CryptoRng>(
75    rng: &mut R,
76    poly: &mut Polynomial<M>,
77) -> Result<()> {
78    let q = M::Q;
79    let n = M::N;
80
81    // Use 3 bytes for sampling
82    let threshold = ((1u32 << 24) / q) * q;
83
84    for i in 0..n {
85        loop {
86            let mut bytes = Zeroizing::new([0u8; 3]);
87            try_fill_bytes_zeroing_on_error(rng, &mut bytes[..])?;
88            let sample =
89                u32::from(bytes[0]) | (u32::from(bytes[1]) << 8) | (u32::from(bytes[2]) << 16);
90
91            if sample < threshold {
92                poly.coeffs[i] = sample % q;
93                break;
94            }
95        }
96    }
97
98    Ok(())
99}
100
101/// Rejection sampling for large moduli (2^24 < Q <= 2^31)
102fn sample_uniform_large<M: Modulus, R: RngCore + CryptoRng>(
103    rng: &mut R,
104    poly: &mut Polynomial<M>,
105) -> Result<()> {
106    let q = M::Q;
107    let n = M::N;
108
109    // Use full u32 with MSB clear to ensure < 2^31
110    let threshold = ((1u32 << 31) / q) * q;
111
112    for i in 0..n {
113        loop {
114            let mut bytes = Zeroizing::new([0u8; 4]);
115            try_fill_bytes_zeroing_on_error(rng, &mut bytes[..])?;
116            bytes[3] &= 0x7F; // Clear MSB
117            let sample = u32::from(bytes[0])
118                | (u32::from(bytes[1]) << 8)
119                | (u32::from(bytes[2]) << 16)
120                | (u32::from(bytes[3]) << 24);
121
122            if sample < threshold {
123                poly.coeffs[i] = sample % q;
124                break;
125            }
126        }
127    }
128
129    Ok(())
130}
131
132impl<M: Modulus> CbdSampler<M> for DefaultSamplers {
133    fn sample_cbd<R: RngCore + CryptoRng>(rng: &mut R, eta: u8) -> Result<Polynomial<M>> {
134        if eta == 0 || eta > 16 {
135            return Err(Error::Parameter {
136                name: "CBD sampling".into(),
137                reason: format!("eta must be in range [1, 16], got {}", eta).into(),
138            });
139        }
140
141        let mut poly = Polynomial::<M>::zero();
142        let n = M::N;
143        let q = M::Q;
144
145        // CBD(eta): sample 2*eta bits, compute sum of first eta bits minus sum of second eta bits
146        let bytes_per_sample = (2 * eta as usize).div_ceil(8); // FIXED: Use div_ceil
147        let mut buffer = Zeroizing::new([0u8; 4]); // Max 32 bits for eta=16
148
149        for i in 0..n {
150            try_fill_bytes_zeroing_on_error(rng, &mut buffer[..bytes_per_sample])?;
151
152            let mut a = 0i32;
153            let mut b = 0i32;
154
155            // Extract eta bits for positive contribution
156            for j in 0..eta {
157                let byte_idx = j as usize / 8;
158                let bit_idx = j as usize % 8;
159                a += ((buffer[byte_idx] >> bit_idx) & 1) as i32;
160            }
161
162            // Extract eta bits for negative contribution
163            for j in 0..eta {
164                let bit_pos = (eta + j) as usize;
165                let byte_idx = bit_pos / 8;
166                let bit_idx = bit_pos % 8;
167                b += ((buffer[byte_idx] >> bit_idx) & 1) as i32;
168            }
169
170            // CBD sample is in range [-eta, eta]
171            let sample = a - b;
172
173            // Convert to [0, q) range
174            poly.coeffs[i] = ((sample + q as i32) % q as i32) as u32;
175        }
176
177        Ok(poly)
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    use dcrypt_internal::random::ChaCha20Rng;
185
186    #[derive(Clone)]
187    struct TestModulus;
188    impl Modulus for TestModulus {
189        const Q: u32 = 3329;
190        const N: usize = 256;
191    }
192
193    #[test]
194    fn test_uniform_sampling() {
195        let mut rng = ChaCha20Rng::from_seed([42u8; 32]);
196        let poly =
197            <DefaultSamplers as UniformSampler<TestModulus>>::sample_uniform(&mut rng).unwrap();
198
199        // Check all coefficients are in valid range
200        for &coeff in poly.as_coeffs_slice() {
201            assert!(coeff < TestModulus::Q);
202        }
203    }
204
205    #[test]
206    fn test_cbd_sampling() {
207        let mut rng = ChaCha20Rng::from_seed([42u8; 32]);
208
209        for eta in 1..=8 {
210            let poly =
211                <DefaultSamplers as CbdSampler<TestModulus>>::sample_cbd(&mut rng, eta).unwrap();
212
213            // Check all coefficients are in valid range
214            for &coeff in poly.as_coeffs_slice() {
215                assert!(coeff < TestModulus::Q);
216            }
217        }
218    }
219
220    #[test]
221    fn test_cbd_distribution() {
222        // Simple statistical test for CBD
223        let mut rng = ChaCha20Rng::from_seed([42u8; 32]);
224        let eta = 2;
225        let num_samples = 10000;
226        let mut histogram = vec![0u32; (2 * eta + 1) as usize];
227
228        for _ in 0..num_samples {
229            let poly =
230                <DefaultSamplers as CbdSampler<TestModulus>>::sample_cbd(&mut rng, eta).unwrap();
231
232            // Check first coefficient distribution
233            let coeff = poly.coeffs[0];
234            let centered = (coeff as i32 + eta as i32) % TestModulus::Q as i32;
235            if centered <= 2 * eta as i32 {
236                histogram[centered as usize] += 1;
237            }
238        }
239
240        // CBD(2) should have distribution:
241        // P(X = -2) = 1/16, P(X = -1) = 4/16, P(X = 0) = 6/16,
242        // P(X = 1) = 4/16, P(X = 2) = 1/16
243        let expected = [625, 2500, 3750, 2500, 625]; // Out of 10000
244
245        // Chi-squared test with reasonable tolerance
246        let mut chi_squared = 0.0;
247        for i in 0..histogram.len() {
248            let observed = histogram[i] as f64;
249            let expected_val = expected[i] as f64;
250            chi_squared += (observed - expected_val).powi(2) / expected_val;
251        }
252
253        // Degrees of freedom = 4, critical value at 0.05 significance ≈ 9.488
254        assert!(
255            chi_squared < 15.0,
256            "Chi-squared test failed: {}",
257            chi_squared
258        );
259    }
260}