mirror of
https://github.com/uutils/coreutils.git
synced 2026-06-10 15:48:22 -07:00
shuf: correctness: Make --random-source compatible with GNU shuf
When the --random-source option is used uutils shuf now gives identical output to GNU shuf in many (but not all) cases. This is helpful to users who use it to get deterministic output, e.g. by combining it with `openssl` as suggested in the GNU info pages. I reverse engineered the algorithm from GNU shuf's output. There may be bugs. Other modes of shuffling still use `rand`'s `ThreadRng`, though they now sample a uniform distribution directly without going through the slice helper trait. Additionally, switch from `usize` to `u64` for `--input-range` and `--head-count`. This way the same range of numbers can be generated on 32-bit platforms as on 64-bit platforms.
This commit is contained in:
Jan Verbeek
@@ -93,6 +93,7 @@ mergeable
|
||||
microbenchmark
|
||||
microbenchmarks
|
||||
microbenchmarking
|
||||
monomorphized
|
||||
multibyte
|
||||
multicall
|
||||
nmerge
|
||||
@@ -107,6 +108,7 @@ nolinks
|
||||
nonblock
|
||||
nonportable
|
||||
nonprinting
|
||||
nonrepeating
|
||||
nonseekable
|
||||
notrunc
|
||||
nowrite
|
||||
|
||||
@@ -20,6 +20,7 @@ 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 '-'
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
use std::io::BufRead;
|
||||
|
||||
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> {
|
||||
pub fn get_value(&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.get_value(at_most)
|
||||
}
|
||||
}
|
||||
|
||||
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.get_value((vals.len() - idx - 1) as u64)? as usize + idx;
|
||||
vals.swap(idx, other_idx);
|
||||
}
|
||||
Ok(&mut vals[..amount])
|
||||
}
|
||||
}
|
||||
@@ -1,32 +1,30 @@
|
||||
// spell-checker:ignore nonrepeating
|
||||
|
||||
// TODO: this iterator is not compatible with GNU when --random-source is used
|
||||
|
||||
use std::{collections::HashSet, ops::RangeInclusive};
|
||||
|
||||
use rand::{Rng, seq::SliceRandom};
|
||||
use uucore::error::UResult;
|
||||
|
||||
use crate::WrappedRng;
|
||||
|
||||
enum NumberSet {
|
||||
AlreadyListed(HashSet<usize>),
|
||||
Remaining(Vec<usize>),
|
||||
AlreadyListed(HashSet<u64>),
|
||||
Remaining(Vec<u64>),
|
||||
}
|
||||
|
||||
pub(crate) struct NonrepeatingIterator<'a> {
|
||||
range: RangeInclusive<usize>,
|
||||
range: RangeInclusive<u64>,
|
||||
rng: &'a mut WrappedRng,
|
||||
remaining_count: usize,
|
||||
remaining_count: u64,
|
||||
buf: NumberSet,
|
||||
}
|
||||
|
||||
impl<'a> NonrepeatingIterator<'a> {
|
||||
pub(crate) fn new(
|
||||
range: RangeInclusive<usize>,
|
||||
rng: &'a mut WrappedRng,
|
||||
amount: usize,
|
||||
) -> Self {
|
||||
pub(crate) fn new(range: RangeInclusive<u64>, rng: &'a mut WrappedRng, amount: u64) -> Self {
|
||||
let capped_amount = if range.start() > range.end() {
|
||||
0
|
||||
} else if range == (0..=usize::MAX) {
|
||||
} else if range == (0..=u64::MAX) {
|
||||
amount
|
||||
} else {
|
||||
amount.min(range.end() - range.start() + 1)
|
||||
@@ -39,12 +37,12 @@ impl<'a> NonrepeatingIterator<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
fn produce(&mut self) -> usize {
|
||||
fn produce(&mut self) -> UResult<u64> {
|
||||
debug_assert!(self.range.start() <= self.range.end());
|
||||
match &mut self.buf {
|
||||
NumberSet::AlreadyListed(already_listed) => {
|
||||
let chosen = loop {
|
||||
let guess = self.rng.random_range(self.range.clone());
|
||||
let guess = self.rng.choose_from_range(self.range.clone())?;
|
||||
let newly_inserted = already_listed.insert(guess);
|
||||
if newly_inserted {
|
||||
break guess;
|
||||
@@ -54,32 +52,32 @@ impl<'a> NonrepeatingIterator<'a> {
|
||||
// the number of attempts to find a number that hasn't been chosen yet increases.
|
||||
// Therefore, we need to switch at some point from "set of already returned values" to "list of remaining values".
|
||||
let range_size = (self.range.end() - self.range.start()).saturating_add(1);
|
||||
if number_set_should_list_remaining(already_listed.len(), range_size) {
|
||||
if number_set_should_list_remaining(already_listed.len() as u64, range_size) {
|
||||
let mut remaining = self
|
||||
.range
|
||||
.clone()
|
||||
.filter(|n| !already_listed.contains(n))
|
||||
.collect::<Vec<_>>();
|
||||
assert!(remaining.len() >= self.remaining_count);
|
||||
remaining.partial_shuffle(&mut self.rng, self.remaining_count);
|
||||
remaining.truncate(self.remaining_count);
|
||||
assert!(remaining.len() as u64 >= self.remaining_count);
|
||||
remaining.truncate(self.remaining_count as usize);
|
||||
self.rng.shuffle(&mut remaining, usize::MAX)?;
|
||||
self.buf = NumberSet::Remaining(remaining);
|
||||
}
|
||||
chosen
|
||||
Ok(chosen)
|
||||
}
|
||||
NumberSet::Remaining(remaining_numbers) => {
|
||||
debug_assert!(!remaining_numbers.is_empty());
|
||||
// We only enter produce() when there is at least one actual element remaining, so popping must always return an element.
|
||||
remaining_numbers.pop().unwrap()
|
||||
Ok(remaining_numbers.pop().unwrap())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Iterator for NonrepeatingIterator<'_> {
|
||||
type Item = usize;
|
||||
type Item = UResult<u64>;
|
||||
|
||||
fn next(&mut self) -> Option<usize> {
|
||||
fn next(&mut self) -> Option<UResult<u64>> {
|
||||
if self.range.is_empty() || self.remaining_count == 0 {
|
||||
return None;
|
||||
}
|
||||
@@ -89,7 +87,7 @@ impl Iterator for NonrepeatingIterator<'_> {
|
||||
}
|
||||
|
||||
// This could be a method, but it is much easier to test as a stand-alone function.
|
||||
fn number_set_should_list_remaining(listed_count: usize, range_size: usize) -> bool {
|
||||
fn number_set_should_list_remaining(listed_count: u64, range_size: u64) -> bool {
|
||||
// Arbitrarily determine the switchover point to be around 25%. This is because:
|
||||
// - HashSet has a large space overhead for the hash table load factor.
|
||||
// - This means that somewhere between 25-40%, the memory required for a "positive" HashSet and a "negative" Vec should be the same.
|
||||
@@ -107,17 +105,17 @@ mod test_number_set_decision {
|
||||
|
||||
#[test]
|
||||
fn test_stay_positive_large_remaining_first() {
|
||||
assert_eq!(false, number_set_should_list_remaining(0, usize::MAX));
|
||||
assert_eq!(false, number_set_should_list_remaining(0, u64::MAX));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stay_positive_large_remaining_second() {
|
||||
assert_eq!(false, number_set_should_list_remaining(1, usize::MAX));
|
||||
assert_eq!(false, number_set_should_list_remaining(1, u64::MAX));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stay_positive_large_remaining_tenth() {
|
||||
assert_eq!(false, number_set_should_list_remaining(9, usize::MAX));
|
||||
assert_eq!(false, number_set_should_list_remaining(9, u64::MAX));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -161,22 +159,19 @@ mod test_number_set_decision {
|
||||
// Ensure that we are overflow-free:
|
||||
#[test]
|
||||
fn test_no_crash_exceed_max_size1() {
|
||||
assert_eq!(false, number_set_should_list_remaining(12345, usize::MAX));
|
||||
assert_eq!(false, number_set_should_list_remaining(12345, u64::MAX));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_crash_exceed_max_size2() {
|
||||
assert_eq!(
|
||||
true,
|
||||
number_set_should_list_remaining(usize::MAX - 1, usize::MAX)
|
||||
number_set_should_list_remaining(u64::MAX - 1, u64::MAX)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_crash_exceed_max_size3() {
|
||||
assert_eq!(
|
||||
true,
|
||||
number_set_should_list_remaining(usize::MAX, usize::MAX)
|
||||
);
|
||||
assert_eq!(true, number_set_should_list_remaining(u64::MAX, u64::MAX));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
// 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.
|
||||
// Copyright 2018 Developers of the Rand project.
|
||||
// Copyright 2013 The Rust Project Developers.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
|
||||
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
|
||||
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
|
||||
// option. This file may not be copied, modified, or distributed
|
||||
// except according to those terms.
|
||||
|
||||
//! A wrapper around any Read to treat it as an RNG.
|
||||
|
||||
use std::cell::Cell;
|
||||
use std::io::{Error, Read};
|
||||
use std::rc::Rc;
|
||||
|
||||
use rand_core::{RngCore, impls};
|
||||
|
||||
/// An RNG that reads random bytes straight from any type supporting
|
||||
/// [`std::io::Read`], for example files.
|
||||
///
|
||||
/// This will work best with an infinite reader, but that is not required.
|
||||
///
|
||||
/// This can be used with `/dev/urandom` on Unix but it is recommended to use
|
||||
/// [`OsRng`] instead.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// `ReadRng` uses [`std::io::Read::read_exact`], which retries on interrupts.
|
||||
/// All other errors from the underlying reader, including when it does not
|
||||
/// have enough data, will be reported via the public error field (which can
|
||||
/// be cloned in advance, as it uses [`Rc`]). This field must be checked for
|
||||
/// errors after every operation.
|
||||
///
|
||||
/// [`OsRng`]: rand::rngs::OsRng
|
||||
pub struct ReadRng<R> {
|
||||
reader: R,
|
||||
pub error: ErrorCell,
|
||||
}
|
||||
|
||||
pub type ErrorCell = Rc<Cell<Option<std::io::Error>>>;
|
||||
|
||||
impl<R: Read> ReadRng<R> {
|
||||
/// Create a new `ReadRng` from a `Read`.
|
||||
pub fn new(r: R) -> Self {
|
||||
Self {
|
||||
reader: r,
|
||||
error: Rc::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
|
||||
if dest.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
// Use `std::io::read_exact`, which retries on `ErrorKind::Interrupted`.
|
||||
self.reader.read_exact(dest)
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read> RngCore for ReadRng<R> {
|
||||
fn next_u32(&mut self) -> u32 {
|
||||
impls::next_u32_via_fill(self)
|
||||
}
|
||||
|
||||
fn next_u64(&mut self) -> u64 {
|
||||
impls::next_u64_via_fill(self)
|
||||
}
|
||||
|
||||
fn fill_bytes(&mut self, dest: &mut [u8]) {
|
||||
if let Err(err) = self.try_fill_bytes(dest) {
|
||||
// Failed to deliver random data, so the caller must check the error
|
||||
// cell before using the result.
|
||||
self.error.set(Some(err));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use std::println;
|
||||
|
||||
use super::ReadRng;
|
||||
use rand::RngCore;
|
||||
|
||||
#[test]
|
||||
fn test_reader_rng_u64() {
|
||||
// transmute from the target to avoid endianness concerns.
|
||||
#[rustfmt::skip]
|
||||
let v = [0u8, 0, 0, 0, 0, 0, 0, 1,
|
||||
0, 4, 0, 0, 3, 0, 0, 2,
|
||||
5, 0, 0, 0, 0, 0, 0, 0];
|
||||
let mut rng = ReadRng::new(&v[..]);
|
||||
|
||||
assert_eq!(rng.next_u64(), 1 << 56);
|
||||
assert_eq!(rng.next_u64(), (2 << 56) + (3 << 32) + (4 << 8));
|
||||
assert_eq!(rng.next_u64(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reader_rng_u32() {
|
||||
let v = [0u8, 0, 0, 1, 0, 0, 2, 0, 3, 0, 0, 0];
|
||||
let mut rng = ReadRng::new(&v[..]);
|
||||
|
||||
assert_eq!(rng.next_u32(), 1 << 24);
|
||||
assert_eq!(rng.next_u32(), 2 << 16);
|
||||
assert_eq!(rng.next_u32(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reader_rng_fill_bytes() {
|
||||
let v = [1u8, 2, 3, 4, 5, 6, 7, 8];
|
||||
let mut w = [0u8; 8];
|
||||
|
||||
let mut rng = ReadRng::new(&v[..]);
|
||||
rng.fill_bytes(&mut w);
|
||||
|
||||
assert_eq!(v, w);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reader_rng_insufficient_bytes() {
|
||||
let v = [1u8, 2, 3, 4, 5, 6, 7, 8];
|
||||
let mut w = [0u8; 9];
|
||||
|
||||
let mut rng = ReadRng::new(&v[..]);
|
||||
|
||||
let result = rng.try_fill_bytes(&mut w);
|
||||
assert!(result.is_err());
|
||||
println!("Error: {}", result.unwrap_err());
|
||||
}
|
||||
}
|
||||
+65
-73
@@ -7,12 +7,11 @@
|
||||
|
||||
use clap::builder::ValueParser;
|
||||
use clap::{Arg, ArgAction, Command};
|
||||
use rand::prelude::SliceRandom;
|
||||
use rand::seq::IndexedRandom;
|
||||
use rand::{Rng, RngCore};
|
||||
use rand::Rng;
|
||||
use rand::seq::{IndexedRandom, SliceRandom};
|
||||
use std::ffi::{OsStr, OsString};
|
||||
use std::fs::File;
|
||||
use std::io::{BufWriter, Error, Read, Write, stdin, stdout};
|
||||
use std::io::{BufReader, BufWriter, Error, Read, Write, stdin, stdout};
|
||||
use std::ops::RangeInclusive;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::str::FromStr;
|
||||
@@ -21,21 +20,21 @@ use uucore::error::{FromIo, UResult, USimpleError, UUsageError};
|
||||
use uucore::format_usage;
|
||||
use uucore::translate;
|
||||
|
||||
mod compat_random_source;
|
||||
mod nonrepeating_iterator;
|
||||
mod rand_read_adapter;
|
||||
|
||||
use nonrepeating_iterator::NonrepeatingIterator;
|
||||
|
||||
enum Mode {
|
||||
Default(PathBuf),
|
||||
Echo(Vec<OsString>),
|
||||
InputRange(RangeInclusive<usize>),
|
||||
InputRange(RangeInclusive<u64>),
|
||||
}
|
||||
|
||||
const BUF_SIZE: usize = 64 * 1024;
|
||||
|
||||
struct Options {
|
||||
head_count: usize,
|
||||
head_count: u64,
|
||||
output: Option<PathBuf>,
|
||||
random_source: Option<PathBuf>,
|
||||
repeat: bool,
|
||||
@@ -87,11 +86,11 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
|
||||
// Busybox takes the final value which is more typical: later
|
||||
// options override earlier options.
|
||||
head_count: matches
|
||||
.get_many::<usize>(options::HEAD_COUNT)
|
||||
.get_many::<u64>(options::HEAD_COUNT)
|
||||
.unwrap_or_default()
|
||||
.copied()
|
||||
.min()
|
||||
.unwrap_or(usize::MAX),
|
||||
.unwrap_or(u64::MAX),
|
||||
output: matches.get_one(options::OUTPUT).cloned(),
|
||||
random_source: matches.get_one(options::RANDOM_SOURCE).cloned(),
|
||||
repeat: matches.get_flag(options::REPEAT),
|
||||
@@ -125,7 +124,8 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
|
||||
let file = File::open(r).map_err_context(
|
||||
|| translate!("shuf-error-failed-to-open-random-source", "file" => r.quote()),
|
||||
)?;
|
||||
WrappedRng::RngFile(rand_read_adapter::ReadRng::new(file))
|
||||
let file = BufReader::new(file);
|
||||
WrappedRng::RngFile(compat_random_source::RandomSourceAdapter::new(file))
|
||||
}
|
||||
None => WrappedRng::RngDefault(rand::rng()),
|
||||
};
|
||||
@@ -180,7 +180,7 @@ pub fn uu_app() -> Command {
|
||||
.value_name("COUNT")
|
||||
.action(ArgAction::Append)
|
||||
.help(translate!("shuf-help-head-count"))
|
||||
.value_parser(usize::from_str),
|
||||
.value_parser(u64::from_str),
|
||||
)
|
||||
.arg(
|
||||
Arg::new(options::OUTPUT)
|
||||
@@ -250,12 +250,15 @@ fn split_seps(data: &[u8], sep: u8) -> Vec<&[u8]> {
|
||||
trait Shufable {
|
||||
type Item: Writable;
|
||||
fn is_empty(&self) -> bool;
|
||||
fn choose(&self, rng: &mut WrappedRng) -> Self::Item;
|
||||
fn choose(&self, rng: &mut WrappedRng) -> UResult<Self::Item>;
|
||||
// In some modes we shuffle ahead of time and in some as we generate
|
||||
// so we unfortunately need to double-wrap UResult.
|
||||
// But it's monomorphized so the optimizer will hopefully Take Care Of It™.
|
||||
fn partial_shuffle<'b>(
|
||||
&'b mut self,
|
||||
rng: &'b mut WrappedRng,
|
||||
amount: usize,
|
||||
) -> impl Iterator<Item = Self::Item>;
|
||||
amount: u64,
|
||||
) -> UResult<impl Iterator<Item = UResult<Self::Item>>>;
|
||||
}
|
||||
|
||||
impl<'a> Shufable for Vec<&'a [u8]> {
|
||||
@@ -265,20 +268,22 @@ impl<'a> Shufable for Vec<&'a [u8]> {
|
||||
(**self).is_empty()
|
||||
}
|
||||
|
||||
fn choose(&self, rng: &mut WrappedRng) -> Self::Item {
|
||||
// Note: "copied()" only copies the reference, not the entire [u8].
|
||||
// Returns None if the slice is empty. We checked this before, so
|
||||
// this is safe.
|
||||
(**self).choose(rng).unwrap()
|
||||
fn choose(&self, rng: &mut WrappedRng) -> UResult<Self::Item> {
|
||||
rng.choose(self)
|
||||
}
|
||||
|
||||
fn partial_shuffle<'b>(
|
||||
&'b mut self,
|
||||
rng: &'b mut WrappedRng,
|
||||
amount: usize,
|
||||
) -> impl Iterator<Item = Self::Item> {
|
||||
// Note: "copied()" only copies the reference, not the entire [u8].
|
||||
(**self).partial_shuffle(rng, amount).0.iter().copied()
|
||||
amount: u64,
|
||||
) -> UResult<impl Iterator<Item = UResult<Self::Item>>> {
|
||||
// On 32-bit platforms it's possible that amount > usize::MAX.
|
||||
// We saturate as usize::MAX since all of our shuffling modes require storing
|
||||
// elements in memory so more than usize::MAX elements won't fit anyway.
|
||||
// (With --repeat an output larger than usize::MAX is possible. But --repeat
|
||||
// uses `choose()`.)
|
||||
let amount = usize::try_from(amount).unwrap_or(usize::MAX);
|
||||
Ok(rng.shuffle(self, amount)?.iter().copied().map(Ok))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -289,36 +294,37 @@ impl<'a> Shufable for Vec<&'a OsStr> {
|
||||
(**self).is_empty()
|
||||
}
|
||||
|
||||
fn choose(&self, rng: &mut WrappedRng) -> Self::Item {
|
||||
(**self).choose(rng).unwrap()
|
||||
fn choose(&self, rng: &mut WrappedRng) -> UResult<Self::Item> {
|
||||
rng.choose(self)
|
||||
}
|
||||
|
||||
fn partial_shuffle<'b>(
|
||||
&'b mut self,
|
||||
rng: &'b mut WrappedRng,
|
||||
amount: usize,
|
||||
) -> impl Iterator<Item = Self::Item> {
|
||||
(**self).partial_shuffle(rng, amount).0.iter().copied()
|
||||
amount: u64,
|
||||
) -> UResult<impl Iterator<Item = UResult<Self::Item>>> {
|
||||
let amount = usize::try_from(amount).unwrap_or(usize::MAX);
|
||||
Ok(rng.shuffle(self, amount)?.iter().copied().map(Ok))
|
||||
}
|
||||
}
|
||||
|
||||
impl Shufable for RangeInclusive<usize> {
|
||||
type Item = usize;
|
||||
impl Shufable for RangeInclusive<u64> {
|
||||
type Item = u64;
|
||||
|
||||
fn is_empty(&self) -> bool {
|
||||
self.is_empty()
|
||||
}
|
||||
|
||||
fn choose(&self, rng: &mut WrappedRng) -> usize {
|
||||
rng.random_range(self.clone())
|
||||
fn choose(&self, rng: &mut WrappedRng) -> UResult<Self::Item> {
|
||||
rng.choose_from_range(self.clone())
|
||||
}
|
||||
|
||||
fn partial_shuffle<'b>(
|
||||
&'b mut self,
|
||||
rng: &'b mut WrappedRng,
|
||||
amount: usize,
|
||||
) -> impl Iterator<Item = Self::Item> {
|
||||
NonrepeatingIterator::new(self.clone(), rng, amount)
|
||||
amount: u64,
|
||||
) -> UResult<impl Iterator<Item = UResult<Self::Item>>> {
|
||||
Ok(NonrepeatingIterator::new(self.clone(), rng, amount))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -338,7 +344,7 @@ impl Writable for &OsStr {
|
||||
}
|
||||
}
|
||||
|
||||
impl Writable for usize {
|
||||
impl Writable for u64 {
|
||||
fn write_all_to(&self, output: &mut impl OsWrite) -> Result<(), Error> {
|
||||
// The itoa crate is surprisingly much more efficient than a formatted write.
|
||||
// It speeds up `shuf -r -n1000000 -i1-1024` by 1.8×.
|
||||
@@ -354,7 +360,6 @@ fn shuf_exec(
|
||||
output: &mut BufWriter<Box<dyn OsWrite>>,
|
||||
) -> UResult<()> {
|
||||
let ctx = || translate!("shuf-error-write-failed");
|
||||
let error_cell = rng.get_error_cell();
|
||||
if opts.repeat {
|
||||
if input.is_empty() {
|
||||
return Err(USimpleError::new(
|
||||
@@ -363,17 +368,16 @@ fn shuf_exec(
|
||||
));
|
||||
}
|
||||
for _ in 0..opts.head_count {
|
||||
let r = input.choose(rng);
|
||||
WrappedRng::check_error(error_cell.as_ref())?;
|
||||
let r = input.choose(rng)?;
|
||||
|
||||
r.write_all_to(output).map_err_context(ctx)?;
|
||||
output.write_all(&[opts.sep]).map_err_context(ctx)?;
|
||||
}
|
||||
} else {
|
||||
let shuffled = input.partial_shuffle(rng, opts.head_count);
|
||||
WrappedRng::check_error(error_cell.as_ref())?;
|
||||
let shuffled = input.partial_shuffle(rng, opts.head_count)?;
|
||||
|
||||
for r in shuffled {
|
||||
let r = r?;
|
||||
r.write_all_to(output).map_err_context(ctx)?;
|
||||
output.write_all(&[opts.sep]).map_err_context(ctx)?;
|
||||
}
|
||||
@@ -383,10 +387,10 @@ fn shuf_exec(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_range(input_range: &str) -> Result<RangeInclusive<usize>, String> {
|
||||
fn parse_range(input_range: &str) -> Result<RangeInclusive<u64>, String> {
|
||||
if let Some((from, to)) = input_range.split_once('-') {
|
||||
let begin = from.parse::<usize>().map_err(|e| e.to_string())?;
|
||||
let end = to.parse::<usize>().map_err(|e| e.to_string())?;
|
||||
let begin = from.parse::<u64>().map_err(|e| e.to_string())?;
|
||||
let end = to.parse::<u64>().map_err(|e| e.to_string())?;
|
||||
if begin <= end || begin == end + 1 {
|
||||
Ok(begin..=end)
|
||||
} else {
|
||||
@@ -398,48 +402,36 @@ fn parse_range(input_range: &str) -> Result<RangeInclusive<usize>, String> {
|
||||
}
|
||||
|
||||
enum WrappedRng {
|
||||
RngFile(rand_read_adapter::ReadRng<File>),
|
||||
RngDefault(rand::rngs::ThreadRng),
|
||||
RngFile(compat_random_source::RandomSourceAdapter<BufReader<File>>),
|
||||
}
|
||||
|
||||
impl WrappedRng {
|
||||
fn get_error_cell(&self) -> Option<rand_read_adapter::ErrorCell> {
|
||||
if let Self::RngFile(adapter) = self {
|
||||
Some(adapter.error.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn check_error(error_cell: Option<&rand_read_adapter::ErrorCell>) -> UResult<()> {
|
||||
if let Some(cell) = error_cell {
|
||||
if let Some(err) = cell.take() {
|
||||
return Err(err.map_err_context(|| translate!("shuf-error-read-random-bytes")));
|
||||
fn choose<T: Copy>(&mut self, vals: &[T]) -> UResult<T> {
|
||||
match self {
|
||||
Self::RngDefault(rng) => Ok(*vals.choose(rng).unwrap()),
|
||||
Self::RngFile(adapter) => {
|
||||
assert!(!vals.is_empty());
|
||||
let idx = adapter.get_value(vals.len() as u64 - 1)? as usize;
|
||||
Ok(vals[idx])
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl RngCore for WrappedRng {
|
||||
fn next_u32(&mut self) -> u32 {
|
||||
fn shuffle<'a, T>(&mut self, vals: &'a mut [T], amount: usize) -> UResult<&'a mut [T]> {
|
||||
match self {
|
||||
Self::RngFile(r) => r.next_u32(),
|
||||
Self::RngDefault(r) => r.next_u32(),
|
||||
Self::RngDefault(rng) => Ok(vals.partial_shuffle(rng, amount).0),
|
||||
Self::RngFile(adapter) => adapter.shuffle(vals, amount),
|
||||
}
|
||||
}
|
||||
|
||||
fn next_u64(&mut self) -> u64 {
|
||||
fn choose_from_range(&mut self, range: RangeInclusive<u64>) -> UResult<u64> {
|
||||
match self {
|
||||
Self::RngFile(r) => r.next_u64(),
|
||||
Self::RngDefault(r) => r.next_u64(),
|
||||
}
|
||||
}
|
||||
|
||||
fn fill_bytes(&mut self, dest: &mut [u8]) {
|
||||
match self {
|
||||
Self::RngFile(r) => r.fill_bytes(dest),
|
||||
Self::RngDefault(r) => r.fill_bytes(dest),
|
||||
Self::RngDefault(rng) => Ok(rng.random_range(range)),
|
||||
Self::RngFile(adapter) => {
|
||||
let offset = adapter.get_value(*range.end() - *range.start())?;
|
||||
Ok(*range.start() + offset)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -859,3 +859,162 @@ fn write_errors_are_reported() {
|
||||
.no_stdout()
|
||||
.stderr_is("shuf: write failed: No space left on device\n");
|
||||
}
|
||||
|
||||
// On 32-bit platforms, if we cast carelessly, this will give no output.
|
||||
#[test]
|
||||
fn test_head_count_does_not_overflow_file() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
|
||||
at.append("input.txt", "hello\n");
|
||||
|
||||
ucmd.arg(format!("-n{}", u64::from(u32::MAX) + 1))
|
||||
.arg("input.txt")
|
||||
.succeeds()
|
||||
.stdout_is("hello\n")
|
||||
.no_stderr();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_head_count_does_not_overflow_args() {
|
||||
new_ucmd!()
|
||||
.arg(format!("-n{}", u64::from(u32::MAX) + 1))
|
||||
.arg("-e")
|
||||
.arg("goodbye")
|
||||
.succeeds()
|
||||
.stdout_is("goodbye\n")
|
||||
.no_stderr();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_head_count_does_not_overflow_range() {
|
||||
new_ucmd!()
|
||||
.arg(format!("-n{}", u64::from(u32::MAX) + 1))
|
||||
.arg("-i1-1")
|
||||
.succeeds()
|
||||
.stdout_is("1\n")
|
||||
.no_stderr();
|
||||
}
|
||||
|
||||
// Test reproducibility and compatibility of --random-source.
|
||||
// These hard-coded results match those of GNU shuf. They should not be changed.
|
||||
|
||||
#[test]
|
||||
fn test_gnu_compat_range_repeat() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xfb\x83\x8f\x21\x9b\x3c\x2d\xc5\x73\xa5\x58\x6c\x54\x2f\x59\xf8",
|
||||
);
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.arg("-r")
|
||||
.arg("-i1-99")
|
||||
.fails_with_code(1)
|
||||
.stderr_is("shuf: end of random source\n")
|
||||
.stdout_is("38\n30\n10\n26\n23\n61\n46\n99\n75\n43\n10\n89\n10\n44\n24\n59\n22\n51\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gnu_compat_args_no_repeat() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xd1\xfd\xb9\x9a\xf5\x81\x71\x42\xf9\x7a\x59\x79\xd4\x9c\x8c\x7d",
|
||||
);
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.arg("-e")
|
||||
.args(&["1", "2", "3", "4", "5", "6", "7"][..])
|
||||
.succeeds()
|
||||
.no_stderr()
|
||||
.stdout_is("7\n1\n2\n5\n3\n4\n6\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gnu_compat_from_stdin() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xd1\xfd\xb9\x9a\xf5\x81\x71\x42\xf9\x7a\x59\x79\xd4\x9c\x8c\x7d",
|
||||
);
|
||||
|
||||
at.append("input.txt", "1\n2\n3\n4\n5\n6\n7\n");
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.set_stdin(at.open("input.txt"))
|
||||
.succeeds()
|
||||
.no_stderr()
|
||||
.stdout_is("7\n1\n2\n5\n3\n4\n6\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gnu_compat_from_file() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xd1\xfd\xb9\x9a\xf5\x81\x71\x42\xf9\x7a\x59\x79\xd4\x9c\x8c\x7d",
|
||||
);
|
||||
|
||||
at.append("input.txt", "1\n2\n3\n4\n5\n6\n7\n");
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.arg("input.txt")
|
||||
.succeeds()
|
||||
.no_stderr()
|
||||
.stdout_is("7\n1\n2\n5\n3\n4\n6\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gnu_compat_limited_from_file() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xd1\xfd\xb9\x9a\xf5\x81\x71\x42\xf9\x7a\x59\x79\xd4\x9c\x8c\x7d",
|
||||
);
|
||||
|
||||
at.append("input.txt", "1\n2\n3\n4\n5\n6\n7\n");
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.arg("-n5")
|
||||
.arg("input.txt")
|
||||
.succeeds()
|
||||
.no_stderr()
|
||||
.stdout_is("7\n1\n2\n5\n3\n");
|
||||
}
|
||||
|
||||
// This specific case causes GNU to give different results than other modes.
|
||||
#[ignore = "disabled until fixed"]
|
||||
#[test]
|
||||
fn test_gnu_compat_limited_from_stdin() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xd1\xfd\xb9\x9a\xf5\x81\x71\x42\xf9\x7a\x59\x79\xd4\x9c\x8c\x7d",
|
||||
);
|
||||
|
||||
at.append("input.txt", "1\n2\n3\n4\n5\n6\n7\n");
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.arg("-n7")
|
||||
.set_stdin(at.open("input.txt"))
|
||||
.succeeds()
|
||||
.no_stderr()
|
||||
.stdout_is("6\n5\n1\n3\n2\n7\n4\n");
|
||||
}
|
||||
|
||||
// We haven't reverse-engineered GNU's nonrepeating integer sampling yet.
|
||||
#[ignore = "disabled until fixed"]
|
||||
#[test]
|
||||
fn test_gnu_compat_range_no_repeat() {
|
||||
let (at, mut ucmd) = at_and_ucmd!();
|
||||
at.append_bytes(
|
||||
"random_bytes.bin",
|
||||
b"\xd1\xfd\xb9\x9a\xf5\x81\x71\x42\xf9\x7a\x59\x79\xd4\x9c\x8c\x7d",
|
||||
);
|
||||
|
||||
ucmd.arg("--random-source=random_bytes.bin")
|
||||
.arg("-i1-10")
|
||||
.succeeds()
|
||||
.no_stderr()
|
||||
.stdout_is("10\n2\n8\n7\n3\n9\n6\n5\n1\n4\n");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user