chacha.rs

1use std::collections::HashSet;23const CONSTANTS: [u32; 4] = [0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574];4const BUF_WORDS: usize = 64;5const BLOCK_WORDS: usize = 16;67/// An eight-round ChaCha keystream serving as a seekable source of random words.8#[derive(Clone)]9pub struct ChaCha8 {10    key: [u32; 8],11    block_pos: u64,12    buf: [u32; BUF_WORDS],13    index: usize,14}1516impl ChaCha8 {17    /// Builds the generator from a 32-byte key.18    pub fn from_seed(seed: [u8; 32]) -> ChaCha8 {19        let mut key = [0u32; 8];20        for (word, chunk) in key.iter_mut().zip(seed.chunks_exact(4)) {21            *word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);22        }23        ChaCha8 {24            key,25            block_pos: 0,26            buf: [0; BUF_WORDS],27            index: BUF_WORDS,28        }29    }3031    /// Builds the generator by stretching a word seed into a full key.32    pub fn from_u64(seed: u64) -> ChaCha8 {33        let mut state = seed;34        let mut bytes = [0u8; 32];35        for chunk in bytes.chunks_exact_mut(4) {36            chunk.copy_from_slice(&pcg32(&mut state));37        }38        ChaCha8::from_seed(bytes)39    }4041    /// Returns the next 32-bit word of the keystream.42    pub fn next_u32(&mut self) -> u32 {43        if self.index >= BUF_WORDS {44            self.generate_and_set(0);45        }46        let value = self.buf[self.index];47        self.index += 1;48        value49    }5051    /// Returns the next two keystream words joined into 64 bits, low word first.52    pub fn next_u64(&mut self) -> u64 {53        let index = self.index;54        if index < BUF_WORDS - 1 {55            self.index += 2;56            (u64::from(self.buf[index + 1]) << 32) | u64::from(self.buf[index])57        } else if index >= BUF_WORDS {58            self.generate_and_set(2);59            (u64::from(self.buf[1]) << 32) | u64::from(self.buf[0])60        } else {61            let x = u64::from(self.buf[BUF_WORDS - 1]);62            self.generate_and_set(1);63            let y = u64::from(self.buf[0]);64            (y << 32) | x65        }66    }6768    /// Returns a uniform double in the half-open unit interval.69    pub fn unit(&mut self) -> f64 {70        let scale = 1.0 / ((1u64 << 53) as f64);71        scale * ((self.next_u64() >> 11) as f64)72    }7374    /// Returns a fair coin flip.75    pub fn boolean(&mut self) -> bool {76        (self.next_u32() as i32) < 077    }7879    /// Returns an unbiased uniform integer between low and high inclusive.80    pub fn range_i64(&mut self, low: i64, high: i64) -> i64 {81        assert!(low <= high, "range_i64: low > high");82        let range = high.wrapping_sub(low).wrapping_add(1) as u64;83        if range == 0 {84            return self.next_u64() as i64;85        }86        let zone = (range << range.leading_zeros()).wrapping_sub(1);87        loop {88            let (hi, lo) = wide_mul_u64(self.next_u64(), range);89            if lo <= zone {90                return low.wrapping_add(hi as i64);91            }92        }93    }9495    /// Returns an unbiased uniform integer below n.96    pub fn below_u64(&mut self, n: u64) -> u64 {97        assert!(n > 0, "below_u64: empty range");98        let zone = (n << n.leading_zeros()).wrapping_sub(1);99        loop {100            let (hi, lo) = wide_mul_u64(self.next_u64(), n);101            if lo <= zone {102                return hi;103            }104        }105    }106107    /// Shuffles the slice in place by Fisher-Yates.108    pub fn shuffle<T>(&mut self, seq: &mut [T]) {109        for i in (1..seq.len()).rev() {110            let j = self.index_below(i + 1);111            seq.swap(i, j);112        }113    }114115    /// Draws amount distinct indices below length, choosing its strategy by the sizes.116    pub fn sample_indices(&mut self, length: usize, amount: usize) -> Vec<usize> {117        assert!(amount <= length, "sample_indices: amount > length");118        if length > u32::MAX as usize {119            return self.sample_rejection_u64(length as u64, amount as u64);120        }121        let length = length as u32;122        let amount = amount as u32;123        if amount < 163 {124            const C: [[f32; 2]; 2] = [[1.6, 8.0 / 45.0], [10.0, 70.0 / 9.0]];125            let j = if length < 500_000 { 0 } else { 1 };126            let amount_fp = amount as f32;127            let m4 = C[0][j] * amount_fp;128            if amount > 11 && (length as f32) < (C[1][j] + m4) * amount_fp {129                self.sample_inplace(length, amount)130            } else {131                self.sample_floyd(length, amount)132            }133        } else {134            const C: [f32; 2] = [270.0, 330.0 / 9.0];135            let j = if length < 500_000 { 0 } else { 1 };136            if (length as f32) < C[j] * (amount as f32) {137                self.sample_inplace(length, amount)138            } else {139                self.sample_rejection_u32(length, amount)140            }141        }142    }143144    fn index_below(&mut self, ubound: usize) -> usize {145        if ubound <= u32::MAX as usize {146            self.range_u32(0, ubound as u32 - 1) as usize147        } else {148            self.below_u64(ubound as u64) as usize149        }150    }151152    fn range_u32(&mut self, low: u32, high: u32) -> u32 {153        let range = high.wrapping_sub(low).wrapping_add(1);154        if range == 0 {155            return self.next_u32();156        }157        let zone = (range << range.leading_zeros()).wrapping_sub(1);158        loop {159            let (hi, lo) = wide_mul_u32(self.next_u32(), range);160            if lo <= zone {161                return low.wrapping_add(hi);162            }163        }164    }165166    fn sample_floyd(&mut self, length: u32, amount: u32) -> Vec<usize> {167        let floyd_shuffle = amount < 50;168        let mut indices: Vec<u32> = Vec::with_capacity(amount as usize);169        for j in length - amount..length {170            let t = self.range_u32(0, j);171            if floyd_shuffle {172                if let Some(pos) = indices.iter().position(|&x| x == t) {173                    indices.insert(pos, j);174                    continue;175                }176            } else if indices.contains(&t) {177                indices.push(j);178                continue;179            }180            indices.push(t);181        }182        if !floyd_shuffle {183            for i in (1..amount).rev() {184                let j = self.range_u32(0, i);185                indices.swap(i as usize, j as usize);186            }187        }188        indices.into_iter().map(|i| i as usize).collect()189    }190191    fn sample_inplace(&mut self, length: u32, amount: u32) -> Vec<usize> {192        let mut indices: Vec<u32> = (0..length).collect();193        for i in 0..amount {194            let j = self.range_u32(i, length - 1);195            indices.swap(i as usize, j as usize);196        }197        indices.truncate(amount as usize);198        indices.into_iter().map(|i| i as usize).collect()199    }200201    fn sample_rejection_u32(&mut self, length: u32, amount: u32) -> Vec<usize> {202        let ints_to_reject = (u32::MAX - length + 1) % length;203        let zone = u32::MAX - ints_to_reject;204        let mut cache = HashSet::with_capacity(amount as usize);205        let mut indices = Vec::with_capacity(amount as usize);206        for _ in 0..amount {207            let mut pos = self.distr_u32(length, zone);208            while !cache.insert(pos) {209                pos = self.distr_u32(length, zone);210            }211            indices.push(pos as usize);212        }213        indices214    }215216    fn sample_rejection_u64(&mut self, length: u64, amount: u64) -> Vec<usize> {217        let ints_to_reject = (u64::MAX - length + 1) % length;218        let zone = u64::MAX - ints_to_reject;219        let mut cache = HashSet::with_capacity(amount as usize);220        let mut indices = Vec::with_capacity(amount as usize);221        for _ in 0..amount {222            let mut pos = self.distr_u64(length, zone);223            while !cache.insert(pos) {224                pos = self.distr_u64(length, zone);225            }226            indices.push(pos as usize);227        }228        indices229    }230231    fn distr_u32(&mut self, range: u32, zone: u32) -> u32 {232        loop {233            let (hi, lo) = wide_mul_u32(self.next_u32(), range);234            if lo <= zone {235                return hi;236            }237        }238    }239240    fn distr_u64(&mut self, range: u64, zone: u64) -> u64 {241        loop {242            let (hi, lo) = wide_mul_u64(self.next_u64(), range);243            if lo <= zone {244                return hi;245            }246        }247    }248249    fn generate_and_set(&mut self, index: usize) {250        self.generate();251        self.index = index;252    }253254    fn generate(&mut self) {255        for block in 0..4 {256            let counter = self.block_pos.wrapping_add(block as u64);257            let initial = [258                CONSTANTS[0],259                CONSTANTS[1],260                CONSTANTS[2],261                CONSTANTS[3],262                self.key[0],263                self.key[1],264                self.key[2],265                self.key[3],266                self.key[4],267                self.key[5],268                self.key[6],269                self.key[7],270                counter as u32,271                (counter >> 32) as u32,272                0,273                0,274            ];275            let mut state = initial;276            for _ in 0..4 {277                quarter(&mut state, 0, 4, 8, 12);278                quarter(&mut state, 1, 5, 9, 13);279                quarter(&mut state, 2, 6, 10, 14);280                quarter(&mut state, 3, 7, 11, 15);281                quarter(&mut state, 0, 5, 10, 15);282                quarter(&mut state, 1, 6, 11, 12);283                quarter(&mut state, 2, 7, 8, 13);284                quarter(&mut state, 3, 4, 9, 14);285            }286            let out = &mut self.buf[block * BLOCK_WORDS..(block + 1) * BLOCK_WORDS];287            for ((slot, word), start) in out.iter_mut().zip(state).zip(initial) {288                *slot = word.wrapping_add(start);289            }290        }291        self.block_pos = self.block_pos.wrapping_add(4);292    }293}294295fn quarter(s: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {296    s[a] = s[a].wrapping_add(s[b]);297    s[d] = (s[d] ^ s[a]).rotate_left(16);298    s[c] = s[c].wrapping_add(s[d]);299    s[b] = (s[b] ^ s[c]).rotate_left(12);300    s[a] = s[a].wrapping_add(s[b]);301    s[d] = (s[d] ^ s[a]).rotate_left(8);302    s[c] = s[c].wrapping_add(s[d]);303    s[b] = (s[b] ^ s[c]).rotate_left(7);304}305306fn pcg32(state: &mut u64) -> [u8; 4] {307    const MUL: u64 = 6364136223846793005;308    const INC: u64 = 11634580027462260723;309    *state = state.wrapping_mul(MUL).wrapping_add(INC);310    let s = *state;311    let xorshifted = (((s >> 18) ^ s) >> 27) as u32;312    let rot = (s >> 59) as u32;313    xorshifted.rotate_right(rot).to_le_bytes()314}315316fn wide_mul_u64(a: u64, b: u64) -> (u64, u64) {317    let t = u128::from(a) * u128::from(b);318    ((t >> 64) as u64, t as u64)319}320321fn wide_mul_u32(a: u32, b: u32) -> (u32, u32) {322    let t = u64::from(a) * u64::from(b);323    ((t >> 32) as u32, t as u32)324}325326#[cfg(test)]327mod tests {328    use super::*;329330    #[test]331    fn words_match_rand_chacha() {332        let mut c = ChaCha8::from_u64(7);333        let words: Vec<u32> = (0..8).map(|_| c.next_u32()).collect();334        assert_eq!(335            words,336            [337                601310139, 677729076, 781920570, 721508819, 1160025393, 3024842937, 155073251,338                3121330102339            ]340        );341        let doubles: Vec<u64> = (0..4).map(|_| c.next_u64()).collect();342        assert_eq!(343            doubles,344            [345                11091271176959810440,346                6629102542470643238,347                1532177833826564251,348                15666632647712377470349            ]350        );351        let mut z = ChaCha8::from_u64(0);352        let zwords: Vec<u32> = (0..8).map(|_| z.next_u32()).collect();353        assert_eq!(354            zwords,355            [356                2811902828, 3045455719, 3134767159, 2001118559, 2179114726, 3002797362, 2409334908,357                258433188358            ]359        );360    }361362    #[test]363    fn distributions_match_rand() {364        let mut c = ChaCha8::from_u64(7);365        let units: Vec<u64> = (0..4).map(|_| c.unit().to_bits()).collect();366        assert_eq!(367            units,368            [369                4594853223840476064,370                4595220474943196564,371                4604518774960858258,372                4604721123211421639373            ]374        );375        let bools: Vec<bool> = (0..8).map(|_| c.boolean()).collect();376        assert_eq!(bools, [false, true, true, false, false, false, false, true]);377        let ranges: Vec<i64> = (0..8).map(|_| c.range_i64(-1000, 1000)).collect();378        assert_eq!(ranges, [-271, 980, -600, -232, 43, -486, -147, 53]);379        let belows: Vec<u64> = (0..8).map(|_| c.below_u64(97)).collect();380        assert_eq!(belows, [29, 7, 31, 21, 42, 79, 2, 7]);381    }382383    #[test]384    fn sequences_match_rand() {385        let mut s = ChaCha8::from_u64(7);386        let mut perm: Vec<usize> = (0..16).collect();387        s.shuffle(&mut perm);388        assert_eq!(perm, [11, 4, 9, 8, 10, 1, 12, 6, 5, 7, 0, 3, 13, 14, 15, 2]);389        let floyd = ChaCha8::from_u64(7).sample_indices(100, 5);390        assert_eq!(floyd, [13, 15, 16, 26, 70]);391        let inplace = ChaCha8::from_u64(7).sample_indices(100, 20);392        assert_eq!(393            inplace,394            [14, 16, 18, 29, 71, 8, 74, 13, 63, 41, 17, 86, 43, 99, 53, 32, 48, 62, 91, 0]395        );396        let rejection = ChaCha8::from_u64(7).sample_indices(100_000, 163);397        assert_eq!(398            &rejection[..10],399            [14000, 15779, 18205, 16798, 27008, 70427, 3610, 72674, 7116, 60125]400        );401    }402}