shuf: Add --random-seed, make --random-source GNU-compatible, report write failures, optimize

.vscode/cspell.dictionaries/jargon.wordlist.txt

@@ -55,6 +55,7 @@ fileio
 filesystem
 filesystems
 flamegraph
+footgun
 freeram
 fsxattr
 fullblock
@@ -92,6 +93,7 @@ mergeable
 microbenchmark
 microbenchmarks
 microbenchmarking
+monomorphized
 multibyte
 multicall
 nmerge
@@ -106,6 +108,7 @@ nolinks
 nonblock
 nonportable
 nonprinting
+nonrepeating
 nonseekable
 notrunc
 nowrite

.vscode/cspell.dictionaries/people.wordlist.txt

@@ -37,6 +37,9 @@ Boden Garman
 Chirag B Jadwani
     Chirag
     Jadwani
+Daniel Lemire
+    Daniel
+    Lemire
 Derek Chiang
     Derek
     Chiang

.vscode/cspell.dictionaries/workspace.wordlist.txt

@@ -38,6 +38,7 @@ getrandom
 globset
 indicatif
 itertools
+itoa
 iuse
 langid
 lscolors

Cargo.toml

@@ -335,6 +335,7 @@ icu_locale = "2.0.0"
 icu_provider = "2.0.0"
 indicatif = "0.18.0"
 itertools = "0.14.0"
+itoa = "1.0.15"
 jiff = { version = "0.2.10", default-features = false, features = [
   "std",
   "alloc",
@@ -360,6 +361,7 @@ phf_codegen = "0.13.1"
 platform-info = "2.0.3"
 procfs = "0.18"
 rand = { version = "0.9.0", features = ["small_rng"] }
+rand_chacha = { version = "0.9.0" }
 rand_core = "0.9.0"
 rayon = "1.10"
 regex = "1.10.4"

src/uu/shuf/Cargo.toml

@@ -19,8 +19,11 @@ path = "src/shuf.rs"
 
 [dependencies]
 clap = { workspace = true }
+itoa = { workspace = true }
 rand = { workspace = true }
+rand_chacha = { workspace = true }
 rand_core = { workspace = true }
+sha3 = { workspace = true }
 uucore = { workspace = true }
 fluent = { workspace = true }
 

src/uu/shuf/locales/en-US.ftl

@@ -10,6 +10,7 @@ shuf-help-echo = treat each ARG as an input line
 shuf-help-input-range = treat each number LO through HI as an input line
 shuf-help-head-count = output at most COUNT lines
 shuf-help-output = write result to FILE instead of standard output
+shuf-help-random-seed = seed with STRING for reproducible output
 shuf-help-random-source = get random bytes from FILE
 shuf-help-repeat = output lines can be repeated
 shuf-help-zero-terminated = line delimiter is NUL, not newline
@@ -19,6 +20,8 @@ shuf-error-unexpected-argument = unexpected argument { $arg } found
 shuf-error-failed-to-open-for-writing = failed to open { $file } for writing
 shuf-error-failed-to-open-random-source = failed to open random source { $file }
 shuf-error-read-error = read error
+shuf-error-read-random-bytes = reading random bytes failed
+shuf-error-end-of-random-bytes = end of random source
 shuf-error-no-lines-to-repeat = no lines to repeat
 shuf-error-start-exceeds-end = start exceeds end
 shuf-error-missing-dash = missing '-'

src/uu/shuf/src/compat_random_source.rs

@@ -0,0 +1,123 @@
+// This file is part of the uutils coreutils package.
+//
+// For the full copyright and license information, please view the LICENSE
+// file that was distributed with this source code.
+
+use std::{io::BufRead, ops::RangeInclusive};
+
+use uucore::error::{FromIo, UResult, USimpleError};
+use uucore::translate;
+
+/// A uniform integer generator that tries to exactly match GNU shuf's --random-source.
+///
+/// It's not particularly efficient and possibly not quite uniform. It should *only* be
+/// used for compatibility with GNU: other modes shouldn't touch this code.
+///
+/// All the logic here was black box reverse engineered. It might not match up in all edge
+/// cases but it gives identical results on many different large and small inputs.
+///
+/// It seems that GNU uses fairly textbook rejection sampling to generate integers, reading
+/// one byte at a time until it has enough entropy, and recycling leftover entropy after
+/// accepting or rejecting a value.
+///
+/// To do your own experiments, start with commands like these:
+///
+///   printf '\x01\x02\x03\x04' | shuf -i0-255 -r --random-source=/dev/stdin
+///
+/// Then vary the integer range and the input and the input length. It can be useful to
+/// see when exactly shuf crashes with an "end of file" error.
+///
+/// To spot small inconsistencies it's useful to run:
+///
+///   diff -y <(my_shuf ...) <(shuf -i0-{MAX} -r --random-source={INPUT}) | head -n 50
+pub struct RandomSourceAdapter<R> {
+    reader: R,
+    state: u64,
+    entropy: u64,
+}
+
+impl<R> RandomSourceAdapter<R> {
+    pub fn new(reader: R) -> Self {
+        Self {
+            reader,
+            state: 0,
+            entropy: 0,
+        }
+    }
+}
+
+impl<R: BufRead> RandomSourceAdapter<R> {
+    fn generate_at_most(&mut self, at_most: u64) -> UResult<u64> {
+        while self.entropy < at_most {
+            let buf = self
+                .reader
+                .fill_buf()
+                .map_err_context(|| translate!("shuf-error-read-random-bytes"))?;
+            let Some(&byte) = buf.first() else {
+                return Err(USimpleError::new(
+                    1,
+                    translate!("shuf-error-end-of-random-bytes"),
+                ));
+            };
+            self.reader.consume(1);
+            // Is overflow OK here? Won't it cause bias? (Seems to work out...)
+            self.state = self.state.wrapping_mul(256).wrapping_add(byte as u64);
+            self.entropy = self.entropy.wrapping_mul(256).wrapping_add(255);
+        }
+
+        if at_most == u64::MAX {
+            // at_most + 1 would overflow but this case is easy.
+            let val = self.state;
+            self.entropy = 0;
+            self.state = 0;
+            return Ok(val);
+        }
+
+        let num_possibilities = at_most + 1;
+
+        // If the generated number falls within this margin at the upper end of the
+        // range then we retry to avoid modulo bias.
+        let margin = ((self.entropy as u128 + 1) % num_possibilities as u128) as u64;
+        let safe_zone = self.entropy - margin;
+
+        if self.state <= safe_zone {
+            let val = self.state % num_possibilities;
+            // Reuse the rest of the state.
+            self.state /= num_possibilities;
+            // We need this subtraction, otherwise we consume new input slightly more
+            // slowly than GNU. Not sure if it checks out mathematically.
+            self.entropy -= at_most;
+            self.entropy /= num_possibilities;
+            Ok(val)
+        } else {
+            self.state %= num_possibilities;
+            self.entropy %= num_possibilities;
+            // I sure hope the compiler optimizes this tail call.
+            self.generate_at_most(at_most)
+        }
+    }
+
+    pub fn choose_from_range(&mut self, range: RangeInclusive<u64>) -> UResult<u64> {
+        let offset = self.generate_at_most(*range.end() - *range.start())?;
+        Ok(*range.start() + offset)
+    }
+
+    pub fn choose_from_slice<T: Copy>(&mut self, vals: &[T]) -> UResult<T> {
+        assert!(!vals.is_empty());
+        let idx = self.generate_at_most(vals.len() as u64 - 1)? as usize;
+        Ok(vals[idx])
+    }
+
+    pub fn shuffle<'a, T>(&mut self, vals: &'a mut [T], amount: usize) -> UResult<&'a mut [T]> {
+        // Fisher-Yates shuffle.
+        // TODO: GNU does something different if amount <= vals.len() and the input is stdin.
+        // The order changes completely and depends on --head-count.
+        // No clue what they might do differently and why.
+        let amount = amount.min(vals.len());
+        for idx in 0..amount {
+            let other_idx = self.generate_at_most((vals.len() - idx - 1) as u64)? as usize + idx;
+            vals.swap(idx, other_idx);
+        }
+        Ok(&mut vals[..amount])
+    }
+}

src/uu/shuf/src/nonrepeating_iterator.rs

@@ -0,0 +1,111 @@
+use std::collections::HashMap;
+use std::ops::RangeInclusive;
+
+use uucore::error::UResult;
+
+use crate::WrappedRng;
+
+/// An iterator that samples from an integer range without repetition.
+///
+/// This is based on Fisher-Yates, and it's required for backward compatibility
+/// that it behaves exactly like Fisher-Yates if --random-source or --random-seed
+/// is used. But we have a few tricks:
+///
+/// - In the beginning we use a hash table instead of an array. This way we lazily
+///   keep track of swaps without allocating the entire range upfront.
+///
+/// - When the hash table starts to get big relative to the remaining items
+///   we switch over to an array.
+///
+/// - We store the array backwards so that we can shrink it as we go and free excess
+///   memory every now and then.
+///
+/// Both the hash table and the array give the same output.
+///
+/// There's room for optimization:
+///
+/// - Switching over from the hash table to the array is costly. If we happen to know
+///   (through --head-count) that only few draws remain then it would be better not
+///   to switch.
+///
+/// - If the entire range gets used then we might as well allocate an array to start
+///   with. But if the user e.g. pipes through `head` rather than using --head-count
+///   we can't know whether that's the case, so there's a tradeoff.
+///
+///   GNU decides the other way: --head-count is noticeably faster than | head.
+pub(crate) struct NonrepeatingIterator<'a> {
+    rng: &'a mut WrappedRng,
+    values: Values,
+}
+
+enum Values {
+    Full(Vec<u64>),
+    Sparse(RangeInclusive<u64>, HashMap<u64, u64>),
+}
+
+impl<'a> NonrepeatingIterator<'a> {
+    pub(crate) fn new(range: RangeInclusive<u64>, rng: &'a mut WrappedRng) -> Self {
+        let values = Values::Sparse(range, HashMap::default());
+        NonrepeatingIterator { rng, values }
+    }
+
+    fn produce(&mut self) -> UResult<u64> {
+        match &mut self.values {
+            Values::Full(items) => {
+                let this_idx = items.len() - 1;
+
+                let other_idx = self.rng.choose_from_range(0..=items.len() as u64 - 1)? as usize;
+                // Flip the index to pretend we're going left-to-right
+                let other_idx = items.len() - other_idx - 1;
+
+                items.swap(this_idx, other_idx);
+
+                let val = items.pop().unwrap();
+                if items.len().is_power_of_two() && items.len() >= 512 {
+                    items.shrink_to_fit();
+                }
+                Ok(val)
+            }
+            Values::Sparse(range, items) => {
+                let this_idx = *range.start();
+                let this_val = items.remove(&this_idx).unwrap_or(this_idx);
+
+                let other_idx = self.rng.choose_from_range(range.clone())?;
+
+                let val = if this_idx == other_idx {
+                    this_val
+                } else {
+                    items.insert(other_idx, this_val).unwrap_or(other_idx)
+                };
+                *range = *range.start() + 1..=*range.end();
+
+                Ok(val)
+            }
+        }
+    }
+}
+
+impl Iterator for NonrepeatingIterator<'_> {
+    type Item = UResult<u64>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match &self.values {
+            Values::Full(items) if items.is_empty() => return None,
+            Values::Full(_) => (),
+            Values::Sparse(range, _) if range.is_empty() => return None,
+            Values::Sparse(range, items) => {
+                let range_len = range.size_hint().0 as u64;
+                if items.len() as u64 >= range_len / 8 {
+                    self.values = Values::Full(hashmap_to_vec(range.clone(), items));
+                }
+            }
+        }
+
+        Some(self.produce())
+    }
+}
+
+fn hashmap_to_vec(range: RangeInclusive<u64>, map: &HashMap<u64, u64>) -> Vec<u64> {
+    let lookup = |idx| *map.get(&idx).unwrap_or(&idx);
+    range.rev().map(lookup).collect()
+}