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 the absolute word position in the keystream.69    pub fn word_pos(&self) -> u128 {70        let buf_start_block = self.block_pos.wrapping_sub(4);71        let blocks_part = (self.index / BLOCK_WORDS) as u64;72        let words_part = (self.index % BLOCK_WORDS) as u64;73        let pos_block = buf_start_block.wrapping_add(blocks_part);74        u128::from(pos_block) * BLOCK_WORDS as u128 + u128::from(words_part)75    }7677    /// Seeks the keystream to an absolute word position.78    pub fn set_word_pos(&mut self, word_offset: u128) {79        self.block_pos = (word_offset / BLOCK_WORDS as u128) as u64;80        self.generate_and_set((word_offset % BLOCK_WORDS as u128) as usize);81    }8283    /// Returns a uniform double in the half-open unit interval.84    pub fn unit(&mut self) -> f64 {85        let scale = 1.0 / ((1u64 << 53) as f64);86        scale * ((self.next_u64() >> 11) as f64)87    }8889    /// Returns a fair coin flip.90    pub fn boolean(&mut self) -> bool {91        (self.next_u32() as i32) < 092    }9394    /// Returns an unbiased uniform integer between low and high inclusive.95    pub fn range_i64(&mut self, low: i64, high: i64) -> i64 {96        assert!(low <= high, "range_i64: low > high");97        let range = high.wrapping_sub(low).wrapping_add(1) as u64;98        if range == 0 {99            return self.next_u64() as i64;100        }101        let zone = (range << range.leading_zeros()).wrapping_sub(1);102        loop {103            let (hi, lo) = wide_mul_u64(self.next_u64(), range);104            if lo <= zone {105                return low.wrapping_add(hi as i64);106            }107        }108    }109110    /// Returns an unbiased uniform integer below n.111    pub fn below_u64(&mut self, n: u64) -> u64 {112        assert!(n > 0, "below_u64: empty range");113        let zone = (n << n.leading_zeros()).wrapping_sub(1);114        loop {115            let (hi, lo) = wide_mul_u64(self.next_u64(), n);116            if lo <= zone {117                return hi;118            }119        }120    }121122    /// Shuffles the slice in place by Fisher-Yates.123    pub fn shuffle<T>(&mut self, seq: &mut [T]) {124        for i in (1..seq.len()).rev() {125            let j = self.index_below(i + 1);126            seq.swap(i, j);127        }128    }129130    /// Draws amount distinct indices below length, choosing its strategy by the sizes.131    pub fn sample_indices(&mut self, length: usize, amount: usize) -> Vec<usize> {132        assert!(amount <= length, "sample_indices: amount > length");133        if length > u32::MAX as usize {134            return self.sample_rejection_u64(length as u64, amount as u64);135        }136        let length = length as u32;137        let amount = amount as u32;138        if amount < 163 {139            const C: [[f32; 2]; 2] = [[1.6, 8.0 / 45.0], [10.0, 70.0 / 9.0]];140            let j = if length < 500_000 { 0 } else { 1 };141            let amount_fp = amount as f32;142            let m4 = C[0][j] * amount_fp;143            if amount > 11 && (length as f32) < (C[1][j] + m4) * amount_fp {144                self.sample_inplace(length, amount)145            } else {146                self.sample_floyd(length, amount)147            }148        } else {149            const C: [f32; 2] = [270.0, 330.0 / 9.0];150            let j = if length < 500_000 { 0 } else { 1 };151            if (length as f32) < C[j] * (amount as f32) {152                self.sample_inplace(length, amount)153            } else {154                self.sample_rejection_u32(length, amount)155            }156        }157    }158159    fn index_below(&mut self, ubound: usize) -> usize {160        if ubound <= u32::MAX as usize {161            self.range_u32(0, ubound as u32 - 1) as usize162        } else {163            self.below_u64(ubound as u64) as usize164        }165    }166167    fn range_u32(&mut self, low: u32, high: u32) -> u32 {168        let range = high.wrapping_sub(low).wrapping_add(1);169        if range == 0 {170            return self.next_u32();171        }172        let zone = (range << range.leading_zeros()).wrapping_sub(1);173        loop {174            let (hi, lo) = wide_mul_u32(self.next_u32(), range);175            if lo <= zone {176                return low.wrapping_add(hi);177            }178        }179    }180181    fn sample_floyd(&mut self, length: u32, amount: u32) -> Vec<usize> {182        let floyd_shuffle = amount < 50;183        let mut indices: Vec<u32> = Vec::with_capacity(amount as usize);184        for j in length - amount..length {185            let t = self.range_u32(0, j);186            if floyd_shuffle {187                if let Some(pos) = indices.iter().position(|&x| x == t) {188                    indices.insert(pos, j);189                    continue;190                }191            } else if indices.contains(&t) {192                indices.push(j);193                continue;194            }195            indices.push(t);196        }197        if !floyd_shuffle {198            for i in (1..amount).rev() {199                let j = self.range_u32(0, i);200                indices.swap(i as usize, j as usize);201            }202        }203        indices.into_iter().map(|i| i as usize).collect()204    }205206    fn sample_inplace(&mut self, length: u32, amount: u32) -> Vec<usize> {207        let mut indices: Vec<u32> = (0..length).collect();208        for i in 0..amount {209            let j = self.range_u32(i, length - 1);210            indices.swap(i as usize, j as usize);211        }212        indices.truncate(amount as usize);213        indices.into_iter().map(|i| i as usize).collect()214    }215216    fn sample_rejection_u32(&mut self, length: u32, amount: u32) -> Vec<usize> {217        let ints_to_reject = (u32::MAX - length + 1) % length;218        let zone = u32::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_u32(length, zone);223            while !cache.insert(pos) {224                pos = self.distr_u32(length, zone);225            }226            indices.push(pos as usize);227        }228        indices229    }230231    fn sample_rejection_u64(&mut self, length: u64, amount: u64) -> Vec<usize> {232        let ints_to_reject = (u64::MAX - length + 1) % length;233        let zone = u64::MAX - ints_to_reject;234        let mut cache = HashSet::with_capacity(amount as usize);235        let mut indices = Vec::with_capacity(amount as usize);236        for _ in 0..amount {237            let mut pos = self.distr_u64(length, zone);238            while !cache.insert(pos) {239                pos = self.distr_u64(length, zone);240            }241            indices.push(pos as usize);242        }243        indices244    }245246    fn distr_u32(&mut self, range: u32, zone: u32) -> u32 {247        loop {248            let (hi, lo) = wide_mul_u32(self.next_u32(), range);249            if lo <= zone {250                return hi;251            }252        }253    }254255    fn distr_u64(&mut self, range: u64, zone: u64) -> u64 {256        loop {257            let (hi, lo) = wide_mul_u64(self.next_u64(), range);258            if lo <= zone {259                return hi;260            }261        }262    }263264    fn generate_and_set(&mut self, index: usize) {265        self.generate();266        self.index = index;267    }268269    fn generate(&mut self) {270        for block in 0..4 {271            let counter = self.block_pos.wrapping_add(block as u64);272            let initial = [273                CONSTANTS[0],274                CONSTANTS[1],275                CONSTANTS[2],276                CONSTANTS[3],277                self.key[0],278                self.key[1],279                self.key[2],280                self.key[3],281                self.key[4],282                self.key[5],283                self.key[6],284                self.key[7],285                counter as u32,286                (counter >> 32) as u32,287                0,288                0,289            ];290            let mut state = initial;291            for _ in 0..4 {292                quarter(&mut state, 0, 4, 8, 12);293                quarter(&mut state, 1, 5, 9, 13);294                quarter(&mut state, 2, 6, 10, 14);295                quarter(&mut state, 3, 7, 11, 15);296                quarter(&mut state, 0, 5, 10, 15);297                quarter(&mut state, 1, 6, 11, 12);298                quarter(&mut state, 2, 7, 8, 13);299                quarter(&mut state, 3, 4, 9, 14);300            }301            let out = &mut self.buf[block * BLOCK_WORDS..(block + 1) * BLOCK_WORDS];302            for ((slot, word), start) in out.iter_mut().zip(state).zip(initial) {303                *slot = word.wrapping_add(start);304            }305        }306        self.block_pos = self.block_pos.wrapping_add(4);307    }308}309310fn quarter(s: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {311    s[a] = s[a].wrapping_add(s[b]);312    s[d] = (s[d] ^ s[a]).rotate_left(16);313    s[c] = s[c].wrapping_add(s[d]);314    s[b] = (s[b] ^ s[c]).rotate_left(12);315    s[a] = s[a].wrapping_add(s[b]);316    s[d] = (s[d] ^ s[a]).rotate_left(8);317    s[c] = s[c].wrapping_add(s[d]);318    s[b] = (s[b] ^ s[c]).rotate_left(7);319}320321fn pcg32(state: &mut u64) -> [u8; 4] {322    const MUL: u64 = 6364136223846793005;323    const INC: u64 = 11634580027462260723;324    *state = state.wrapping_mul(MUL).wrapping_add(INC);325    let s = *state;326    let xorshifted = (((s >> 18) ^ s) >> 27) as u32;327    let rot = (s >> 59) as u32;328    xorshifted.rotate_right(rot).to_le_bytes()329}330331fn wide_mul_u64(a: u64, b: u64) -> (u64, u64) {332    let t = u128::from(a) * u128::from(b);333    ((t >> 64) as u64, t as u64)334}335336fn wide_mul_u32(a: u32, b: u32) -> (u32, u32) {337    let t = u64::from(a) * u64::from(b);338    ((t >> 32) as u32, t as u32)339}340341#[cfg(test)]342mod tests {343    use super::*;344345    #[test]346    fn words_match_rand_chacha() {347        let mut c = ChaCha8::from_u64(7);348        let words: Vec<u32> = (0..8).map(|_| c.next_u32()).collect();349        assert_eq!(350            words,351            [352                601310139, 677729076, 781920570, 721508819, 1160025393, 3024842937, 155073251,353                3121330102354            ]355        );356        let doubles: Vec<u64> = (0..4).map(|_| c.next_u64()).collect();357        assert_eq!(358            doubles,359            [360                11091271176959810440,361                6629102542470643238,362                1532177833826564251,363                15666632647712377470364            ]365        );366        let mut z = ChaCha8::from_u64(0);367        let zwords: Vec<u32> = (0..8).map(|_| z.next_u32()).collect();368        assert_eq!(369            zwords,370            [371                2811902828, 3045455719, 3134767159, 2001118559, 2179114726, 3002797362, 2409334908,372                258433188373            ]374        );375    }376377    #[test]378    fn distributions_match_rand() {379        let mut c = ChaCha8::from_u64(7);380        let units: Vec<u64> = (0..4).map(|_| c.unit().to_bits()).collect();381        assert_eq!(382            units,383            [384                4594853223840476064,385                4595220474943196564,386                4604518774960858258,387                4604721123211421639388            ]389        );390        let bools: Vec<bool> = (0..8).map(|_| c.boolean()).collect();391        assert_eq!(bools, [false, true, true, false, false, false, false, true]);392        let ranges: Vec<i64> = (0..8).map(|_| c.range_i64(-1000, 1000)).collect();393        assert_eq!(ranges, [-271, 980, -600, -232, 43, -486, -147, 53]);394        let belows: Vec<u64> = (0..8).map(|_| c.below_u64(97)).collect();395        assert_eq!(belows, [29, 7, 31, 21, 42, 79, 2, 7]);396    }397398    #[test]399    fn sequences_match_rand() {400        let mut s = ChaCha8::from_u64(7);401        let mut perm: Vec<usize> = (0..16).collect();402        s.shuffle(&mut perm);403        assert_eq!(perm, [11, 4, 9, 8, 10, 1, 12, 6, 5, 7, 0, 3, 13, 14, 15, 2]);404        let floyd = ChaCha8::from_u64(7).sample_indices(100, 5);405        assert_eq!(floyd, [13, 15, 16, 26, 70]);406        let inplace = ChaCha8::from_u64(7).sample_indices(100, 20);407        assert_eq!(408            inplace,409            [14, 16, 18, 29, 71, 8, 74, 13, 63, 41, 17, 86, 43, 99, 53, 32, 48, 62, 91, 0]410        );411        let rejection = ChaCha8::from_u64(7).sample_indices(100_000, 163);412        assert_eq!(413            &rejection[..10],414            [14000, 15779, 18205, 16798, 27008, 70427, 3610, 72674, 7116, 60125]415        );416    }417418    #[test]419    fn word_pos_seeks_the_stream() {420        let mut c = ChaCha8::from_u64(7);421        assert_eq!(c.word_pos(), 0);422        for _ in 0..70 {423            c.next_u32();424        }425        assert_eq!(c.word_pos(), 70);426        c.next_u64();427        assert_eq!(c.word_pos(), 72);428        c.set_word_pos(5);429        let a = c.next_u64();430        c.set_word_pos(5);431        assert_eq!(c.next_u64(), a);432    }433}