Upgrade fst to 0.4 and swap memmap for memmap2

Cargo.toml

@@ -21,7 +21,8 @@ rust-version = "1.88"
 clap = "2.33.0"
 csv = "1.1.2"
 error-chain = "0.12.4"
-fst = "0.3.5"
+fst = "0.4.7"
+memmap2 = "0.9.9"
 lazy_static = "1.4.0"
 ordered-float = "1.0.2"
 rayon = "1.3.0"

src/commands/pept2lca.rs

@@ -1,13 +1,13 @@
 //! The `umgap pept2lca` command.
 
-use std::fs;
 use std::io;
 use std::path::PathBuf;
 
 use rayon::iter::{ParallelBridge, ParallelIterator};
 
 use crate::errors;
 use crate::io::fasta;
+use crate::utils;
 
 #[derive(Debug, StructOpt)]
 #[structopt(verbatim_doc_comment)]
@@ -71,12 +71,7 @@ pub struct PeptToLca {
 
 /// Implements the pept2lca command
 pub fn pept2lca(args: PeptToLca) -> errors::Result<()> {
-    let fst = if args.fst_in_memory {
-        let bytes = fs::read(args.fst_file)?;
-        fst::Map::from_bytes(bytes)?
-    } else {
-        unsafe { fst::Map::from_path(args.fst_file) }?
-    };
+    let fst = utils::load_fst(&args.fst_file, args.fst_in_memory)?;
 
     let default = if args.one_on_one { Some(0) } else { None };
 

src/commands/printindex.rs

@@ -6,6 +6,7 @@ use std::path::PathBuf;
 use fst::Streamer;
 
 use crate::errors;
+use crate::utils;
 
 #[derive(Debug, StructOpt)]
 #[structopt(verbatim_doc_comment)]
@@ -40,7 +41,7 @@ pub fn printindex(args: PrintIndex) -> errors::Result<()> {
         .delimiter(b'\t')
         .from_writer(io::stdout());
 
-    let index = unsafe { fst::Map::from_path(args.fst_file) }?;
+    let index = utils::load_fst(&args.fst_file, false)?;
     let mut stream = index.stream();
 
     while let Some((k, v)) = stream.next() {

src/commands/prot2kmer2lca.rs

@@ -1,7 +1,6 @@
 //! The `umgap prot2kmer2lca` command.
 #![cfg(target_family = "unix")]
 
-use std::fs;
 use std::io;
 use std::io::Read;
 use std::io::Write;
@@ -13,6 +12,7 @@ use rayon::iter::{ParallelBridge, ParallelIterator};
 
 use crate::errors;
 use crate::io::fasta;
+use crate::utils;
 
 #[derive(Debug, StructOpt)]
 #[structopt(verbatim_doc_comment)]
@@ -106,12 +106,7 @@ pub struct ProtToKmerToLca {
 
 /// Implements the prot2kmer2lca command
 pub fn prot2kmer2lca(args: ProtToKmerToLca) -> errors::Result<()> {
-    let fst = if args.fst_in_memory {
-        let bytes = fs::read(&args.fst_file)?;
-        fst::Map::from_bytes(bytes)?
-    } else {
-        unsafe { fst::Map::from_path(&args.fst_file) }?
-    };
+    let fst = utils::load_fst(&args.fst_file, args.fst_in_memory)?;
     let default = if args.one_on_one { Some(0) } else { None };
     if let Some(socket_addr) = &args.socket {
         let listener = UnixListener::bind(socket_addr)?;
@@ -150,7 +145,7 @@ pub fn prot2kmer2lca(args: ProtToKmerToLca) -> errors::Result<()> {
 fn stream_prot2kmer2lca<R, W>(
     input: R,
     output: W,
-    fst: &fst::Map,
+    fst: &fst::Map<utils::FstData>,
     k: usize,
     chunk_size: usize,
     default: Option<u64>,

src/commands/prot2tryp2lca.rs

@@ -1,18 +1,16 @@
 //! The `umgap prot2tryp2lca` command.
 
 use std::collections::HashSet;
-use std::fs;
 use std::io;
 use std::path::PathBuf;
 
-use fst;
-
 use regex;
 
 use rayon::iter::{ParallelBridge, ParallelIterator};
 
 use crate::errors;
 use crate::io::fasta;
+use crate::utils;
 
 #[derive(Debug, StructOpt)]
 #[structopt(verbatim_doc_comment)]
@@ -86,12 +84,7 @@ pub struct ProtToTrypToLca {
 
 /// Implements the prot2tryp2lca command.
 pub fn prot2tryp2lca(args: ProtToTrypToLca) -> errors::Result<()> {
-    let fst = if args.fst_in_memory {
-        let bytes = fs::read(&args.fst_file)?;
-        fst::Map::from_bytes(bytes)?
-    } else {
-        unsafe { fst::Map::from_path(&args.fst_file) }?
-    };
+    let fst = utils::load_fst(&args.fst_file, args.fst_in_memory)?;
     let default = if args.one_on_one { Some(0) } else { None };
     let pattern = regex::Regex::new(&args.pattern)?;
     let contains = args.contains.chars().collect::<HashSet<char>>();

src/utils.rs

@@ -19,3 +19,37 @@ impl<E, I: Iterator<Item = E>> Iterator for Zip<E, I> {
         self.parts.iter_mut().map(|part| part.next()).collect()
     }
 }
+
+/// The bytes backing an FST index.
+pub enum FstData {
+    /// The whole index, read into memory.
+    Memory(Vec<u8>),
+    /// The index file, memory mapped.
+    Mapped(memmap2::Mmap),
+}
+
+impl AsRef<[u8]> for FstData {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            FstData::Memory(bytes) => bytes,
+            FstData::Mapped(mmap) => mmap,
+        }
+    }
+}
+
+/// Opens an FST index, reading it into memory if `in_memory` is set and memory mapping
+/// it otherwise.
+///
+/// fst dropped its own mmap support in 0.4, so this carries the same caveat its
+/// `Map::from_path` did: the file must not be modified while the index is in use.
+pub fn load_fst<P: AsRef<std::path::Path>>(
+    path: P,
+    in_memory: bool,
+) -> crate::errors::Result<fst::Map<FstData>> {
+    let data = if in_memory {
+        FstData::Memory(std::fs::read(path)?)
+    } else {
+        FstData::Mapped(unsafe { memmap2::Mmap::map(&std::fs::File::open(path)?)? })
+    };
+    Ok(fst::Map::new(data)?)
+}