perf: speedup fasta reader

Here I restructured `FastaReader::read()` function to avoid extra allocations. This is basically a complete rewrite.

This involved changing error handling, so I thought while we are at it, why not improve error messages.

While working on the new algo, I also discovered a few important cases missing from tests, so I added them.

Speedup on mpox-500 run: ~29%

Command:
```
$ cargo -q build --release --target-dir=/workdir/.build/docker --bin=treetime
$ hyperfine --warmup 1 --show-output '/workdir/.build/docker/release/treetime ancestral --method-anc=parsimony --dense=false --tree=data/mpox/clade-ii/500/tree.nwk --outdir=tmp/smoke-tests/ancestral/marginal/mpox/clade-ii/500 data/mpox/clade-ii/500/aln.fasta.xz'
```

Before (branch: rust, commit f676392):
```
  Time (mean ± σ):      1.866 s ±  0.113 s    [User: 4.547 s, System: 0.576 s]
  Range (min … max):    1.768 s …  2.041 s    10 runs
```

After:
```
  Time (mean ± σ):      1.444 s ±  0.111 s    [User: 4.181 s, System: 0.527 s]
  Range (min … max):    1.341 s …  1.650 s    10 runs
```

Profiler shows that `treetime::io::fasta::FastaReader::read()` now takes ~18% of the program, down from 25% (in single-threaded mode).

Author: ivan-aksamentov

packages/treetime/src/io/fasta.rs

@@ -3,8 +3,10 @@ use crate::io::compression::Decompressor;
 use crate::io::concat::Concat;
 use crate::io::file::{create_file_or_stdout, open_file_or_stdin, open_stdin};
 use crate::make_error;
-use eyre::Report;
-use log::info;
+use crate::utils::string::quote_single;
+use eyre::{Context, Report};
+use itertools::Itertools;
+use log::{info, warn};
 use serde::{Deserialize, Serialize};
 use std::io::{BufRead, BufReader};
 use std::path::Path;
@@ -33,12 +35,21 @@ impl FastaRecord {
   pub fn is_empty(&self) -> bool {
     self.seq_name.is_empty() && self.seq_name.is_empty() && self.desc.is_none() && self.index == 0
   }
+
+  pub fn header(&self) -> String {
+    match &self.desc {
+      Some(desc) => format!(">{} {}", self.seq_name, desc),
+      None => format!(">{}", self.seq_name),
+    }
+  }
 }
 
 pub struct FastaReader<'a, 'b> {
   reader: Box<dyn BufRead + 'a>,
   alphabet: &'b Alphabet,
   line: String,
+  n_lines: usize,
+  n_chars: usize,
   index: usize,
 }
 
@@ -48,6 +59,8 @@ impl<'a, 'b> FastaReader<'a, 'b> {
       reader,
       alphabet,
       line: String::new(),
+      n_lines: 0,
+      n_chars: 0,
       index: 0,
     }
   }
@@ -94,100 +107,71 @@ impl<'a, 'b> FastaReader<'a, 'b> {
     record.clear();
 
     if self.line.is_empty() {
+      // Read lines until we find the next record or EOF
       loop {
         self.line.clear();
+        if self.reader.read_line(&mut self.line)? == 0 {
+          if self.index > 0 {
+            // We have read at least one record  by the end of input - this is normal operation mode
+            return Ok(());
+          }
+
+          if self.index == 0 && self.n_chars == 0 {
+            // We have read no records and no non-whitespace characters by the end of input
+            warn!("FASTA input is empty or consists entirely from whitespace: this is allowed but might not be what's intended");
+            return Ok(());
+          }
 
-        let n_bytes = self.reader.read_line(&mut self.line)?;
-        if n_bytes == 0 {
-          return Ok(());
+          // We have read some characters, but no records detected by the end of input
+          return make_error!(
+            "FASTA input is incorrectly formatted: expected at least one FASTA record starting with character '>', but none found"
+          );
         }
 
         let trimmed = self.line.trim();
-        if !trimmed.is_empty() {
-          self.line = trimmed.to_owned();
-          break;
+        self.n_lines += 1;
+        self.n_chars += trimmed.len();
+        if trimmed.starts_with('>') {
+          break; // Found the header of the next record
         }
       }
     }
 
-    if !self.line.starts_with('>') {
-      return make_error!("Expected character '>' at record start.");
-    }
-
-    let s = self.line[1..].trim();
-
-    let mut parts = s.splitn(2, ' ');
-    let name = parts.next().unwrap_or_default().to_owned();
-    let desc = parts.next().map(ToOwned::to_owned);
-
-    record.seq_name = name;
-    record.desc = desc;
-
-    loop {
-      self.line.clear();
-
-      let n_bytes = self.reader.read_line(&mut self.line)?;
-      if n_bytes == 0 {
-        record.index = self.index;
-        self.index += 1;
-        return Ok(());
-      }
+    let header_line = self.line.trim();
+    let (name, desc) = header_line[1..].split_once(' ').unwrap_or((&header_line[1..], ""));
+    record.seq_name = name.to_owned();
+    record.desc = if desc.is_empty() { None } else { Some(desc.to_owned()) };
+    record.index = self.index;
+    self.index += 1;
 
+    // Read sequence lines until the next header or EOF
+    self.line.clear();
+    while self.reader.read_line(&mut self.line)? > 0 {
       let trimmed = self.line.trim();
-      if !trimmed.is_empty() {
-        self.line = trimmed.to_owned();
+      self.n_lines += 1;
+      self.n_chars += trimmed.len();
+      if trimmed.starts_with('>') {
+        // We have reached the next record
         break;
       }
-    }
-
-    if self.line.is_empty() || self.line.starts_with('>') {
-      record.index = self.index;
-      self.index += 1;
-      return Ok(());
-    }
 
-    let fragment = self
-      .line
-      .chars()
-      .map(|c| {
-        let c = c.to_ascii_uppercase();
-        if !self.alphabet.contains(c) {
-          make_error!("Character is not in the alphabet: '{c}'")
+      record.seq.reserve(trimmed.len());
+      for c in trimmed.chars() {
+        let uc = c.to_ascii_uppercase();
+        if self.alphabet.contains(uc) {
+          record.seq.push(uc);
         } else {
-          Ok(c)
+          return make_error!(
+            "FASTA input is incorrect: character \"{c}\" is not in the alphabet. Expected characters: {}",
+            self.alphabet.chars().map(quote_single).join(", ")
+          )
+          .wrap_err_with(|| format!("When processing sequence #{}: \"{}\"", self.index, record.header()));
         }
-      })
-      .collect::<Result<Vec<char>, Report>>()?;
-
-    record.seq.extend(&fragment);
-
-    loop {
-      self.line.clear();
-      self.reader.read_line(&mut self.line)?;
-      self.line = self.line.trim().to_owned();
-      if self.line.is_empty() || self.line.starts_with('>') {
-        break;
       }
 
-      let fragment = self
-        .line
-        .chars()
-        .map(|c| {
-          let c = c.to_ascii_uppercase();
-          if !self.alphabet.contains(c) {
-            make_error!("Character is not in the alphabet: '{c}'")
-          } else {
-            Ok(c)
-          }
-        })
-        .collect::<Result<Vec<char>, Report>>()?;
-
-      record.seq.extend(&fragment);
+      self.line.clear();
     }
 
-    record.index = self.index;
-    self.index += 1;
-
     Ok(())
   }
 }
@@ -294,6 +278,7 @@ mod tests {
   use super::*;
   use crate::alphabet::alphabet::AlphabetName;
   use crate::o;
+  use crate::utils::error::report_to_string;
   use indoc::indoc;
   use lazy_static::lazy_static;
   use pretty_assertions::assert_eq;
@@ -312,10 +297,20 @@ mod tests {
     let mut record = FastaRecord::new();
     assert_eq!(
       reader.read(&mut record).unwrap_err().to_string(),
-      "Expected character '>' at record start."
+      "FASTA input is incorrectly formatted: expected at least one FASTA record starting with character '>', but none found"
     );
   }
 
+  #[test]
+  fn test_fasta_reader_fail_on_unknown_char() {
+    let data = b">seq%1\nACGT%ACGT\n";
+    let mut reader = FastaReader::new(Box::new(Cursor::new(data)), &NUC_ALPHABET);
+    let mut record = FastaRecord::new();
+    let actual = report_to_string(&reader.read(&mut record).unwrap_err());
+    let expected = r#"When processing sequence #1: ">seq%1": FASTA input is incorrect: character "%" is not in the alphabet. Expected characters: '-', 'A', 'B', 'C', 'D', 'G', 'H', 'K', 'M', 'N', 'R', 'S', 'T', 'V', 'W', 'Y'"#;
+    assert_eq!(expected, actual);
+  }
+
   #[test]
   fn test_fasta_reader_read_empty() {
     let data = b"";
@@ -568,6 +563,8 @@ mod tests {
       ACGT
       >Identifier Description with spaces
       ACGT
+
+
     "#},
       &NUC_ALPHABET,
     )?;
@@ -646,8 +643,8 @@ mod tests {
         index: 4,
       },
       FastaRecord {
-        seq_name: o!("SneezeC-19"),
-        desc: None,
+        seq_name: o!(""),
+        desc: Some(o!("SneezeC-19")),
         seq: o!("CCGGCGATGTRTTG--"),
         index: 5,
       },

packages/treetime/src/utils/string.rs

@@ -16,3 +16,7 @@ pub fn vec_to_string(v: Vec<char>) -> String {
 pub fn quote(x: impl Display) -> String {
   format!("\"{x}\"")
 }
+
+pub fn quote_single(x: impl Display) -> String {
+  format!("'{x}'")
+}