1use rand::Rng;
  2
  3#[cfg(not(feature = "std"))]
  4use num::Float;
  5
  6/// Samples an integer from {0, ..., 18} according to the distribution χ, which is close to
  7/// the half-Gaussian distribution on the natural numbers with mean 0 and standard deviation
  8/// equal to sigma_max.
  9fn base_sampler(bytes: [u8; 9]) -> i16 {
 10    const RCDT: [u128; 18] = [
 11        3024686241123004913666,
 12        1564742784480091954050,
 13        636254429462080897535,
 14        199560484645026482916,
 15        47667343854657281903,
 16        8595902006365044063,
 17        1163297957344668388,
 18        117656387352093658,
 19        8867391802663976,
 20        496969357462633,
 21        20680885154299,
 22        638331848991,
 23        14602316184,
 24        247426747,
 25        3104126,
 26        28824,
 27        198,
 28        1,
 29    ];
 30
 31    let mut tmp = bytes.to_vec();
 32    tmp.extend_from_slice(&[0u8; 7]);
 33    tmp.reverse();
 34    let u = u128::from_be_bytes(tmp.try_into().expect("should have length 16"));
 35    RCDT.into_iter().filter(|r| u < *r).count() as i16
 36}
 37
 38/// Computes an integer approximation of 2^63 * ccs * exp(-x).
 39fn approx_exp(x: f64, ccs: f64) -> u64 {
 40    // The constants C are used to approximate exp(-x); these
 41    // constants are taken from FACCT (up to a scaling factor
 42    // of 2^63):
 43    //   https://eprint.iacr.org/2018/1234
 44    //   https://github.com/raykzhao/gaussian
 45    const C: [u64; 13] = [
 46        0x00000004741183A3u64,
 47        0x00000036548CFC06u64,
 48        0x0000024FDCBF140Au64,
 49        0x0000171D939DE045u64,
 50        0x0000D00CF58F6F84u64,
 51        0x000680681CF796E3u64,
 52        0x002D82D8305B0FEAu64,
 53        0x011111110E066FD0u64,
 54        0x0555555555070F00u64,
 55        0x155555555581FF00u64,
 56        0x400000000002B400u64,
 57        0x7FFFFFFFFFFF4800u64,
 58        0x8000000000000000u64,
 59    ];
 60
 61    let mut z: u64;
 62    let mut y: u64;
 63    let twoe63 = 1u64 << 63;
 64
 65    y = C[0];
 66    z = f64::floor(x * (twoe63 as f64)) as u64;
 67    for cu in C.iter().skip(1) {
 68        let zy = (z as u128) * (y as u128);
 69        y = cu - ((zy >> 63) as u64);
 70    }
 71
 72    z = f64::floor((twoe63 as f64) * ccs) as u64;
 73
 74    (((z as u128) * (y as u128)) >> 63) as u64
 75}
 76
 77/// A random bool that is true with probability ≈ ccs · exp(-x).
 78fn ber_exp<R: Rng>(x: f64, ccs: f64, rng: &mut R) -> bool {
 79    const LN2: f64 = core::f64::consts::LN_2;
 80    const ILN2: f64 = 1.0 / LN2;
 81    let s = f64::floor(x * ILN2);
 82    let r = x - s * LN2;
 83    let s = (s as u64).min(63);
 84    let z = ((approx_exp(r, ccs) << 1) - 1) >> s;
 85
 86    let mut w = 0_i32;
 87    for i in (0..=56).rev().step_by(8) {
 88        let mut dest = [0_u8; 1];
 89        rng.fill_bytes(&mut dest);
 90        let p = u8::from_be_bytes(dest);
 91        w = (p as i32) - (z >> i & 0xFF) as i32;
 92        if w != 0 {
 93            break;
 94        }
 95    }
 96    w < 0
 97}
 98
 99/// Samples an integer from the Gaussian distribution with given mean (mu) and standard deviation
100/// (sigma).
101pub(crate) fn sampler_z<R: Rng>(mu: f64, sigma: f64, sigma_min: f64, rng: &mut R) -> i16 {
102    const SIGMA_MAX: f64 = 1.8205;
103    const INV_2SIGMA_MAX_SQ: f64 = 1f64 / (2f64 * SIGMA_MAX * SIGMA_MAX);
104    let isigma = 1f64 / sigma;
105    let dss = 0.5f64 * isigma * isigma;
106    let s = f64::floor(mu);
107    let r = mu - s;
108    let ccs = sigma_min * isigma;
109    loop {
110        let mut dest = [0_u8; 9];
111        rng.fill_bytes(&mut dest);
112        let z0 = base_sampler(dest);
113
114        let mut dest = [0_u8; 1];
115        rng.fill_bytes(&mut dest);
116        let random_byte: u8 = dest[0];
117
118        let b = (random_byte & 1) as i16;
119        let z = b + (2 * b - 1) * z0;
120        let zf_min_r = (z as f64) - r;
121        let x = zf_min_r * zf_min_r * dss - (z0 * z0) as f64 * INV_2SIGMA_MAX_SQ;
122
123        if ber_exp(x, ccs, rng) {
124            return z + (s as i16);
125        }
126    }
127}
128
129#[cfg(all(test, feature = "std"))]
130mod test {
131    use super::approx_exp;
132
133    #[test]
134    fn test_approx_exp() {
135        let precision = 1u64 << 14;
136        // known answers were generated with the following sage script:
137        //```sage
138        // num_samples = 10
139        // precision = 200
140        // R = Reals(precision)
141        //
142        // print(f"let kats : [(f64, f64, u64);{num_samples}] = [")
143        // for i in range(num_samples):
144        //     x = RDF.random_element(0.0, 0.693147180559945)
145        //     ccs = RDF.random_element(0.0, 1.0)
146        //     res = round(2^63 * R(ccs) * exp(R(-x)))
147        //     print(f"({x}, {ccs}, {res}),")
148        // print("];")
149        // ```
150        let kats: [(f64, f64, u64); 10] = [
151            (0.2314993926072656, 0.8148006314615972, 5962140072160879737),
152            (0.2648875572812225, 0.12769669655309035, 903712282351034505),
153            (0.11251957513682391, 0.9264611470305881, 7635725498677341553),
154            (0.04353439307256617, 0.5306497137523327, 4685877322232397936),
155            (0.41834495299784347, 0.879438856118578, 5338392138535350986),
156            (0.32579398973228557, 0.16513412873289002, 1099603299296456803),
157            (0.5939508073919817, 0.029776019144967303, 151637565622779016),
158            (0.2932367999399056, 0.37123847662857923, 2553827649386670452),
159            (0.5005699297417507, 0.31447208863888976, 1758235618083658825),
160            (0.4876437338498085, 0.6159515298936868, 3488632981903743976),
161        ];
162        for (x, ccs, answer) in kats {
163            let difference = (answer as i128) - (approx_exp(x, ccs) as i128);
164            assert!(
165                (difference * difference) as u64 <= precision * precision,
166                "answer: {answer} versus approximation: {}\ndifference: {} whereas precision: {}",
167                approx_exp(x, ccs),
168                difference,
169                precision
170            );
171        }
172    }
173}