From 85896bdb86dbcfbfb38fff85f29bab2fac97b24a Mon Sep 17 00:00:00 2001
From: Paul Campbell <pcampbell@kemitix.net>
Date: Thu, 23 Mar 2023 14:46:02 +0000
Subject: [PATCH] User BufRead to avoid loading entire file into memory and
 also support stdin

---
 src/lib.rs  | 118 +++++++++++++++++++++++++++++-----------------------
 src/main.rs |   1 -
 2 files changed, 66 insertions(+), 53 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index f26afee..304b7a2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,7 +1,7 @@
 use clap::Parser;
 use std::fmt::Debug;
-use std::fs;
-use std::io::Write;
+use std::fs::File;
+use std::io::{BufRead, BufReader, Write};
 use std::path::PathBuf;
 
 pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
@@ -12,7 +12,7 @@ pub struct Cli {
     /// The number of lines to skip
     lines: usize,
     /// The file to read, or stdin if not given
-    file: Option<PathBuf>,
+    pub file: Option<PathBuf>,
     /// Skip until N lines matching this
     #[arg(short, long, conflicts_with = "token")]
     line: Option<String>,
@@ -30,64 +30,75 @@ pub struct Cli {
 }
 
 pub fn skip(cli: &Cli, writer: &mut impl Write) -> Result<()> {
-    match &cli.file {
+    let reader: Box<dyn BufRead> = match &cli.file {
         Some(ref file) => {
-            if let Some(line) = &cli.line {
-                skip_file_lines_matching(&cli, writer, file, line)
-            } else if let Some(ref token) = cli.token {
-                skip_file_tokens(cli, writer, file, token)
-            } else {
-                skip_file_lines(cli, writer, file)
+            let file = File::open(file)?;
+            Box::new(BufReader::new(file))
+        }
+        None => Box::new(BufReader::new(std::io::stdin())),
+    };
+    if let Some(line) = &cli.line {
+        skip_file_lines_matching(&cli, reader, writer, line)
+    } else if let Some(ref token) = cli.token {
+        skip_file_tokens(cli, reader, writer, token)
+    } else {
+        skip_file_lines(cli, reader, writer)
+    }
+}
+
+fn skip_file_lines(cli: &Cli, reader: Box<dyn BufRead>, writer: &mut impl Write) -> Result<()> {
+    let mut counter = 0usize;
+    for current_line in reader.lines() {
+        if let Ok(current_line) = current_line {
+            if counter >= cli.lines {
+                writeln!(writer, "{}", current_line)?;
             }
-        }
-        None => todo!("reading from stdin"),
-    }
-}
-
-fn skip_file_lines(cli: &Cli, writer: &mut impl Write, file: &PathBuf) -> Result<()> {
-    let content = fs::read_to_string(file).expect("Could not read file");
-    let mut counter = 0usize;
-    for current_line in content.lines() {
-        if counter >= cli.lines {
-            writeln!(writer, "{}", current_line)?;
-        }
-        counter += 1;
-    }
-    Ok(())
-}
-
-fn skip_file_lines_matching(
-    cli: &Cli,
-    writer: &mut impl Write,
-    file: &PathBuf,
-    line: &str,
-) -> Result<()> {
-    let content = fs::read_to_string(file).expect("Could not read file");
-    let mut counter = 0usize;
-    for current_line in content.lines() {
-        if counter >= cli.lines {
-            writeln!(writer, "{}", current_line)?;
-        }
-        if line == current_line {
             counter += 1;
         }
     }
     Ok(())
 }
 
-fn skip_file_tokens(cli: &Cli, writer: &mut impl Write, file: &PathBuf, token: &str) -> Result<()> {
-    let content = fs::read_to_string(file).expect("Could not read file");
+fn skip_file_lines_matching(
+    cli: &Cli,
+    reader: Box<dyn BufRead>,
+    writer: &mut impl Write,
+    line: &str,
+) -> Result<()> {
     let mut counter = 0usize;
-    for current_line in content.lines() {
-        if counter >= cli.lines {
-            writeln!(writer, "{}", current_line)?;
-        }
-        if current_line.contains(&token) {
-            if cli.ignore_extras {
+    for current_line in reader.lines() {
+        if let Ok(current_line) = current_line {
+            if counter >= cli.lines {
+                writeln!(writer, "{}", current_line)?;
+            }
+            if line == current_line {
                 counter += 1;
-            } else {
-                let occurances = current_line.matches(&token).count();
-                counter += occurances;
+            }
+        }
+    }
+    Ok(())
+}
+
+fn skip_file_tokens(
+    cli: &Cli,
+    reader: Box<dyn BufRead>,
+    writer: &mut impl Write,
+    token: &str,
+) -> Result<()> {
+    let mut counter = 0usize;
+
+    for current_line in reader.lines() {
+        if let Ok(current_line) = current_line {
+            if counter >= cli.lines {
+                writeln!(writer, "{}", current_line)?;
+            }
+            if current_line.contains(&token) {
+                if cli.ignore_extras {
+                    counter += 1;
+                } else {
+                    let occurances = current_line.matches(&token).count();
+                    counter += occurances;
+                }
             }
         }
     }
@@ -135,7 +146,10 @@ mod tests {
         skip(&cli, &mut lines)?;
 
         //then
-        assert_eq!(String::from_utf8(lines)?, vec!["alpha", "gamma\n"].join("\n"));
+        assert_eq!(
+            String::from_utf8(lines)?,
+            vec!["alpha", "gamma\n"].join("\n")
+        );
         Ok(())
     }
 
diff --git a/src/main.rs b/src/main.rs
index 2df5010..4eb8e7c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,5 @@ use skip::{skip, Cli, Result};
 fn main() -> Result<()> {
     let cli = Cli::parse();
     let mut stdout = std::io::stdout();
-
     skip(&cli, &mut stdout)
 }