diff --git a/src/shuf/randint.v b/src/shuf/randint.v new file mode 100644 index 00000000..8c7f58aa --- /dev/null +++ b/src/shuf/randint.v @@ -0,0 +1,148 @@ +module main + +import os +import rand +import common + +// A port of GNU coreutils' randint/randperm (gl/lib/randint.c and +// gl/lib/randperm.c). +// +// This has to consume the --random-source bytes in exactly the same order and +// quantity as GNU does, otherwise the same file produces a different +// permutation. The algorithm below is a direct translation, including the way +// leftover entropy is carried across calls, because that carry is observable: +// shuf -i 1-5 over a one byte source still consumes a single byte, and only +// because randnum/randmax keep the unused bits. +struct RandintSource { +mut: + // randnum is a buffered random integer, uniformly distributed over + // 0..=randmax. If randmax is 0 then randnum must be 0. + randnum u64 + randmax u64 + // The --random-source file, or none to draw from the system PRNG. + source ?os.File + source_name string +} + +fn new_randint_source(source_name string) &RandintSource { + mut s := &RandintSource{} + if source_name.len > 0 { + s.source = os.open(source_name) or { + common.exit_with_error_message(app_name, + '${source_name}: ${posix_msg()}') + } + s.source_name = source_name + } + return s +} + +// fill puts buf.len random bytes into buf. +fn (mut s RandintSource) fill(mut buf []u8) { + if mut f := s.source { + // GNU reads exactly the requested number of bytes and reports end of + // file once the source runs dry. + n := f.read(mut buf) or { + if err == os.Eof{} { + common.exit_with_error_message(app_name, + '‘${s.source_name}’: end of file') + } + common.exit_with_error_message(app_name, '‘${s.source_name}’: read error') + } + if n < buf.len { + common.exit_with_error_message(app_name, '‘${s.source_name}’: end of file') + } + } else { + rand.read(mut buf) + } +} + +// genmax consumes random bytes to produce a value uniformly distributed over +// 0..=limit. This is randint_genmax() in GNU's randint.c. +fn (mut s RandintSource) genmax(limit u64) u64 { + mut randnum := s.randnum + mut randmax := s.randmax + choices := limit + 1 + for { + if randmax < limit { + // Work out how many bytes are needed to cover the new limit. + mut count := 0 + mut rmax := randmax + for { + rmax = (rmax << 8) + 255 + count++ + if !(rmax < limit) { + break + } + } + mut buf := []u8{len: count} + s.fill(mut buf) + + // Append the bytes to randnum, and 255 to randmax, until randmax + // covers the limit. Up to 8 bits of information are dropped here, + // which GNU accepts as not worth the extra bookkeeping. + count = 0 + for { + randnum = (randnum << 8) + u64(buf[count]) + randmax = (randmax << 8) + 255 + count++ + if !(randmax < limit) { + break + } + } + } + + if randmax == limit { + s.randnum = 0 + s.randmax = 0 + return randnum + } + + // limit < randmax, so randnum % choices is only fair while randnum + // stays inside an integral multiple of choices. Outside that range the + // attempt is discarded, but the partial randomness is kept so no byte + // is thrown away for nothing. + excess_choices := randmax - limit + unusable_choices := excess_choices % choices + last_usable_choice := randmax - unusable_choices + reduced_randnum := randnum % choices + if randnum <= last_usable_choice { + s.randnum = randnum / choices + s.randmax = excess_choices / choices + return reduced_randnum + } + randnum = reduced_randnum + randmax = unusable_choices - 1 + } + panic('unreachable') +} + +// choose returns a value uniformly distributed over 0..=n-1, matching GNU's +// randint_choose(), which is genmax(n - 1). +fn (mut s RandintSource) choose(n u64) u64 { + return s.genmax(n - 1) +} + +// randperm_new returns the first h elements of a random permutation of +// 0..=n-1, matching GNU's randperm_new(). +// +// GNU switches to a hash based "sparse" representation for large, sparse +// permutations. That is purely a memory optimisation: it performs the same +// swaps in the same order, so the permutation is identical and is not needed +// here. +fn randperm_new(mut s RandintSource, h int, n int) []u64 { + if h == 0 { + return []u64{} + } + if h == 1 { + return [s.choose(u64(n))] + } + mut v := []u64{len: n, init: 0} + for i in 0 .. n { + v[i] = u64(i) + } + for i in 0 .. h { + j := i + int(s.choose(u64(n - i))) + v[i], v[j] = v[j], v[i] + } + return v[..h] +} diff --git a/src/shuf/shuf.v b/src/shuf/shuf.v index 5fce1649..9af0e5e0 100644 --- a/src/shuf/shuf.v +++ b/src/shuf/shuf.v @@ -1,201 +1,229 @@ module main import os -import io -import rand +import flag +import strconv import common const app_name = 'shuf' -const app_description = 'Shuffles its input by outputting a random permutation of its input lines' + +const app_description = 'Write a random permutation of the input lines to standard output.' struct Settings { - echo bool - input_range string - head_count int - output string - random_source string - repeat bool - zero_terminated bool - fnames []string +mut: + echo bool + input_range string + head_count u64 + head_count_given bool + output string + output_given bool + random_source string + random_source_given bool + repeat bool + zero_terminated bool + operands []string } -// main fn main() { - shuf(args())! + settings := args() + run(settings) } -fn shuf(settings Settings) ! { +// read_lines returns the lines of f, keeping empty lines. A final line without +// a terminator is returned as is; callers add the terminator back. +fn read_lines(mut f os.File, delim u8) []string { mut lines := []string{} - - lines = set_lines(lines.clone(), settings) - lines = shuffle_lines(lines, settings) - - if settings.repeat { - if settings.output.len > 0 { - mut file := os.open_file(settings.output, 'w+', 0o666)! - for { - output_lines_file(lines, settings, mut file)! - } - file.close() - } else { - for { - output_lines(lines, settings) + mut buf := []u8{len: 64 * 1024} + mut line := []u8{} + for { + n := f.read(mut buf) or { break } + for i in 0 .. n { + b := buf[i] + if b == delim { + lines << line.bytestr() + line = []u8{} + } else { + line << b } } - } else { - if settings.output.len > 0 { - mut file := os.open_file(settings.output, 'w+', 0o666)! - output_lines_file(lines, settings, mut file)! - file.close() - } else { - output_lines(lines, settings) - } } + if line.len > 0 { + lines << line.bytestr() + } + return lines } -fn output_lines(lines []string, settings Settings) { - for i in 0 .. lines.len { - if settings.head_count > 0 && i >= settings.head_count { - break - } - if settings.zero_terminated { - print(lines[i]) - } else { - println(lines[i]) +fn open_operands(operands []string, delim u8) []string { + mut lines := []string{} + for operand in operands { + mut file := os.stdin() + if operand != '-' { + file = os.open(operand) or { + common.exit_with_error_message(app_name, + '${operand}: ${posix_msg()}') + } } + lines << read_lines(mut file, delim) + file.close() } + return lines } -fn output_lines_file(lines []string, settings Settings, mut file os.File) ! { - for i in 0 .. lines.len { - if settings.head_count > 0 && i >= settings.head_count { - break - } - if settings.zero_terminated { - file.write_string('${lines[i]}')! - } else { - file.write_string('${lines[i]}\n')! - } +// parse_range splits LO-HI and returns lo and hi, rejecting anything GNU +// rejects. +fn parse_range(arg string) []u64 { + parts := arg.split('-') + if parts.len != 2 || parts[0].len == 0 || parts[1].len == 0 { + common.exit_with_error_message(app_name, 'invalid input range: ‘${arg}’') } + lo := strconv.atou64(parts[0]) or { + common.exit_with_error_message(app_name, 'invalid input range: ‘${arg}’') + } + hi := strconv.atou64(parts[1]) or { + common.exit_with_error_message(app_name, 'invalid input range: ‘${arg}’') + } + if lo > hi { + common.exit_with_error_message(app_name, 'invalid input range: ‘${arg}’') + } + return [lo, hi] } -fn shuffle_lines(lines []string, settings Settings) []string { - mut new_lines := lines.clone() - if settings.random_source.len > 0 { - mut file := os.File{} - file = os.open(settings.random_source) or { - eprintln('${app_name}: ${settings.random_source}: No such file') - exit(1) - } - mut bytes := io.read_all(io.ReadAllConfig{ reader: file, read_to_end_of_stream: true }) or { - eprintln("${app_name}: ${settings.random_source}: Can't read file") - exit(1) - } - mut seed := u32(0) - for b in bytes { - seed += b.hex().u32() +// posix_msg returns the strerror text for the failure that just happened, in +// the form GNU prints it. os.error_posix() returns an error value whose +// interpolation carries V's trailing "; code: N". +fn posix_msg() string { + return os.error_posix().msg() +} + +fn run(settings Settings) { + if settings.echo && settings.input_range.len > 0 { + common.exit_with_error_message(app_name, 'cannot combine -e and -i options') + } + if !settings.echo && settings.input_range.len == 0 && settings.operands.len > 1 { + common.exit_with_error_message(app_name, 'extra operand ‘${settings.operands[1]}’') + } + if settings.input_range.len > 0 && settings.operands.len > 0 { + common.exit_with_error_message(app_name, 'extra operand ‘${settings.operands[0]}’') + } + + delim := if settings.zero_terminated { u8(0) } else { u8(`\n`) } + mut source := new_randint_source(settings.random_source) + mut out := os.stdout() + + if settings.output_given { + out = os.create(settings.output) or { + common.exit_with_error_message(app_name, + '${settings.output}: ${posix_msg()}') } - rand.seed([u32(0), seed]) + } - for i in 0 .. new_lines.len { - tmp := new_lines[i] - random := rand.intn(new_lines.len) or { - common.exit_with_error_message(app_name, err.msg()) - } - new_lines[i] = new_lines[random] - new_lines[random] = tmp + // Collect the input up front, except when -n 0 means nothing is needed. + mut lines := []string{} + mut lo := u64(0) + if settings.echo { + lines = settings.operands.clone() + } else if settings.input_range.len > 0 { + r := parse_range(settings.input_range) + lo = r[0] + hi := r[1] + for i in lo .. hi + 1 { + lines << i.str() } } else { - for i in 0 .. new_lines.len { - tmp := new_lines[i] - random := rand.intn(new_lines.len) or { - common.exit_with_error_message(app_name, err.msg()) - } - new_lines[i] = new_lines[random] - new_lines[random] = tmp - } + operands := if settings.operands.len == 0 { ['-'] } else { settings.operands } + lines = open_operands(operands, delim) } - return new_lines -} + // -n 0 means output nothing; without -n every line is used. + count := if settings.head_count_given { settings.head_count } else { u64(lines.len) } + if count == 0 { + return + } -fn set_lines(lines []string, settings Settings) []string { - mut new_lines := lines.clone() - mut fnames := settings.fnames.clone() - input_range := settings.input_range.split('-') - if input_range.len == 2 { - for i in input_range[0].int() .. input_range[1].int() + 1 { - new_lines << i.str() - } - } else { - if fnames.len < 1 { - fnames = ['-'] + if settings.repeat { + if lines.len == 0 { + common.exit_with_error_message(app_name, 'no lines to repeat') } - - for fname in fnames { - if settings.echo { - new_lines << fname - } else { - new_lines = register_lines_by_file(new_lines.clone(), fname, - settings.zero_terminated) + for _ in 0 .. count { + j := int(source.choose(u64(lines.len))) + out.write_string(lines[j] + delim.ascii_str()) or { + common.exit_with_error_message(app_name, err.msg()) } } + return } - return new_lines -} - -fn register_lines_by_file(lines []string, fname string, zero_terminated bool) []string { - mut new_lines := lines.clone() - mut file := os.File{} - if fname == '-' { - file = os.stdin() - } else { - file = os.open(fname) or { - eprintln('${app_name}: ${fname}: No such file or directory') - exit(1) - } + ahead := if count < u64(lines.len) { int(count) } else { lines.len } + if ahead == 0 { + return } - if zero_terminated { - mut bytes := io.read_all(io.ReadAllConfig{ reader: file, read_to_end_of_stream: true }) or { - eprintln("${app_name}: ${fname}: Can't read file") - exit(1) + permutation := randperm_new(mut source, ahead, lines.len) + for idx in permutation { + line := lines[int(idx)] + if line.len > 0 || !settings.zero_terminated { + out.write_string(line) or { + common.exit_with_error_message(app_name, err.msg()) + } } - - new_lines << bytes.bytestr() - } else { - mut br := io.new_buffered_reader(io.BufferedReaderConfig{ reader: file }) - for { - line := br.read_line() or { break } - if line.len > 0 { - new_lines << line + if !line.ends_with(delim.ascii_str()) { + out.write_string(delim.ascii_str()) or { + common.exit_with_error_message(app_name, err.msg()) } } } - - return new_lines } -// args +// The _opt flag variants are used throughout so that V does not append a +// "(default ...)" note to every option in --help, which GNU does not print. fn args() Settings { mut fp := common.flag_parser(os.args) fp.application(app_name) fp.description(app_description) + fp.usage_example('[OPTION]... [FILE]') + fp.usage_example('-e [OPTION]... [ARG]...') + fp.usage_example('-i LO-HI [OPTION]...') + fp.description('') + fp.description('With no FILE, or when FILE is -, read standard input.') + + mut settings := Settings{} + settings.echo = fp.bool_opt('echo', `e`, 'treat each ARG as an input line', + flag.FlagConfig{}) or { false } + settings.input_range = fp.string_opt('input-range', `i`, + 'treat each number LO through HI as an input line', flag.FlagConfig{ + val_desc: 'LO-HI' + }) or { '' } + + head_count := fp.string_opt('head-count', `n`, 'output at most COUNT lines', + flag.FlagConfig{ + val_desc: 'COUNT' + }) or { '' } + settings.head_count_given = head_count.len > 0 + if settings.head_count_given { + settings.head_count = strconv.atou64(head_count) or { + common.exit_with_error_message(app_name, 'invalid line count: ‘${head_count}’') + } + } - echo := fp.bool('echo', `e`, false, 'Treat each command-line operand as an input line') - input_range := fp.string('input-range', `i`, '', - 'Act as if input came from a file containing the range of unsigned decimal integers lo…hi, one per line') - head_count := fp.int('head-count', `n`, 0, - 'Output at most count lines. By default, all input lines are output') - output := fp.string('output', `o`, '', 'Write output to output-file instead of standard output') - random_source := fp.string('random-source', 0, '', - 'Use file as a source of random data used to determine which permutation to generate') - repeat := fp.bool('repeat', `r`, false, 'Repeat output values') - zero_terminated := fp.bool('zero-terminated', `z`, false, - 'Delimit items with a zero byte rather than a newline') - - fnames := fp.remaining_parameters() - - return Settings{echo, input_range, head_count, output, random_source, repeat, zero_terminated, fnames} + output := fp.string_opt('output', `o`, 'write result to FILE instead of standard output', + flag.FlagConfig{ + val_desc: 'FILE' + }) or { '' } + settings.output_given = output.len > 0 + settings.output = output + + random_source := fp.string_opt('random-source', 0, 'get random bytes from FILE', + flag.FlagConfig{ + val_desc: 'FILE' + }) or { '' } + settings.random_source_given = random_source.len > 0 + settings.random_source = random_source + + settings.repeat = fp.bool_opt('repeat', `r`, 'output lines can be repeated', + flag.FlagConfig{}) or { false } + settings.zero_terminated = fp.bool_opt('zero-terminated', `z`, + 'line delimiter is NUL, not newline', flag.FlagConfig{}) or { false } + + settings.operands = fp.remaining_parameters() + return settings } diff --git a/src/shuf/shuf_test.v b/src/shuf/shuf_test.v index fe7b2d2c..81fdc1b4 100644 --- a/src/shuf/shuf_test.v +++ b/src/shuf/shuf_test.v @@ -26,14 +26,13 @@ fn test_file() { } fn test_zero_terminated_echo() { - res := os.execute('${executable_under_test} -z -e aa bb') - assert res.output == 'aabb' || res.output == 'bbaa' + os.write_file(test_txt_path, 'hello\nworld!')! + rig.assert_same_results('-z -e --random-source ${test_txt_path} aa bb') } fn test_zero_terminated_file() { - os.write_file(test_txt_path, 'hello\nworld!')! - res := os.execute('${executable_under_test} -z ${test_txt_path}') - assert res.output == 'hello${eol}world!' + os.write_file(test_txt_path, 'aa\x00bb\x00')! + rig.assert_same_results('-z --random-source ${test_txt_path} ${test_txt_path}') } fn test_head_count() { @@ -47,10 +46,75 @@ fn test_input_range() { assert res.output.split_into_lines().len == 10 } +// The permutation produced by a given --random-source file is fully +// determined, so it is compared against the platform shuf rather than being +// hardcoded: the exact sequence of bytes consumed changed between GNU 8.32 and +// later releases, and the platform utility is the reference this test rig +// already trusts everywhere else. fn test_random_source() { os.write_file(test_txt_path, 'hello\nworld!')! - res := os.execute('${executable_under_test} -i 1-5 --random-source ${test_txt_path}') - assert res.output == '1${eol}4${eol}5${eol}2${eol}3${eol}' + rig.assert_same_results('-i 1-5 --random-source ${test_txt_path}') +} + +fn test_random_source_lines() { + os.write_file(test_txt_path, 'a\nb\nc\nd\ne\n')! + rig.assert_same_results('--random-source ${test_txt_path} ${test_txt_path}') +} + +fn test_random_source_head_count() { + os.write_file(test_txt_path, 'a\nb\nc\nd\ne\n')! + rig.assert_same_results('-n 3 --random-source ${test_txt_path} ${test_txt_path}') +} + +fn test_random_source_exhausted() { + os.write_file(test_txt_path, 'hello\nworld!')! + // Only the exit status is compared: common.exit_with_error_message adds a + // "Try 'shuf --help'" line that GNU does not print. + rig.cmd.expected_failure('-i 1-1000 --random-source ${test_txt_path}')! +} + +fn test_random_source_missing_file() { + rig.cmd.expected_failure('-i 1-5 --random-source ${rig.temp_dir}/no-such-source')! +} + +fn test_repeat_with_random_source() { + rig.assert_same_results('-r -n 6 -e --random-source ${test_txt_path} aa bb') +} + +// Empty lines are part of the input; the previous implementation dropped them. +fn test_empty_lines_are_kept() { + os.write_file(test_txt_path, 'x\n\ny\n')! + rig.assert_same_results('--random-source ${test_txt_path} ${test_txt_path}') +} + +fn test_output_file_matches() { + os.write_file(test_txt_path, 'hello\nworld!')! + rig.call_new('-i 1-5 --random-source ${test_txt_path} -o out1') + rig.call_orig('-i 1-5 --random-source ${test_txt_path} -o out2') + assert os.read_file('out1')! == os.read_file('out2')! + os.rm('out1')! + os.rm('out2')! +} + +fn test_head_count_zero_outputs_nothing() { + os.write_file(test_txt_path, 'a\nb\nc\n')! + res := os.execute('${executable_under_test} -n 0 ${test_txt_path}') + assert res.exit_code == 0 + assert res.output == '' +} + +fn test_input_range_errors() { + for bad in ['5-1', 'abc', '1-x', '-', '1-2-3'] { + rig.cmd.expected_failure('-i ${bad}')! + } +} + +fn test_extra_operand() { + rig.cmd.expected_failure('a b')! +} + +fn test_echo_and_input_range_conflict() { + rig.cmd.expected_failure('-e -i 1-5')! } fn test_unknown_option() ? {