diff --git a/Cargo.lock b/Cargo.lock index ce9a51d97b0..dfad6f922a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4312,6 +4312,7 @@ dependencies = [ "clap", "codspeed-divan-compat", "fluent", + "memchr", "nom 8.0.0", "rustix", "uucore", diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock index acbe03a8ece..9416a44a368 100644 --- a/fuzz/Cargo.lock +++ b/fuzz/Cargo.lock @@ -1954,6 +1954,7 @@ dependencies = [ "bytecount", "clap", "fluent", + "memchr", "nom", "uucore", ] diff --git a/src/uu/tr/Cargo.toml b/src/uu/tr/Cargo.toml index 8077597d903..d17b671ca68 100644 --- a/src/uu/tr/Cargo.toml +++ b/src/uu/tr/Cargo.toml @@ -20,6 +20,7 @@ doctest = false bytecount = { workspace = true, features = ["runtime-dispatch-simd"] } clap = { workspace = true } fluent = { workspace = true } +memchr = { workspace = true } nom = { workspace = true } uucore = { workspace = true, features = ["fs", "signals"] } diff --git a/src/uu/tr/benches/tr_bench.rs b/src/uu/tr/benches/tr_bench.rs index 8f751c72581..a26fc3bc11f 100644 --- a/src/uu/tr/benches/tr_bench.rs +++ b/src/uu/tr/benches/tr_bench.rs @@ -73,6 +73,13 @@ mod benches { let data = text_data::generate_by_size(SIZE_MB, 80); bench_tr_with_stdin(bencher, &data, &["-d", "a-z"]); } + + /// Delete a single character (the newlines). + #[divan::bench] + fn tr_delete_single_char(bencher: Bencher) { + let data = text_data::generate_by_size(SIZE_MB, 80); + bench_tr_with_stdin(bencher, &data, &["-d", "\\n"]); + } } fn main() { diff --git a/src/uu/tr/src/operation.rs b/src/uu/tr/src/operation.rs index 813a94b5215..36a2acc0728 100644 --- a/src/uu/tr/src/operation.rs +++ b/src/uu/tr/src/operation.rs @@ -29,7 +29,9 @@ use uucore::show_warning; /// Common trait for operations that can process chunks of data pub trait ChunkProcessor { - fn process_chunk(&self, input: &[u8], output: &mut Vec); + /// Return the bytes to write: `input` itself when it is left unchanged, + /// `output` otherwise. + fn process_chunk<'a>(&self, input: &'a [u8], output: &'a mut Vec) -> &'a [u8]; } #[derive(Debug, Clone)] @@ -809,13 +811,20 @@ fn set_to_bitmap(set: &[u8]) -> [bool; 256] { #[derive(Debug)] pub struct DeleteOperation { - pub(crate) delete_table: [bool; 256], + pub(crate) keep_table: [bool; 256], + /// The byte to delete, when it is the only one. + single_delete: Option, } impl DeleteOperation { pub fn new(set: Vec) -> Self { + use crate::simd::find_single_change; + + let keep_table = set_to_bitmap(&set).map(|delete| !delete); + let single_delete = find_single_change(&keep_table, |_, &keep| !keep).map(|(b, _)| b); Self { - delete_table: set_to_bitmap(&set), + keep_table, + single_delete, } } } @@ -823,27 +832,20 @@ impl DeleteOperation { impl SymbolTranslator for DeleteOperation { fn translate(&mut self, current: u8) -> Option { // keep if not present in the delete set - (!self.delete_table[current as usize]).then_some(current) + self.keep_table[current as usize].then_some(current) } } impl ChunkProcessor for DeleteOperation { - fn process_chunk(&self, input: &[u8], output: &mut Vec) { - use crate::simd::{find_single_change, process_single_delete}; + fn process_chunk<'a>(&self, input: &'a [u8], output: &'a mut Vec) -> &'a [u8] { + use crate::simd::{process_delete, process_single_delete}; - // Check if this is single character deletion - if let Some((delete_char, _)) = - find_single_change(&self.delete_table, |_, &should_delete| should_delete) - { - process_single_delete(input, output, delete_char); + if let Some(delete_char) = self.single_delete { + process_single_delete(input, output, delete_char, &self.keep_table) } else { // Standard deletion - output.extend( - input - .iter() - .filter(|&&b| !self.delete_table[b as usize]) - .copied(), - ); + process_delete(input, output, &self.keep_table); + output } } } @@ -851,6 +853,8 @@ impl ChunkProcessor for DeleteOperation { #[derive(Debug)] pub struct TranslateOperation { pub(crate) translation_table: [u8; 256], + /// The byte to replace and its replacement, when it is the only one. + single_change: Option<(u8, u8)>, } impl TranslateOperation { @@ -867,16 +871,26 @@ impl TranslateOperation { translation_table[from as usize] = to; } - Ok(Self { translation_table }) + Ok(Self::from_table(translation_table)) } else if set1.is_empty() && set2.is_empty() { // Identity mapping for empty sets - Ok(Self { translation_table }) + Ok(Self::from_table(translation_table)) } else { // Raised against the solved sets rather than what was typed, so // there is nothing to point a caret at. Err(BadSequence::EmptySet2WhenNotTruncatingSet1) } } + + fn from_table(translation_table: [u8; 256]) -> Self { + use crate::simd::find_single_change; + + let single_change = find_single_change(&translation_table, |i, &val| val != i as u8); + Self { + translation_table, + single_change, + } + } } impl SymbolTranslator for TranslateOperation { @@ -886,18 +900,16 @@ impl SymbolTranslator for TranslateOperation { } impl ChunkProcessor for TranslateOperation { - fn process_chunk(&self, input: &[u8], output: &mut Vec) { - use crate::simd::{find_single_change, process_single_char_replace}; + fn process_chunk<'a>(&self, input: &'a [u8], output: &'a mut Vec) -> &'a [u8] { + use crate::simd::process_single_char_replace; - // Check if this is a simple single-character translation - if let Some((source, target)) = - find_single_change(&self.translation_table, |i, &val| val != i as u8) - { + if let Some((source, target)) = self.single_change { // Use SIMD-optimized single character replacement - process_single_char_replace(input, output, source, target); + process_single_char_replace(input, output, source, target) } else { // Standard translation using table lookup output.extend(input.iter().map(|&b| self.translation_table[b as usize])); + output } } } diff --git a/src/uu/tr/src/simd.rs b/src/uu/tr/src/simd.rs index 4af7760a415..0f31e242008 100644 --- a/src/uu/tr/src/simd.rs +++ b/src/uu/tr/src/simd.rs @@ -28,16 +28,17 @@ where /// SIMD-optimized single character replacement #[inline] -pub fn process_single_char_replace( - input: &[u8], - output: &mut Vec, +pub fn process_single_char_replace<'a>( + input: &'a [u8], + output: &'a mut Vec, source_char: u8, target_char: u8, -) { +) -> &'a [u8] { let count = bytecount::count(input, source_char); if count == 0 { - output.extend_from_slice(input); - } else if count == input.len() { + return input; + } + if count == input.len() { output.resize(output.len() + input.len(), target_char); } else { output.extend( @@ -46,17 +47,70 @@ pub fn process_single_char_replace( .map(|&b| if b == source_char { target_char } else { b }), ); } + output } /// SIMD-optimized delete operation for single character -pub fn process_single_delete(input: &[u8], output: &mut Vec, delete_char: u8) { +/// +/// `keep` must be false for `delete_char` only. +pub fn process_single_delete<'a>( + input: &'a [u8], + output: &'a mut Vec, + delete_char: u8, + keep: &[bool; 256], +) -> &'a [u8] { let count = bytecount::count(input, delete_char); if count == 0 { - output.extend_from_slice(input); + return input; + } + if count < input.len() / 128 { + // Below one match per 128 bytes, copying the runs between matches + // beats `process_delete`. + let mut start = 0; + for pos in memchr::memchr_iter(delete_char, input) { + output.extend_from_slice(&input[start..pos]); + start = pos + 1; + } + output.extend_from_slice(&input[start..]); } else if count < input.len() { - output.extend(input.iter().filter(|&&b| b != delete_char).copied()); + process_delete(input, output, keep); } // If count == input.len(), all deleted, output nothing + output +} + +/// Append to `output` the bytes of `input` whose `keep` entry is true. +pub fn process_delete(input: &[u8], output: &mut Vec, keep: &[bool; 256]) { + // The index is always below `BLOCK` (a power of two), so the modulo is a + // mask that only serves to drop the bounds check. + const BLOCK: usize = 1024; + // Below one kept byte in `FEW`, a branch per byte is well predicted and + // stores less. + const FEW: usize = 64; + let mut block = [0; BLOCK]; + // Guess for the first block from its start, then go by the previous one. + let start = &input[..input.len().min(256)]; + let mut few_kept = start.iter().filter(|&&b| keep[b as usize]).count() * FEW < start.len(); + for chunk in input.chunks(BLOCK) { + let mut kept = 0; + if few_kept { + // Only store the kept bytes. + for &b in chunk { + if keep[b as usize] { + block[kept % BLOCK] = b; + kept += 1; + } + } + } else { + // Store every byte and only advance past kept ones: no branch. + for &b in chunk { + block[kept % BLOCK] = b; + kept += usize::from(keep[b as usize]); + } + } + output.extend_from_slice(&block[..kept]); + few_kept = kept * FEW < chunk.len(); + } } /// Unified I/O processing for all operations @@ -79,10 +133,10 @@ where }; output_buf.clear(); - processor.process_chunk(&buf[..length], &mut output_buf); + let processed = processor.process_chunk(&buf[..length], &mut output_buf); - if !output_buf.is_empty() { - write_output(output, &output_buf)?; + if !processed.is_empty() { + write_output(output, processed)?; } } @@ -90,7 +144,9 @@ where } /// Helper function to handle platform-specific write operations -#[inline] +// Kept out of line: inlined into `translate_input`, the raw `write` made the +// translator state go to memory on every byte, doubling the time of `tr -s`. +#[inline(never)] pub fn write_output(output: &mut W, buf: &[u8]) -> UResult<()> { #[cfg(not(windows))] return output diff --git a/src/uu/tr/src/tr.rs b/src/uu/tr/src/tr.rs index 05bdfa7ad44..5aadc652683 100644 --- a/src/uu/tr/src/tr.rs +++ b/src/uu/tr/src/tr.rs @@ -97,7 +97,12 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> { let stdin = stdin(); let mut locked_stdin = stdin.lock(); - let mut locked_stdout = stdout().lock(); + // Write straight to the file descriptor: `Stdout` is line buffered, which + // costs a search for the last newline and an extra write per chunk. + #[cfg(any(unix, target_os = "wasi"))] + let mut output = uucore::io::RawWriter(stdout()); + #[cfg(not(any(unix, target_os = "wasi")))] + let mut output = stdout().lock(); // According to the man page: translating only happens if deleting or if a second set is given let translating = !delete_flag && sets.len() > 1; @@ -131,27 +136,27 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> { let delete_op = DeleteOperation::new(set1); let squeeze_op = SqueezeOperation::new(set2); let op = delete_op.chain(squeeze_op); - translate_input(&mut locked_stdin, &mut locked_stdout, op)?; + translate_input(&mut locked_stdin, &mut output, op)?; } else { let op = DeleteOperation::new(set1); - process_input(&mut locked_stdin, &mut locked_stdout, &op)?; + process_input(&mut locked_stdin, &mut output, &op)?; } } else if squeeze_flag { if sets_len == 1 { let op = SqueezeOperation::new(set1); - translate_input(&mut locked_stdin, &mut locked_stdout, op)?; + translate_input(&mut locked_stdin, &mut output, op)?; } else { let translate_op = TranslateOperation::new(set1, set2.clone())?; let squeeze_op = SqueezeOperation::new(set2); let op = translate_op.chain(squeeze_op); - translate_input(&mut locked_stdin, &mut locked_stdout, op)?; + translate_input(&mut locked_stdin, &mut output, op)?; } } else { let op = TranslateOperation::new(set1, set2)?; - process_input(&mut locked_stdin, &mut locked_stdout, &op)?; + process_input(&mut locked_stdin, &mut output, &op)?; } - flush_output(&mut locked_stdout)?; + flush_output(&mut output)?; Ok(()) } diff --git a/tests/by-util/test_tr.rs b/tests/by-util/test_tr.rs index 30009241d76..b2655de2e05 100644 --- a/tests/by-util/test_tr.rs +++ b/tests/by-util/test_tr.rs @@ -156,6 +156,87 @@ fn test_delete_complement_2() { .stdout_is("01234567890"); } +#[test] +fn test_delete_one_char_large_input() { + // Longer than one read, with no, few, many and only matches. `few` is + // sparse enough for the path that copies the runs between matches. + let none = vec![b'x'; 40_000]; + let few = [[b'x'; 999].as_slice(), b","].concat().repeat(40); + let many = b"a,".repeat(20_000); + let only = vec![b','; 40_000]; + for input in [none, few, many, only] { + let expected: Vec = input.iter().copied().filter(|&b| b != b',').collect(); + new_ucmd!() + .args(&["-d", ","]) + .pipe_in(input) + .succeeds() + .stdout_is_bytes(expected); + } +} + +#[test] +fn test_delete_set_large_input() { + let input: Vec = (0..=u8::MAX).cycle().take(40_000).collect(); + + let expected: Vec = input + .iter() + .copied() + .filter(|&b| !matches!(b, 0 | b'a'..=b'f' | u8::MAX)) + .collect(); + new_ucmd!() + .args(&["-d", "\\000a-f\\377"]) + .pipe_in(input.clone()) + .succeeds() + .stdout_is_bytes(expected); + + let expected: Vec = input + .iter() + .copied() + .filter(u8::is_ascii_lowercase) + .collect(); + new_ucmd!() + .args(&["-cd", "a-z"]) + .pipe_in(input) + .succeeds() + .stdout_is_bytes(expected); +} + +#[test] +fn test_delete_set_mostly_deleted() { + // Almost everything deleted, then half, then almost everything again. + let few_kept = [[b'a'; 99].as_slice(), b"\n"].concat().repeat(100); + let half_kept = b"a\n".repeat(5_000); + let input = [few_kept.as_slice(), &half_kept, &few_kept].concat(); + let expected: Vec = input + .iter() + .copied() + .filter(|b| !b.is_ascii_lowercase()) + .collect(); + new_ucmd!() + .args(&["-d", "a-z"]) + .pipe_in(input) + .succeeds() + .stdout_is_bytes(expected); +} + +#[test] +fn test_translate_one_char_large_input() { + // Longer than one read, with no and some matches. + let none = vec![b'x'; 40_000]; + let some = b"a,".repeat(20_000); + for input in [none, some] { + let expected: Vec = input + .iter() + .map(|&b| if b == b',' { b';' } else { b }) + .collect(); + new_ucmd!() + .args(&[",", ";"]) + .pipe_in(input) + .succeeds() + .stdout_is_bytes(expected); + } +} + #[test] fn test_complement1() { new_ucmd!() @@ -1688,9 +1769,11 @@ fn test_failed_write_is_reported() { #[test] #[cfg_attr(wasi_runner, ignore = "WASI: no pipe/signal support")] fn test_broken_pipe_no_error() { + // More than a pipe buffer, so that the output blocks until the reader is gone. new_ucmd!() .args(&["e", "a"]) - .pipe_in("hello".repeat(100)) + .pipe_in("hello".repeat(100_000)) + .ignore_stdin_write_error() .run_stdout_starts_with(b"") .fails_silently(); }