use anyhow::Result;
use memchr::{memchr_iter, memrchr};
#[cfg(test)]
use std::borrow::Cow;
use std::io::BufRead;
use std::ops::Range;
pub const BUF_READER_CAPACITY: usize = 1024 * 1024;
#[derive(Debug)]
pub struct LineBatch {
pub buf: Vec<u8>,
pub lines: Vec<Range<usize>>,
}
impl LineBatch {
fn from_chunk(chunk: &[u8]) -> Self {
let buf = chunk.to_vec();
let mut lines = Vec::new();
let mut start = 0;
for newline in memchr_iter(b'\n', &buf) {
lines.push(strip_cr(&buf, start..newline));
start = newline + 1;
}
lines.push(strip_cr(&buf, start..buf.len()));
Self { buf, lines }
}
pub fn single_line(line: &[u8]) -> Self {
let line = line.strip_suffix(b"\n").unwrap_or(line);
let line = line.strip_suffix(b"\r").unwrap_or(line);
Self {
buf: line.to_vec(),
lines: std::iter::once(0..line.len()).collect(),
}
}
#[cfg(test)]
pub fn iter(&self) -> impl Iterator<Item = Cow<'_, str>> {
self.lines
.iter()
.map(|range| String::from_utf8_lossy(&self.buf[range.clone()]))
}
}
fn strip_cr(buf: &[u8], range: Range<usize>) -> Range<usize> {
if buf[range.clone()].ends_with(b"\r") {
range.start..range.end - 1
} else {
range
}
}
#[derive(Debug)]
pub enum ReadResult {
Eof,
Batch(LineBatch),
}
pub fn read_batch<R>(reader: &mut R) -> Result<ReadResult>
where
R: BufRead,
{
let buf = reader.fill_buf()?;
if buf.is_empty() {
return Ok(ReadResult::Eof);
}
let Some(last) = memrchr(b'\n', buf) else {
let mut line = Vec::new();
reader.read_until(b'\n', &mut line)?;
return Ok(ReadResult::Batch(LineBatch::single_line(&line)));
};
let batch = LineBatch::from_chunk(&buf[..last]);
reader.consume(last + 1);
Ok(ReadResult::Batch(batch))
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{BufReader, Read};
use std::sync::mpsc;
use std::time::Duration;
fn texts(batch: &LineBatch) -> Vec<String> {
batch.iter().map(Cow::into_owned).collect()
}
#[test]
fn test_read_complete_lines_multiple_lines() {
let input = b"line 1\nline 2\nline 3\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["line 1", "line 2", "line 3"]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_do_not_omit_empty_lines() {
let input = b"line 1\n\nline 3\n\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["line 1", "", "line 3", ""]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_read_complete_lines_multiple_lines_no_final_newline() {
let input = b"line 1\nline 2\nline 3";
let mut reader = BufReader::new(&input[..]);
let mut lines = vec![];
let lines_first_read = read_batch(&mut reader).unwrap();
let lines_second_read = read_batch(&mut reader).unwrap();
if let ReadResult::Batch(batch) = lines_first_read {
lines.extend(texts(&batch));
}
if let ReadResult::Batch(batch) = lines_second_read {
lines.extend(texts(&batch));
}
assert_eq!(lines, vec!["line 1", "line 2", "line 3"]);
}
#[test]
fn test_read_eof_without_new_line_should_still_return_the_whole_string() {
let input = b"incomplete line without newline";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["incomplete line without newline"]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_read_complete_lines_single_complete_line() {
let input = b"single complete line\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["single complete line"]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_read_complete_lines_empty_input() {
let input = b"";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Eof) = read_batch(&mut reader) {
} else {
panic!("Expected Ok(ReadResult::Eof), got something else");
}
}
#[test]
fn test_read_complete_lines_partial_line_pending() {
let rx = read_batch_in_background(b"incomplete line without newline");
let result = rx.recv_timeout(Duration::from_millis(100));
assert!(result.is_err(), "Expected read_batch to block without newline or EOF");
}
#[test]
fn test_read_available_lines() {
let rx = read_batch_in_background(b"extract lines\n that are ready\n to be extracted\n but not more than that");
match rx.recv_timeout(Duration::from_secs(5)) {
Ok(Ok(ReadResult::Batch(batch))) => {
assert_eq!(
texts(&batch),
vec!["extract lines", " that are ready", " to be extracted"]
);
}
other => panic!("Expected Ok(ReadResult::Batch), got {other:?}"),
}
}
fn read_batch_in_background(input: &[u8]) -> mpsc::Receiver<Result<ReadResult>> {
let source = BlocksAfterInitialRead::new(input);
let (tx, rx) = mpsc::channel();
std::thread::spawn(move || {
let mut reader = BufReader::new(source);
let _ = tx.send(read_batch(&mut reader));
});
rx
}
struct BlocksAfterInitialRead {
data: Vec<u8>,
position: usize,
}
impl BlocksAfterInitialRead {
fn new(initial_data: &[u8]) -> Self {
Self {
data: initial_data.to_vec(),
position: 0,
}
}
}
impl Read for BlocksAfterInitialRead {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
if self.position < self.data.len() {
let remaining = &self.data[self.position..];
let to_read = remaining.len().min(buf.len());
buf[..to_read].copy_from_slice(&remaining[..to_read]);
self.position += to_read;
Ok(to_read)
} else {
loop {
std::thread::park();
}
}
}
}
#[test]
fn test_crlf_batch() {
let input = b"line1\r\nline2\r\nline3\r\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["line1", "line2", "line3"]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_crlf_single_line() {
let input = b"line1\r\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["line1"]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_non_utf8_single_line() {
let input = b"caf\xe9\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
let lines = texts(&batch);
assert_eq!(lines.len(), 1);
assert!(lines[0].starts_with("caf"));
assert!(lines[0].contains('\u{FFFD}'));
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
#[test]
fn test_mixed_line_endings() {
let input = b"unix\nwindows\r\nunix2\n";
let mut reader = BufReader::new(&input[..]);
if let Ok(ReadResult::Batch(batch)) = read_batch(&mut reader) {
assert_eq!(texts(&batch), vec!["unix", "windows", "unix2"]);
} else {
panic!("Expected Ok(ReadResult::Batch), got something else");
}
}
}