Avoid allocating on every input line.

By using `BufRead::read_line` instead of `BufRead::lines`. Now
allocations are only required (a) for the first occurrence of a
particular line, and (b) for lines that are modified by `-e`.

This requires some changes to handle cases where the line is a `&str` vs
a `String` (e.g. after modification due to `-e`).

This roughly doubles the speed of `counts` on files where every line is
the same, and leaves it unchanged on files where every line is
different. Typical cases will be between those two extremes.

This idea came from @Shnatsel via
nnethercote/perf-book#68.

src/main.rs

@@ -85,23 +85,23 @@ fn do_main() -> io::Result<()> {
 
  let erased_label = if erase { ", erased" } else { "" };
  match weights {
- Unit => process(readers, "", |line| (line, 1i64)),
+ Unit => process(readers, "", |_line| (None, 1i64)),
  Integral => {
  let re = Regex::new(r"(([+-]?)\d+)(\D*)$").unwrap();
  process(
  readers,
  &format!(" (weighted integral{})", erased_label),
  |line| {
- if let Some(captures) = re.captures(&line) {
+ if let Some(captures) = re.captures(line) {
  let weight = i64::from_str(&captures[1]).unwrap();
- let line = if erase {
- re.replace(&line, "NNN${3}").to_string()
+ if erase {
+ let line = re.replace(line, "NNN${3}").to_string();
+ (Some(line), weight)
  } else {
- line
- };
- (line, weight)
+ (None, weight)
+ }
  } else {
- (line, 1i64)
+ (None, 1i64)
  }
  },
  )
@@ -112,16 +112,16 @@ fn do_main() -> io::Result<()> {
  readers,
  &format!(" (weighted fractional{})", erased_label),
  |line| {
- if let Some(captures) = re.captures(&line) {
+ if let Some(captures) = re.captures(line) {
  let weight = f64::from_str(&captures[1]).unwrap();
- let line = if erase {
- re.replace(&line, "NNN${4}").to_string()
+ if erase {
+ let line = re.replace(line, "NNN${4}").to_string();
+ (Some(line), weight)
  } else {
- line
- };
- (line, weight)
+ (None, weight)
+ }
  } else {
- (line, 1f64)
+ (None, 1f64)
  }
  },
  )
@@ -137,22 +137,54 @@ fn process<F, N>(
  get_line_and_weight: F,
 ) -> io::Result<()>
 where
- F: Fn(String) -> (String, N),
+ F: Fn(&str) -> (Option<String>, N),
  N: Total,
 {
- let mut counts: FxHashMap<String, N> = FxHashMap::default();
  let mut total = N::from(0u32);
+ let mut counts: FxHashMap<String, N> = FxHashMap::default();
 
- for reader in readers {
- for line in reader.lines() {
- let Ok(line) = line else {
- eprintln!("counts: non-UTF-8 input detected, aborting");
- std::process::exit(1);
- };
- let (line, weight) = get_line_and_weight(line);
- let entry = counts.entry(line).or_insert_with(|| N::from(0u32));
- *entry += weight;
+ // `reader.lines()` is the obvious way to do this loop, but that requires
+ // allocating every line into a `String`. Instead we use
+ // `reader.read_line()` and use a single string for every iteration. On the
+ // first occurrence of a line we need to do a `to_string` to insert it into
+ // the table. On subsequent occurrences we don't. Most `counts` input tends
+ // to have significant numbers of repeated lines, so this approach reduces
+ // allocation counts greatly.
+ let mut line_with_nl = String::new();
+ for mut reader in readers {
+ loop {
+ match reader.read_line(&mut line_with_nl) {
+ Ok(0) => break,
+ Ok(_) => {}
+ Err(err) => {
+ eprintln!("counts: {}", err);
+ std::process::exit(1);
+ }
+ }
+
+ let line = &line_with_nl[..line_with_nl.len() - 1];
+ let (modified_line, weight) = get_line_and_weight(line);
+ match modified_line {
+ None => {
+ // The line has not been modified. Only promote to `String`
+ // if it hasn't been seen before.
+ if let Some(entry) = counts.get_mut(line) {
+ *entry += weight;
+ } else {
+ counts.insert(line.to_string(), weight);
+ }
+ }
+ Some(modified_line) => {
+ // The line has been modified, which means it has already
+ // been promoted to a `String`.
+ let entry = counts.entry(modified_line).or_insert_with(|| N::from(0u32));
+ *entry += weight;
+ }
+ }
  total += weight;
+
+ // We are finished with the contents of this line.
+ line_with_nl.clear();
  }
  }
 

tests/cli.rs

@@ -228,7 +228,7 @@ baz 23 - +1
 #[test]
 fn non_utf8() -> Result<(), Box<dyn std::error::Error>> {
  let input = unsafe { std::str::from_utf8_unchecked(&[0x97, 0x98, 0x99, 0xff]) };
- let expected_output = "counts: non-UTF-8 input detected, aborting\n";
+ let expected_output = "counts: stream did not contain valid UTF-8\n";
 
  bad_test(input, expected_output)
 }