|
| 1 | +// This file is part of the uutils coreutils package. |
| 2 | +// |
| 3 | +// For the full copyright and license information, please view the LICENSE |
| 4 | +// file that was distributed with this source code. |
| 5 | + |
| 6 | +use std::{io::BufRead, ops::RangeInclusive}; |
| 7 | + |
| 8 | +use uucore::error::{FromIo, UResult, USimpleError}; |
| 9 | +use uucore::translate; |
| 10 | + |
| 11 | +/// A uniform integer generator that tries to exactly match GNU shuf's --random-source. |
| 12 | +/// |
| 13 | +/// It's not particularly efficient and possibly not quite uniform. It should *only* be |
| 14 | +/// used for compatibility with GNU: other modes shouldn't touch this code. |
| 15 | +/// |
| 16 | +/// All the logic here was black box reverse engineered. It might not match up in all edge |
| 17 | +/// cases but it gives identical results on many different large and small inputs. |
| 18 | +/// |
| 19 | +/// It seems that GNU uses fairly textbook rejection sampling to generate integers, reading |
| 20 | +/// one byte at a time until it has enough entropy, and recycling leftover entropy after |
| 21 | +/// accepting or rejecting a value. |
| 22 | +/// |
| 23 | +/// To do your own experiments, start with commands like these: |
| 24 | +/// |
| 25 | +/// printf '\x01\x02\x03\x04' | shuf -i0-255 -r --random-source=/dev/stdin |
| 26 | +/// |
| 27 | +/// Then vary the integer range and the input and the input length. It can be useful to |
| 28 | +/// see when exactly shuf crashes with an "end of file" error. |
| 29 | +/// |
| 30 | +/// To spot small inconsistencies it's useful to run: |
| 31 | +/// |
| 32 | +/// diff -y <(my_shuf ...) <(shuf -i0-{MAX} -r --random-source={INPUT}) | head -n 50 |
| 33 | +pub struct RandomSourceAdapter<R> { |
| 34 | + reader: R, |
| 35 | + state: u64, |
| 36 | + entropy: u64, |
| 37 | +} |
| 38 | + |
| 39 | +impl<R> RandomSourceAdapter<R> { |
| 40 | + pub fn new(reader: R) -> Self { |
| 41 | + Self { |
| 42 | + reader, |
| 43 | + state: 0, |
| 44 | + entropy: 0, |
| 45 | + } |
| 46 | + } |
| 47 | +} |
| 48 | + |
| 49 | +impl<R: BufRead> RandomSourceAdapter<R> { |
| 50 | + fn generate_at_most(&mut self, at_most: u64) -> UResult<u64> { |
| 51 | + while self.entropy < at_most { |
| 52 | + let buf = self |
| 53 | + .reader |
| 54 | + .fill_buf() |
| 55 | + .map_err_context(|| translate!("shuf-error-read-random-bytes"))?; |
| 56 | + let Some(&byte) = buf.first() else { |
| 57 | + return Err(USimpleError::new( |
| 58 | + 1, |
| 59 | + translate!("shuf-error-end-of-random-bytes"), |
| 60 | + )); |
| 61 | + }; |
| 62 | + self.reader.consume(1); |
| 63 | + // Is overflow OK here? Won't it cause bias? (Seems to work out...) |
| 64 | + self.state = self.state.wrapping_mul(256).wrapping_add(byte as u64); |
| 65 | + self.entropy = self.entropy.wrapping_mul(256).wrapping_add(255); |
| 66 | + } |
| 67 | + |
| 68 | + if at_most == u64::MAX { |
| 69 | + // at_most + 1 would overflow but this case is easy. |
| 70 | + let val = self.state; |
| 71 | + self.entropy = 0; |
| 72 | + self.state = 0; |
| 73 | + return Ok(val); |
| 74 | + } |
| 75 | + |
| 76 | + let num_possibilities = at_most + 1; |
| 77 | + |
| 78 | + // If the generated number falls within this margin at the upper end of the |
| 79 | + // range then we retry to avoid modulo bias. |
| 80 | + let margin = ((self.entropy as u128 + 1) % num_possibilities as u128) as u64; |
| 81 | + let safe_zone = self.entropy - margin; |
| 82 | + |
| 83 | + if self.state <= safe_zone { |
| 84 | + let val = self.state % num_possibilities; |
| 85 | + // Reuse the rest of the state. |
| 86 | + self.state /= num_possibilities; |
| 87 | + // We need this subtraction, otherwise we consume new input slightly more |
| 88 | + // slowly than GNU. Not sure if it checks out mathematically. |
| 89 | + self.entropy -= at_most; |
| 90 | + self.entropy /= num_possibilities; |
| 91 | + Ok(val) |
| 92 | + } else { |
| 93 | + self.state %= num_possibilities; |
| 94 | + self.entropy %= num_possibilities; |
| 95 | + // I sure hope the compiler optimizes this tail call. |
| 96 | + self.generate_at_most(at_most) |
| 97 | + } |
| 98 | + } |
| 99 | + |
| 100 | + pub fn choose_from_range(&mut self, range: RangeInclusive<u64>) -> UResult<u64> { |
| 101 | + let offset = self.generate_at_most(*range.end() - *range.start())?; |
| 102 | + Ok(*range.start() + offset) |
| 103 | + } |
| 104 | + |
| 105 | + pub fn choose_from_slice<T: Copy>(&mut self, vals: &[T]) -> UResult<T> { |
| 106 | + assert!(!vals.is_empty()); |
| 107 | + let idx = self.generate_at_most(vals.len() as u64 - 1)? as usize; |
| 108 | + Ok(vals[idx]) |
| 109 | + } |
| 110 | + |
| 111 | + pub fn shuffle<'a, T>(&mut self, vals: &'a mut [T], amount: usize) -> UResult<&'a mut [T]> { |
| 112 | + // Fisher-Yates shuffle. |
| 113 | + // TODO: GNU does something different if amount <= vals.len() and the input is stdin. |
| 114 | + // The order changes completely and depends on --head-count. |
| 115 | + // No clue what they might do differently and why. |
| 116 | + let amount = amount.min(vals.len()); |
| 117 | + for idx in 0..amount { |
| 118 | + let other_idx = self.generate_at_most((vals.len() - idx - 1) as u64)? as usize + idx; |
| 119 | + vals.swap(idx, other_idx); |
| 120 | + } |
| 121 | + Ok(&mut vals[..amount]) |
| 122 | + } |
| 123 | +} |
0 commit comments