Merge pull request #8587 from ThePseudo/input_reuse_b64

Input reuse on base64
This commit is contained in:
Sylvestre Ledru
2025-09-11 20:17:02 +02:00
committed by GitHub
+62 -229
View File
@@ -8,7 +8,7 @@
use clap::{Arg, ArgAction, Command};
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, ErrorKind, Read, Seek, SeekFrom};
use std::io::{self, ErrorKind, Read, Seek};
use std::path::{Path, PathBuf};
use uucore::display::Quotable;
use uucore::encoding::{
@@ -172,7 +172,7 @@ pub fn get_input(config: &Config) -> UResult<Box<dyn ReadSeek>> {
}
/// Determines if the input buffer ends with padding ('=') after trimming trailing whitespace.
fn has_padding<R: Read + Seek>(input: &mut R) -> UResult<bool> {
fn read_and_has_padding<R: Read>(input: &mut R) -> UResult<(bool, Vec<u8>)> {
let mut buf = Vec::new();
input
.read_to_end(&mut buf)
@@ -184,30 +184,27 @@ fn has_padding<R: Read + Seek>(input: &mut R) -> UResult<bool> {
.rfind(|&&byte| !byte.is_ascii_whitespace())
.is_some_and(|&byte| byte == b'=');
input.seek(SeekFrom::Start(0))?;
Ok(has_padding)
Ok((has_padding, buf))
}
pub fn handle_input<R: Read + Seek>(input: &mut R, format: Format, config: Config) -> UResult<()> {
let has_padding = has_padding(input)?;
let (has_padding, read) = read_and_has_padding(input)?;
let supports_fast_decode_and_encode =
get_supports_fast_decode_and_encode(format, config.decode, has_padding);
let supports_fast_decode_and_encode_ref = supports_fast_decode_and_encode.as_ref();
let mut stdout_lock = io::stdout().lock();
if config.decode {
fast_decode::fast_decode(
input,
read,
&mut stdout_lock,
supports_fast_decode_and_encode_ref,
config.ignore_garbage,
)
} else {
fast_encode::fast_encode(
input,
read,
&mut stdout_lock,
supports_fast_decode_and_encode_ref,
config.wrap_cols,
@@ -292,16 +289,14 @@ pub fn get_supports_fast_decode_and_encode(
}
pub mod fast_encode {
use crate::base_common::{WRAP_DEFAULT, format_read_error};
use crate::base_common::WRAP_DEFAULT;
use std::{
cmp::min,
collections::VecDeque,
io::{self, ErrorKind, Read, Write},
io::{self, Write},
num::NonZeroUsize,
};
use uucore::{
encoding::SupportsFastDecodeAndEncode,
error::{UResult, USimpleError},
};
use uucore::{encoding::SupportsFastDecodeAndEncode, error::UResult};
struct LineWrapping {
line_length: NonZeroUsize,
@@ -311,46 +306,10 @@ pub mod fast_encode {
// Start of helper functions
fn encode_in_chunks_to_buffer(
supports_fast_decode_and_encode: &dyn SupportsFastDecodeAndEncode,
encode_in_chunks_of_size: usize,
bytes_to_steal: usize,
read_buffer: &[u8],
encoded_buffer: &mut VecDeque<u8>,
leftover_buffer: &mut VecDeque<u8>,
) -> UResult<()> {
let bytes_to_chunk = if bytes_to_steal > 0 {
let (stolen_bytes, rest_of_read_buffer) = read_buffer.split_at(bytes_to_steal);
leftover_buffer.extend(stolen_bytes);
// After appending the stolen bytes to `leftover_buffer`, it should be the right size
assert_eq!(leftover_buffer.len(), encode_in_chunks_of_size);
// Encode the old unencoded data and the stolen bytes, and add the result to
// `encoded_buffer`
supports_fast_decode_and_encode
.encode_to_vec_deque(leftover_buffer.make_contiguous(), encoded_buffer)?;
// Reset `leftover_buffer`
leftover_buffer.clear();
rest_of_read_buffer
} else {
// Do not need to steal bytes from `read_buffer`
read_buffer
};
let chunks_exact = bytes_to_chunk.chunks_exact(encode_in_chunks_of_size);
let remainder = chunks_exact.remainder();
for sl in chunks_exact {
assert_eq!(sl.len(), encode_in_chunks_of_size);
supports_fast_decode_and_encode.encode_to_vec_deque(sl, encoded_buffer)?;
}
leftover_buffer.extend(remainder);
supports_fast_decode_and_encode.encode_to_vec_deque(read_buffer, encoded_buffer)?;
Ok(())
}
@@ -440,13 +399,12 @@ pub mod fast_encode {
// End of helper functions
pub fn fast_encode(
input: &mut dyn Read,
input: Vec<u8>,
output: &mut dyn Write,
supports_fast_decode_and_encode: &dyn SupportsFastDecodeAndEncode,
wrap: Option<usize>,
) -> UResult<()> {
// Based on performance testing
const INPUT_BUFFER_SIZE: usize = 32 * 1_024;
const ENCODE_IN_CHUNKS_OF_SIZE_MULTIPLE: usize = 1_024;
@@ -473,12 +431,9 @@ pub mod fast_encode {
}),
};
let input_size = input.len();
// Start of buffers
// Data that was read from `input`
let mut input_buffer = vec![0; INPUT_BUFFER_SIZE];
assert!(!input_buffer.is_empty());
// Data that was read from `input` but has not been encoded yet
let mut leftover_buffer = VecDeque::<u8>::new();
@@ -486,39 +441,33 @@ pub mod fast_encode {
let mut encoded_buffer = VecDeque::<u8>::new();
// End of buffers
loop {
match input.read(&mut input_buffer) {
Ok(bytes_read_from_input) => {
if bytes_read_from_input == 0 {
break;
}
// The part of `input_buffer` that was actually filled by the call to `read`
let read_buffer = &input_buffer[..bytes_read_from_input];
// How many bytes to steal from `read_buffer` to get `leftover_buffer` to the right size
let bytes_to_steal = encode_in_chunks_of_size - leftover_buffer.len();
if bytes_to_steal > bytes_read_from_input {
// Do not have enough data to encode a chunk, so copy data to `leftover_buffer` and read more
leftover_buffer.extend(read_buffer);
assert!(leftover_buffer.len() < encode_in_chunks_of_size);
continue;
}
input
.iter()
.enumerate()
.step_by(encode_in_chunks_of_size)
.map(|(idx, _)| {
// The part of `input_buffer` that was actually filled by the call
// to `read`
&input[idx..min(input_size, idx + encode_in_chunks_of_size)]
})
.map(|buffer| {
if buffer.len() < encode_in_chunks_of_size {
leftover_buffer.extend(buffer);
assert!(leftover_buffer.len() < encode_in_chunks_of_size);
return None;
}
Some(buffer)
})
.for_each(|buffer| {
if let Some(read_buffer) = buffer {
// Encode data in chunks, then place it in `encoded_buffer`
assert_eq!(read_buffer.len(), encode_in_chunks_of_size);
encode_in_chunks_to_buffer(
supports_fast_decode_and_encode,
encode_in_chunks_of_size,
bytes_to_steal,
read_buffer,
&mut encoded_buffer,
&mut leftover_buffer,
)?;
assert!(leftover_buffer.len() < encode_in_chunks_of_size);
)
.unwrap();
// Write all data in `encoded_buffer` to `output`
write_to_output(
&mut line_wrapping,
@@ -526,20 +475,10 @@ pub mod fast_encode {
output,
false,
wrap == Some(0),
)?;
)
.unwrap();
}
Err(er) => {
let kind = er.kind();
if kind == ErrorKind::Interrupted {
// Retry reading
continue;
}
return Err(USimpleError::new(1, format_read_error(kind)));
}
}
}
});
// Cleanup
// `input` has finished producing data, so the data remaining in the buffers needs to be encoded and printed
@@ -558,18 +497,13 @@ pub mod fast_encode {
wrap == Some(0),
)?;
}
Ok(())
}
}
pub mod fast_decode {
use crate::base_common::format_read_error;
use std::io::{self, ErrorKind, Read, Write};
use uucore::{
encoding::SupportsFastDecodeAndEncode,
error::{UResult, USimpleError},
};
use std::io::{self, Write};
use uucore::{encoding::SupportsFastDecodeAndEncode, error::UResult};
// Start of helper functions
fn alphabet_to_table(alphabet: &[u8], ignore_garbage: bool) -> [bool; 256] {
@@ -610,46 +544,10 @@ pub mod fast_decode {
fn decode_in_chunks_to_buffer(
supports_fast_decode_and_encode: &dyn SupportsFastDecodeAndEncode,
decode_in_chunks_of_size: usize,
bytes_to_steal: usize,
read_buffer_filtered: &[u8],
decoded_buffer: &mut Vec<u8>,
leftover_buffer: &mut Vec<u8>,
) -> UResult<()> {
let bytes_to_chunk = if bytes_to_steal > 0 {
let (stolen_bytes, rest_of_read_buffer_filtered) =
read_buffer_filtered.split_at(bytes_to_steal);
leftover_buffer.extend(stolen_bytes);
// After appending the stolen bytes to `leftover_buffer`, it should be the right size
assert_eq!(leftover_buffer.len(), decode_in_chunks_of_size);
// Decode the old un-decoded data and the stolen bytes, and add the result to
// `decoded_buffer`
supports_fast_decode_and_encode.decode_into_vec(leftover_buffer, decoded_buffer)?;
// Reset `leftover_buffer`
leftover_buffer.clear();
rest_of_read_buffer_filtered
} else {
// Do not need to steal bytes from `read_buffer`
read_buffer_filtered
};
let chunks_exact = bytes_to_chunk.chunks_exact(decode_in_chunks_of_size);
let remainder = chunks_exact.remainder();
for sl in chunks_exact {
assert_eq!(sl.len(), decode_in_chunks_of_size);
supports_fast_decode_and_encode.decode_into_vec(sl, decoded_buffer)?;
}
leftover_buffer.extend(remainder);
supports_fast_decode_and_encode.decode_into_vec(read_buffer_filtered, decoded_buffer)?;
Ok(())
}
@@ -664,14 +562,11 @@ pub mod fast_decode {
// End of helper functions
pub fn fast_decode(
input: &mut dyn Read,
input: Vec<u8>,
output: &mut dyn Write,
supports_fast_decode_and_encode: &dyn SupportsFastDecodeAndEncode,
ignore_garbage: bool,
) -> UResult<()> {
// Based on performance testing
const INPUT_BUFFER_SIZE: usize = 32 * 1_024;
const DECODE_IN_CHUNKS_OF_SIZE_MULTIPLE: usize = 1_024;
let alphabet = supports_fast_decode_and_encode.alphabet();
@@ -693,102 +588,40 @@ pub mod fast_decode {
let table = alphabet_to_table(alphabet, ignore_garbage);
// Start of buffers
// Data that was read from `input`
let mut input_buffer = vec![0; INPUT_BUFFER_SIZE];
assert!(!input_buffer.is_empty());
// Data that was read from `input` but has not been decoded yet
let mut leftover_buffer = Vec::<u8>::new();
// Decoded data that needs to be written to `output`
let mut decoded_buffer = Vec::<u8>::new();
// Buffer that will be used when `ignore_garbage` is true, and the chunk read from `input` contains garbage
// data
let mut non_garbage_buffer = Vec::<u8>::new();
// End of buffers
loop {
match input.read(&mut input_buffer) {
Ok(bytes_read_from_input) => {
if bytes_read_from_input == 0 {
break;
}
let read_buffer_filtered = {
// The part of `input_buffer` that was actually filled by the call to `read`
let read_buffer = &input_buffer[..bytes_read_from_input];
// First just scan the data for the happy path
// Yields significant speedup when the input does not contain line endings
let found_garbage = read_buffer.iter().any(|ue| {
// Garbage, since it was not found in the table
!table[usize::from(*ue)]
});
if found_garbage {
non_garbage_buffer.clear();
for ue in read_buffer {
if table[usize::from(*ue)] {
// Not garbage, since it was found in the table
non_garbage_buffer.push(*ue);
}
}
non_garbage_buffer.as_slice()
} else {
read_buffer
}
};
// How many bytes to steal from `read_buffer` to get `leftover_buffer` to the right size
let bytes_to_steal = decode_in_chunks_of_size - leftover_buffer.len();
if bytes_to_steal > read_buffer_filtered.len() {
// Do not have enough data to decode a chunk, so copy data to `leftover_buffer` and read more
leftover_buffer.extend(read_buffer_filtered);
assert!(leftover_buffer.len() < decode_in_chunks_of_size);
continue;
}
let mut buffer = Vec::with_capacity(decode_in_chunks_of_size);
input
.iter()
.filter(|ch| table[usize::from(**ch)])
.for_each(|ch| {
buffer.push(*ch);
// How many bytes to steal from `read_buffer` to get
// `leftover_buffer` to the right size
if buffer.len() == decode_in_chunks_of_size {
assert_eq!(decode_in_chunks_of_size, buffer.len());
// Decode data in chunks, then place it in `decoded_buffer`
decode_in_chunks_to_buffer(
supports_fast_decode_and_encode,
decode_in_chunks_of_size,
bytes_to_steal,
read_buffer_filtered,
&buffer,
&mut decoded_buffer,
&mut leftover_buffer,
)?;
assert!(leftover_buffer.len() < decode_in_chunks_of_size);
)
.unwrap();
// Write all data in `decoded_buffer` to `output`
write_to_output(&mut decoded_buffer, output)?;
write_to_output(&mut decoded_buffer, output).unwrap();
buffer.clear();
}
Err(er) => {
let kind = er.kind();
if kind == ErrorKind::Interrupted {
// Retry reading
continue;
}
return Err(USimpleError::new(1, format_read_error(kind)));
}
}
}
});
// Cleanup
// `input` has finished producing data, so the data remaining in the buffers needs to be decoded and printed
{
// Decode all remaining encoded bytes, placing them in `decoded_buffer`
supports_fast_decode_and_encode
.decode_into_vec(&leftover_buffer, &mut decoded_buffer)?;
supports_fast_decode_and_encode.decode_into_vec(&buffer, &mut decoded_buffer)?;
// Write all data in `decoded_buffer` to `output`
write_to_output(&mut decoded_buffer, output)?;
@@ -819,7 +652,7 @@ fn format_read_error(kind: ErrorKind) -> String {
#[cfg(test)]
mod tests {
use super::*;
use crate::base_common::read_and_has_padding;
use std::io::Cursor;
#[test]
@@ -838,7 +671,7 @@ mod tests {
for (input, expected) in test_cases {
let mut cursor = Cursor::new(input.as_bytes());
assert_eq!(
has_padding(&mut cursor).unwrap(),
read_and_has_padding(&mut cursor).unwrap().0,
expected,
"Failed for input: '{input}'"
);