head: read directly from input buffer in print_n_lines

Author: rossilor95

src/uu/head/src/head.rs

@@ -6,7 +6,7 @@
 // spell-checker:ignore (vars) memrchr
 
 use clap::ArgMatches;
-use memchr::memrchr_iter;
+use memchr::{memchr_iter, memrchr_iter};
 use std::ffi::OsString;
 use std::fs::File;
 use std::io::{self, BufWriter, Read, Seek, SeekFrom, Write};
@@ -31,7 +31,6 @@ mod parse;
 mod take;
 use take::copy_all_but_n_bytes;
 use take::copy_all_but_n_lines;
-use take::take_lines;
 
 #[derive(Error, Debug)]
 enum HeadError {
@@ -203,29 +202,37 @@ fn print_n_lines(
     n: u64,
     separator: u8,
 ) -> Result<u64, HeadFileError> {
-    // Read the first `n` lines from the `input` reader.
-    let mut reader = take_lines(input, n, separator);
-
-    // Write those bytes to `stdout`.
     let stdout = io::stdout();
     let stdout = stdout.lock();
     let mut writer = BufWriter::with_capacity(BUF_SIZE, stdout);
 
-    let mut bytes_written = 0;
-    let mut buf = [0; BUF_SIZE];
-    loop {
-        let bytes_read = reader.read(&mut buf).map_err(HeadFileError::Read)?;
+    let mut bytes_written: u64 = 0;
+    let mut remaining = n;
+    while remaining > 0 {
+        let chunk = input.fill_buf().map_err(HeadFileError::Read)?;
 
-        if bytes_read == 0 {
+        if chunk.is_empty() {
             break;
         }
 
+        let mut take_len = chunk.len(); // default: take everything
+        let mut separators_seen: u64 = 0;
+        for separator_idx in memchr_iter(separator, chunk) {
+            separators_seen += 1;
+            if separators_seen == remaining {
+                take_len = separator_idx + 1; // include the separator itself
+                break;
+            }
+        }
+        remaining -= separators_seen;
+
         writer
-            .write_all(&buf[..bytes_read])
+            .write_all(&chunk[..take_len])
             .map_err(wrap_in_stdout_error)
             .map_err(HeadFileError::WriteStdout)?;
 
-        bytes_written += bytes_read as u64;
+        input.consume(take_len);
+        bytes_written += take_len as u64;
     }
 
     // Make sure we finish writing everything to the target before

src/uu/head/src/take.rs

@@ -303,62 +303,11 @@ pub fn copy_all_but_n_lines<R: Read, W: Write>(
     Ok(total_bytes_copied)
 }
 
-/// Like `std::io::Take`, but for lines instead of bytes.
-///
-/// This struct is generally created by calling [`take_lines`] on a
-/// reader. Please see the documentation of [`take_lines`] for more
-/// details.
-pub struct TakeLines<T> {
-    inner: T,
-    limit: u64,
-    separator: u8,
-}
-
-impl<T: Read> Read for TakeLines<T> {
-    /// Read bytes from a buffer up to the requested number of lines.
-    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-        if self.limit == 0 {
-            return Ok(0);
-        }
-        match self.inner.read(buf) {
-            Ok(0) => Ok(0),
-            Ok(n) => {
-                for i in memchr_iter(self.separator, &buf[..n]) {
-                    self.limit -= 1;
-                    if self.limit == 0 {
-                        return Ok(i + 1);
-                    }
-                }
-                Ok(n)
-            }
-            Err(e) => Err(e),
-        }
-    }
-}
-
-/// Create an adaptor that will read at most `limit` lines from a given reader.
-///
-/// This function returns a new instance of `Read` that will read at
-/// most `limit` lines, after which it will always return EOF
-/// (`Ok(0)`).
-///
-/// The `separator` defines the character to interpret as the line
-/// ending. For the usual notion of "line", set this to `b'\n'`.
-pub fn take_lines<R>(reader: R, limit: u64, separator: u8) -> TakeLines<R> {
-    TakeLines {
-        inner: reader,
-        limit,
-        separator,
-    }
-}
-
 #[cfg(test)]
 mod tests {
 
-    use std::io::{BufRead, BufReader};
-
     use crate::take::{
-        TakeAllBuffer, TakeAllLinesBuffer, copy_all_but_n_bytes, copy_all_but_n_lines, take_lines,
+        TakeAllBuffer, TakeAllLinesBuffer, copy_all_but_n_bytes, copy_all_but_n_lines,
     };
 
     #[test]
@@ -635,33 +584,4 @@ mod tests {
         assert_eq!(bytes_copied, 2);
         assert_eq!(output_reader.get_ref()[..], input_buffer.as_bytes()[0..2]);
     }
-
-    #[test]
-    fn test_zero_lines() {
-        let input_reader = std::io::Cursor::new("a\nb\nc\n");
-        let output_reader = BufReader::new(take_lines(input_reader, 0, b'\n'));
-        let mut iter = output_reader.lines().map(|l| l.unwrap());
-        assert_eq!(None, iter.next());
-    }
-
-    #[test]
-    fn test_fewer_lines() {
-        let input_reader = std::io::Cursor::new("a\nb\nc\n");
-        let output_reader = BufReader::new(take_lines(input_reader, 2, b'\n'));
-        let mut iter = output_reader.lines().map(|l| l.unwrap());
-        assert_eq!(Some(String::from("a")), iter.next());
-        assert_eq!(Some(String::from("b")), iter.next());
-        assert_eq!(None, iter.next());
-    }
-
-    #[test]
-    fn test_more_lines() {
-        let input_reader = std::io::Cursor::new("a\nb\nc\n");
-        let output_reader = BufReader::new(take_lines(input_reader, 4, b'\n'));
-        let mut iter = output_reader.lines().map(|l| l.unwrap());
-        assert_eq!(Some(String::from("a")), iter.next());
-        assert_eq!(Some(String::from("b")), iter.next());
-        assert_eq!(Some(String::from("c")), iter.next());
-        assert_eq!(None, iter.next());
-    }
 }