diff --git a/Cargo.lock b/Cargo.lock index 9eba4bf2..724764dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -101,12 +101,6 @@ version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" - [[package]] name = "bytes" version = "1.12.1" @@ -306,13 +300,9 @@ dependencies = [ [[package]] name = "fst" -version = "0.3.5" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "927fb434ff9f0115b215dc0efd2e4fbdd7448522a92a1aa37c77d6a2f8f1ebd6" -dependencies = [ - "byteorder", - "memmap", -] +checksum = "7ab85b9b05e3978cc9a9cf8fea7f01b494e1a09ed3037e16ba39edc7a29eb61a" [[package]] name = "getrandom" @@ -507,13 +497,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] -name = "memmap" -version = "0.6.2" +name = "memmap2" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2ffa2c986de11a9df78620c01eeaaf27d94d3ff02bf81bfcca953102dd0c6ff" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", - "winapi", ] [[package]] @@ -995,6 +984,7 @@ dependencies = [ "error-chain", "fst", "lazy_static", + "memmap2", "ordered-float", "rayon", "regex", diff --git a/Cargo.toml b/Cargo.toml index 650150b3..d9d37b4e 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/commands/pept2lca.rs b/src/commands/pept2lca.rs index f3f4426c..362ee0c5 100644 --- a/src/commands/pept2lca.rs +++ b/src/commands/pept2lca.rs @@ -1,6 +1,5 @@ //! The `umgap pept2lca` command. -use std::fs; use std::io; use std::path::PathBuf; @@ -8,6 +7,7 @@ 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 }; diff --git a/src/commands/printindex.rs b/src/commands/printindex.rs index 1208ebd7..298d845c 100644 --- a/src/commands/printindex.rs +++ b/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() { diff --git a/src/commands/prot2kmer2lca.rs b/src/commands/prot2kmer2lca.rs index 734cd71b..b397ebf5 100644 --- a/src/commands/prot2kmer2lca.rs +++ b/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( input: R, output: W, - fst: &fst::Map, + fst: &fst::Map, k: usize, chunk_size: usize, default: Option, diff --git a/src/commands/prot2tryp2lca.rs b/src/commands/prot2tryp2lca.rs index f1311bac..a4292959 100644 --- a/src/commands/prot2tryp2lca.rs +++ b/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::>(); diff --git a/src/utils.rs b/src/utils.rs index b79f2f7e..c1f0fa0b 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -19,3 +19,37 @@ impl> Iterator for Zip { 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), + /// 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>( + path: P, + in_memory: bool, +) -> crate::errors::Result> { + 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)?) +}