diff --git a/src/obikindex/src/layer/content_layer.rs b/src/obikindex/src/layer/content_layer.rs index bea0dd5e..a6f54866 100644 --- a/src/obikindex/src/layer/content_layer.rs +++ b/src/obikindex/src/layer/content_layer.rs @@ -28,9 +28,11 @@ use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; use obikseq::CanonicalKmer; use crate::layer::error::OLMResult; -use crate::layer::typed_layer::{LayerContent, TypedLayer, COUNTS_DIR, PRESENCE_DIR}; use crate::layer::meta::IndexMode; -use crate::layer::mphf_layer::{EvidenceKind, EVIDENCE_FILE, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE}; +use crate::layer::mphf_layer::{ + EVIDENCE_FILE, EvidenceKind, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE, +}; +use crate::layer::typed_layer::{COUNTS_DIR, LayerContent, PRESENCE_DIR, TypedLayer}; /// One layer, at any point in its life — see the module docs. Only /// [`Empty`](Layer::Empty) and the two ready-to-read states @@ -42,7 +44,9 @@ pub enum Layer { /// the path accessors panics on this variant: calling them means the /// caller assumed a layer was ready when it wasn't, an implementation /// error to surface loudly, not paper over with a default value. - Empty { dir: PathBuf }, + Empty { + dir: PathBuf, + }, Count(TypedLayer), Presence(TypedLayer), } @@ -53,7 +57,9 @@ impl Layer { /// building a new layer. pub fn create(dir: &Path) -> std::io::Result { std::fs::create_dir_all(dir)?; - Ok(Layer::Empty { dir: dir.to_owned() }) + Ok(Layer::Empty { + dir: dir.to_owned(), + }) } /// Open one layer, auto-detecting count vs. presence from what is @@ -66,7 +72,8 @@ impl Layer { /// here. pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult { if with_counts && layer_dir.join(COUNTS_DIR).exists() { - return TypedLayer::::open(layer_dir, mode).map(Layer::Count); + return TypedLayer::::open(layer_dir, mode) + .map(Layer::Count); } TypedLayer::::open(layer_dir, mode).map(Layer::Presence) } @@ -91,12 +98,24 @@ impl Layer { } } - pub fn mphf_path(&self) -> PathBuf { self.dir().join(MPHF_FILE) } - pub fn unitigs_path(&self) -> PathBuf { self.dir().join(UNITIGS_FILE) } - pub fn evidence_path(&self) -> PathBuf { self.dir().join(EVIDENCE_FILE) } - pub fn fingerprint_path(&self) -> PathBuf { self.dir().join(FINGERPRINT_FILE) } - pub fn counts_dir(&self) -> PathBuf { self.dir().join(COUNTS_DIR) } - pub fn presence_dir(&self) -> PathBuf { self.dir().join(PRESENCE_DIR) } + pub fn mphf_path(&self) -> PathBuf { + self.dir().join(MPHF_FILE) + } + pub fn unitigs_path(&self) -> PathBuf { + self.dir().join(UNITIGS_FILE) + } + pub fn evidence_path(&self) -> PathBuf { + self.dir().join(EVIDENCE_FILE) + } + pub fn fingerprint_path(&self) -> PathBuf { + self.dir().join(FINGERPRINT_FILE) + } + pub fn counts_dir(&self) -> PathBuf { + self.dir().join(COUNTS_DIR) + } + pub fn presence_dir(&self) -> PathBuf { + self.dir().join(PRESENCE_DIR) + } // ── Ready-only surface ─────────────────────────────────────────────── @@ -135,11 +154,11 @@ impl Layer { /// Raw MPHF batch lookup: kmer → slot, no membership check — for /// callers that already know every kmer is a member of *this* layer, so /// the evidence check `find_slot`/`find` would perform is redundant. - pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec { + pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec { match self { - Layer::Count(l) => l.index_batch(kmers), - Layer::Presence(l) => l.index_batch(kmers), - Layer::Empty { .. } => panic!("Layer::index_batch() called on an Empty layer"), + Layer::Count(l) => l.hash_batch(kmers), + Layer::Presence(l) => l.hash_batch(kmers), + Layer::Empty { .. } => panic!("Layer::hash_batch() called on an Empty layer"), } } @@ -168,7 +187,61 @@ impl Layer { o.extend(c.iter().map(|&v| v != 0)); } } - Layer::Empty { .. } => panic!("Layer::fill_sub_matrix_carries() called on an Empty layer"), + Layer::Empty { .. } => { + panic!("Layer::fill_sub_matrix_carries() called on an Empty layer") + } + } + } + + /// Raw MPHF lookup: kmer → slot, no membership check. + pub fn hash(&self, kmer: CanonicalKmer) -> usize { + match self { + Layer::Count(l) => l.hash(kmer), + Layer::Presence(l) => l.hash(kmer), + Layer::Empty { .. } => panic!("Layer::hash() called on an Empty layer"), + } + } + + /// Iterate over all canonical kmers in the layer, in deterministic order. + pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter { + match self { + Layer::Count(l) => l.iter_kmers(), + Layer::Presence(l) => l.iter_kmers(), + Layer::Empty { .. } => panic!("Layer::iter_kmers() called on an Empty layer"), + } + } + + /// Iterate over all canonical kmers, each paired with its zero-based + /// sequence index in `unitigs.bin`. + pub fn enumerate_kmers(&self) -> std::iter::Enumerate { + match self { + Layer::Count(l) => l.enumerate_kmers(), + Layer::Presence(l) => l.enumerate_kmers(), + Layer::Empty { .. } => panic!("Layer::enumerate_kmers() called on an Empty layer"), + } + } + + /// Iterate over the layer's canonical kmers in batches of `n`. + pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter { + match self { + Layer::Count(l) => l.iter_kmers_batch(n), + Layer::Presence(l) => l.iter_kmers_batch(n), + Layer::Empty { .. } => panic!("Layer::iter_kmers_batch() called on an Empty layer"), + } + } + + /// Iterate over batches, each paired with the zero-based index of the + /// first kmer in the batch. + pub fn enumerate_kmers_batch( + &self, + n: usize, + ) -> Box)> + Send + 'static> { + match self { + Layer::Count(l) => Box::new(l.enumerate_kmers_batch(n)), + Layer::Presence(l) => Box::new(l.enumerate_kmers_batch(n)), + Layer::Empty { .. } => { + panic!("Layer::enumerate_kmers_batch() called on an Empty layer") + } } } } diff --git a/src/obikindex/src/layer/mphf_layer.rs b/src/obikindex/src/layer/mphf_layer.rs index 993cc6df..b290d1d9 100644 --- a/src/obikindex/src/layer/mphf_layer.rs +++ b/src/obikindex/src/layer/mphf_layer.rs @@ -15,9 +15,9 @@ use crate::layer::evidence::{Evidence, EvidenceWriter}; use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter}; use crate::layer::meta::IndexMode; -pub(crate) const MPHF_FILE: &str = "mphf.bin"; -pub(crate) const UNITIGS_FILE: &str = "unitigs.bin"; -pub(crate) const EVIDENCE_FILE: &str = "evidence.bin"; +pub(crate) const MPHF_FILE: &str = "mphf.bin"; +pub(crate) const UNITIGS_FILE: &str = "unitigs.bin"; +pub(crate) const EVIDENCE_FILE: &str = "evidence.bin"; pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin"; /// Owned MPHF — used only at build time (construction + store). @@ -30,9 +30,20 @@ type MphfEps = PtrHash, Xx // ── LayerEvidence ───────────────────────────────────────────────────────────── enum LayerEvidence { - Exact { evidence: Evidence, unitigs: Arc }, - Approx { fingerprint: FingerprintVec, unitigs: Arc, unitigs_path: PathBuf }, - Hybrid { evidence: Evidence, unitigs: Arc, fingerprint: FingerprintVec }, + Exact { + evidence: Evidence, + unitigs: Arc, + }, + Approx { + fingerprint: FingerprintVec, + unitigs: Arc, + unitigs_path: PathBuf, + }, + Hybrid { + evidence: Evidence, + unitigs: Arc, + fingerprint: FingerprintVec, + }, } // ── EvidenceKind ────────────────────────────────────────────────────────────── @@ -51,14 +62,15 @@ impl EvidenceKind { /// `evidence.bin`/`fingerprint.bin` alone determines the mode — same /// signal `MphfLayer::open` uses, just without opening either file. pub fn detect(layer_dir: &Path) -> OLMResult { - let has_evidence = layer_dir.join(EVIDENCE_FILE).exists(); + let has_evidence = layer_dir.join(EVIDENCE_FILE).exists(); let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists(); match (has_evidence, has_fingerprint) { - (true, false) => Ok(EvidenceKind::Exact), - (false, true) => Ok(EvidenceKind::Approx), - (true, true) => Ok(EvidenceKind::Hybrid), + (true, false) => Ok(EvidenceKind::Exact), + (false, true) => Ok(EvidenceKind::Approx), + (true, true) => Ok(EvidenceKind::Hybrid), (false, false) => Err(OLMError::InvalidLayer(format!( - "no evidence.bin or fingerprint.bin in {}", layer_dir.display() + "no evidence.bin or fingerprint.bin in {}", + layer_dir.display() ))), } } @@ -74,8 +86,8 @@ impl EvidenceKind { /// O(n) sequential scan on Approx layers. pub struct MphfLayer { mphf: MemCase, - ev: LayerEvidence, - n: usize, + ev: LayerEvidence, + n: usize, } impl MphfLayer { @@ -96,14 +108,28 @@ impl MphfLayer { let n = fingerprint.n(); let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?); let unitigs_path = dir.join(UNITIGS_FILE); - (LayerEvidence::Approx { fingerprint, unitigs, unitigs_path }, n) + ( + LayerEvidence::Approx { + fingerprint, + unitigs, + unitigs_path, + }, + n, + ) } IndexMode::Hybrid { .. } => { - let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?; + let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?; let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?; let n = evidence.len(); let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?); - (LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n) + ( + LayerEvidence::Hybrid { + evidence, + unitigs, + fingerprint, + }, + n, + ) } }; Ok(Self { mphf, ev, n }) @@ -118,9 +144,13 @@ impl MphfLayer { #[inline] pub fn find(&self, kmer: CanonicalKmer) -> Option { let slot = self.mphf.index(&kmer.raw()); - if slot >= self.n { return None; } + if slot >= self.n { + return None; + } match &self.ev { - LayerEvidence::Exact { evidence, unitigs, .. } => { + LayerEvidence::Exact { + evidence, unitigs, .. + } => { let (chunk_id, rank) = evidence.decode(slot); if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) { Some(slot) @@ -128,9 +158,13 @@ impl MphfLayer { None } } - LayerEvidence::Approx { fingerprint, .. } | - LayerEvidence::Hybrid { fingerprint, .. } => { - if fingerprint.matches(slot, kmer.seq_hash()) { Some(slot) } else { None } + LayerEvidence::Approx { fingerprint, .. } + | LayerEvidence::Hybrid { fingerprint, .. } => { + if fingerprint.matches(slot, kmer.seq_hash()) { + Some(slot) + } else { + None + } } } } @@ -142,10 +176,16 @@ impl MphfLayer { /// that owns the slot, then exact comparison. pub fn find_strict(&self, kmer: CanonicalKmer) -> Option { let slot = self.mphf.index(&kmer.raw()); - if slot >= self.n { return None; } + if slot >= self.n { + return None; + } match &self.ev { - LayerEvidence::Exact { evidence, unitigs, .. } | - LayerEvidence::Hybrid { evidence, unitigs, .. } => { + LayerEvidence::Exact { + evidence, unitigs, .. + } + | LayerEvidence::Hybrid { + evidence, unitigs, .. + } => { let (chunk_id, rank) = evidence.decode(slot); if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) { Some(slot) @@ -165,38 +205,16 @@ impl MphfLayer { } } - /// Reconstruct the canonical k-mer stored at `slot` — the inverse of - /// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot). - /// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives - /// `(chunk_id, rank)` directly (no MPHF hashing, no file scan), then - /// `unitigs.canonical_raw_kmer` is a direct-access read. `None` on - /// `Approx` layers — fingerprints alone can't recover a k-mer, and - /// falling back to a full sequential scan here would silently make one - /// call cost O(n); callers needing this on an `Approx` layer should - /// scan `unitigs.bin` themselves and decide how to handle that cost - /// explicitly. - pub fn kmer_at(&self, slot: usize) -> Option { - if slot >= self.n { - return None; - } - match &self.ev { - LayerEvidence::Exact { evidence, unitigs, .. } | - LayerEvidence::Hybrid { evidence, unitigs, .. } => { - let (chunk_id, rank) = evidence.decode(slot); - let raw = unitigs.canonical_raw_kmer(chunk_id as usize, rank as usize); - Some(CanonicalKmer::from_raw_unchecked(raw)) - } - LayerEvidence::Approx { .. } => None, - } + /// Number of slots in the MPHF layer — equal to the number of stored k-mers. + pub fn n(&self) -> usize { + self.n } - pub fn n(&self) -> usize { self.n } - /// This already-open layer's evidence mode — reads the discriminant /// already in memory, no disk access. pub fn evidence_kind(&self) -> EvidenceKind { match &self.ev { - LayerEvidence::Exact { .. } => EvidenceKind::Exact, + LayerEvidence::Exact { .. } => EvidenceKind::Exact, LayerEvidence::Approx { .. } => EvidenceKind::Approx, LayerEvidence::Hybrid { .. } => EvidenceKind::Hybrid, } @@ -208,7 +226,7 @@ impl MphfLayer { /// check against the layer's evidence. The returned slot is only meaningful /// when the kmer is actually present in the layer; callers that need an /// existence test should use [`find`](Self::find) instead. - pub fn index(&self, kmer: CanonicalKmer) -> usize { + pub fn hash(&self, kmer: CanonicalKmer) -> usize { self.mphf.index(&kmer.raw()) } @@ -217,7 +235,7 @@ impl MphfLayer { /// Returns a [`Vec`] of slots, one per input kmer, in the same order as the /// input slice. Like [`index`](Self::index), no membership check is /// performed. - pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec { + pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec { kmers.iter().map(|k| self.mphf.index(&k.raw())).collect() } @@ -237,7 +255,9 @@ impl MphfLayer { LayerEvidence::Approx { unitigs, .. } => unitigs, LayerEvidence::Hybrid { unitigs, .. } => unitigs, }; - KmerIter { inner: Box::new(reader.iter_indexed_canonical_kmers_owned()) } + KmerIter { + inner: Box::new(reader.iter_indexed_canonical_kmers_owned()), + } } /// Iterate over all canonical kmers, each paired with its zero-based @@ -255,7 +275,10 @@ impl MphfLayer { /// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n` /// kmers. The final batch may be shorter when the layer is exhausted. pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter { - KmerBatchIter { inner: self.iter_kmers(), batch_size: n } + KmerBatchIter { + inner: self.iter_kmers(), + batch_size: n, + } } /// Iterate over batches, each paired with the zero-based index of the @@ -265,7 +288,10 @@ impl MphfLayer { /// `batch_start_index` is the iteration-order offset of the first kmer /// in that batch within the full layer sequence — i.e. a multiple of `n` /// except for the final (possibly shorter) batch. - pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator)> + Send + 'static { + pub fn enumerate_kmers_batch( + &self, + n: usize, + ) -> impl Iterator)> + Send + 'static { let mut offset = 0usize; self.iter_kmers_batch(n).map(move |batch| { let base = offset; @@ -320,11 +346,7 @@ impl Iterator for KmerBatchIter { break; } } - if batch.is_empty() { - None - } else { - Some(batch) - } + if batch.is_empty() { None } else { Some(batch) } } } @@ -373,7 +395,7 @@ impl MphfLayer { let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE)) .map_err(|e| OLMError::InvalidLayer(e.to_string()))?; - let mut ev = EvidenceWriter::new(n); + let mut ev = EvidenceWriter::new(n); let mut seen = vec![0u8; (n + 7) / 8]; for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() { @@ -382,7 +404,7 @@ impl MphfLayer { return Err(OLMError::Mphf("slot out of bounds".into())); } let byte = slot / 8; - let bit = 1u8 << (slot % 8); + let bit = 1u8 << (slot % 8); if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } @@ -398,7 +420,9 @@ impl MphfLayer { /// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`. pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult { if b == 0 || b > 64 { - return Err(OLMError::InvalidLayer("fingerprint width must be 1..=64".into())); + return Err(OLMError::InvalidLayer( + "fingerprint width must be 1..=64".into(), + )); } if z == 0 { return Err(OLMError::InvalidLayer("z must be ≥ 1".into())); @@ -452,9 +476,8 @@ impl MphfLayer { // ── Empty layer ─────────────────────────────────────────────────────── if n == 0 { - let mphf: Mphf = - Mphf::try_new(&[] as &[u64], PtrHashParams::::default()) - .ok_or_else(|| OLMError::Mphf("construction failed".into()))?; + let mphf: Mphf = Mphf::try_new(&[] as &[u64], PtrHashParams::::default()) + .ok_or_else(|| OLMError::Mphf("construction failed".into()))?; mphf.store(&dir.join(MPHF_FILE)) .map_err(|e| OLMError::InvalidLayer(e.to_string()))?; match mode { @@ -474,9 +497,11 @@ impl MphfLayer { // ── Pass 1: MPHF via clonable mmap iterator ─────────────────────────── let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?; - let mphf: Mphf = - Mphf::new_from_par_iter(n, keys.map(|k| k.raw()).par_bridge(), - PtrHashParams::::default()); + let mphf: Mphf = Mphf::new_from_par_iter( + n, + keys.map(|k| k.raw()).par_bridge(), + PtrHashParams::::default(), + ); mphf.store(&dir.join(MPHF_FILE)) .map_err(|e| OLMError::InvalidLayer(e.to_string()))?; @@ -489,9 +514,14 @@ impl MphfLayer { let mut ev = EvidenceWriter::new(n); for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() { let slot = mphf.index(&kmer.raw()); - if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } - let byte = slot / 8; let bit = 1u8 << (slot % 8); - if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } + if slot >= n { + return Err(OLMError::Mphf("slot out of bounds".into())); + } + let byte = slot / 8; + let bit = 1u8 << (slot % 8); + if seen[byte] & bit != 0 { + return Err(OLMError::Mphf("duplicate slot".into())); + } seen[byte] |= bit; ev.set(slot, chunk_id as u32, rank as u8); fill_slot(slot, kmer)?; @@ -504,9 +534,14 @@ impl MphfLayer { let mut fw = FingerprintVecWriter::new(n, *b); for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() { let slot = mphf.index(&kmer.raw()); - if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } - let byte = slot / 8; let bit = 1u8 << (slot % 8); - if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } + if slot >= n { + return Err(OLMError::Mphf("slot out of bounds".into())); + } + let byte = slot / 8; + let bit = 1u8 << (slot % 8); + if seen[byte] & bit != 0 { + return Err(OLMError::Mphf("duplicate slot".into())); + } seen[byte] |= bit; fw.set(slot, kmer.seq_hash()); fill_slot(slot, kmer)?; @@ -519,9 +554,14 @@ impl MphfLayer { let mut fw = FingerprintVecWriter::new(n, *b); for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() { let slot = mphf.index(&kmer.raw()); - if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } - let byte = slot / 8; let bit = 1u8 << (slot % 8); - if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); } + if slot >= n { + return Err(OLMError::Mphf("slot out of bounds".into())); + } + let byte = slot / 8; + let bit = 1u8 << (slot % 8); + if seen[byte] & bit != 0 { + return Err(OLMError::Mphf("duplicate slot".into())); + } seen[byte] |= bit; ev.set(slot, chunk_id as u32, rank as u8); fw.set(slot, kmer.seq_hash()); diff --git a/src/obikindex/src/layer/typed_layer.rs b/src/obikindex/src/layer/typed_layer.rs index 073e2789..82986704 100644 --- a/src/obikindex/src/layer/typed_layer.rs +++ b/src/obikindex/src/layer/typed_layer.rs @@ -3,10 +3,8 @@ use std::fs; use std::path::{Path, PathBuf}; use obicompactvec::{ - BinaryMatrix, - PersistentBitMatrix, PersistentBitMatrixBuilder, - PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, - PersistentSparseBitMatrix, + BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix, + PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix, }; use obikseq::CanonicalKmer; use obiskio::{UnitigFileReader, UnitigFileWriter}; @@ -16,7 +14,7 @@ use crate::layer::meta::IndexMode; use crate::layer::mphf_layer::MphfLayer; pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE; -pub(crate) const COUNTS_DIR: &str = "counts"; +pub(crate) const COUNTS_DIR: &str = "counts"; pub(crate) const PRESENCE_DIR: &str = "presence"; // ── Trait ───────────────────────────────────────────────────────────────────── @@ -82,7 +80,9 @@ pub fn open_data(root: &Path, i: usize) -> OLMResult { impl LayerData for () { type Item = (); - fn open(_layer_dir: &Path) -> OLMResult { Ok(()) } + fn open(_layer_dir: &Path) -> OLMResult { + Ok(()) + } fn read(&self, _slot: usize) {} } @@ -91,7 +91,9 @@ impl LayerData for PersistentCompactIntMatrix { fn open(layer_dir: &Path) -> OLMResult { PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io) } - fn read(&self, slot: usize) -> Box<[u32]> { self.row(slot) } + fn read(&self, slot: usize) -> Box<[u32]> { + self.row(slot) + } } impl LayerData for PersistentBitMatrix { @@ -99,7 +101,9 @@ impl LayerData for PersistentBitMatrix { fn open(layer_dir: &Path) -> OLMResult { PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io) } - fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) } + fn read(&self, slot: usize) -> Box<[bool]> { + self.row(slot) + } } impl LayerData for PersistentSparseBitMatrix { @@ -107,7 +111,9 @@ impl LayerData for PersistentSparseBitMatrix { fn open(layer_dir: &Path) -> OLMResult { PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io) } - fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) } + fn read(&self, slot: usize) -> Box<[bool]> { + self.row(slot) + } } // ── LayerContent ───────────────────────────────────────────────────────────── @@ -202,7 +208,10 @@ impl TypedLayer { } pub fn query(&self, kmer: CanonicalKmer) -> Option> { - self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) }) + self.mphf.find(kmer).map(|slot| Hit { + slot, + data: self.data.read(slot), + }) } /// MPHF + evidence membership check only — no data read. For callers @@ -214,16 +223,18 @@ impl TypedLayer { self.mphf.find(kmer) } - pub fn n(&self) -> usize { self.mphf.n() } + pub fn n(&self) -> usize { + self.mphf.n() + } /// Raw MPHF lookup: kmer → slot, no membership check. - pub fn index(&self, kmer: CanonicalKmer) -> usize { - self.mphf.index(kmer) + pub fn hash(&self, kmer: CanonicalKmer) -> usize { + self.mphf.hash(kmer) } /// Batch raw MPHF lookup: kmers → slots, no membership check. - pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec { - self.mphf.index_batch(kmers) + pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec { + self.mphf.hash_batch(kmers) } /// Iterate over all canonical kmers in the layer, in deterministic order. @@ -244,7 +255,10 @@ impl TypedLayer { /// Iterate over batches, each paired with the zero-based index of the /// first kmer in the batch. - pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator)> + Send + 'static { + pub fn enumerate_kmers_batch( + &self, + n: usize, + ) -> impl Iterator)> + Send + 'static { self.mphf.enumerate_kmers_batch(n) } @@ -301,8 +315,12 @@ impl TypedLayer<()> { pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> { let presence_dir = layer_dir.join(PRESENCE_DIR); fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?; - let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?; - mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?; + let mut mb = + PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?; + mb.add_col_ones() + .map_err(OLMError::Io)? + .close() + .map_err(OLMError::Io)?; mb.close().map_err(OLMError::Io) } } @@ -318,8 +336,8 @@ impl TypedLayer { ) -> OLMResult { let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers(); let counts_dir = out_dir.join(COUNTS_DIR); - let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir) - .map_err(OLMError::Io)?; + let mut mb = + PersistentCompactIntMatrixBuilder::new(n, &counts_dir).map_err(OLMError::Io)?; let mut col = mb.add_col().map_err(OLMError::Io)?; let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| { col.set(slot, count_of(kmer)); @@ -336,7 +354,9 @@ impl TypedLayer { mode: &IndexMode, counts: &HashMap, ) -> OLMResult { - Self::build(out_dir, block_bits, mode, |kmer| counts.get(&kmer).copied().unwrap_or(0)) + Self::build(out_dir, block_bits, mode, |kmer| { + counts.get(&kmer).copied().unwrap_or(0) + }) } } diff --git a/src/obikphylo/src/siblings/cache.rs b/src/obikphylo/src/siblings/cache.rs index 03450e47..1ee70260 100644 --- a/src/obikphylo/src/siblings/cache.rs +++ b/src/obikphylo/src/siblings/cache.rs @@ -1,7 +1,7 @@ use rayon::prelude::*; -use obikseq::CanonicalKmer; use obikindex::layer::Layer; +use obikseq::CanonicalKmer; use obisys::progress_bar; use obikindex::{KmerIndex, OKIResult}; @@ -45,11 +45,7 @@ pub(super) struct PartitionCache { } impl PartitionCache { - pub(super) fn build( - index: &KmerIndex, - n_parts: usize, - with_counts: bool, - ) -> OKIResult { + pub(super) fn build(index: &KmerIndex, n_parts: usize, with_counts: bool) -> OKIResult { let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); let built: Vec<(Vec, usize)> = (0..n_parts) .into_par_iter() @@ -200,7 +196,7 @@ impl PartitionCache { let mat = &mats[li]; let variants: Vec = entries.iter().map(|&(variant, _, _)| variant).collect(); - let slots = mat.index_batch(&variants); + let slots = mat.hash_batch(&variants); let hits: Vec<(usize, usize, u8)> = slots .into_iter() .zip(entries.iter()) diff --git a/src/obikphylo/src/siblings/family_scan.rs b/src/obikphylo/src/siblings/family_scan.rs index 678449d6..01719903 100644 --- a/src/obikphylo/src/siblings/family_scan.rs +++ b/src/obikphylo/src/siblings/family_scan.rs @@ -53,19 +53,19 @@ use std::sync::atomic::{AtomicU8, Ordering}; use rayon::prelude::*; -use obikseq::CanonicalKmer; use obikindex::layer::meta::PartitionMeta; +use obikseq::CanonicalKmer; use obipipeline::{ThrottleGuard, throttle}; -use obikindex::{OKIError, OKIResult}; use obikindex::KmerIndex; +use obikindex::{OKIError, OKIResult}; use obikindex::layer::Layer; use super::cache::PartitionCache; use super::helpers::central_base; use super::iter::{SiblingEntry, SiblingLayerExt}; -use super::{olm_to_ok, FamilyMask, SiblingAnnex, ANNEX_FILE_NAME}; +use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex, olm_to_ok}; /// Families per batch — see the module docs for the memory-vs-per-partition- /// density trade-off this picks a point on. At ~90 genomes and a few @@ -191,14 +191,22 @@ pub(super) fn scan_layer_families( selection: &Selection, mut on_family: impl FnMut(usize, FamilyMask, &[u8]), ) -> OKIResult<()> { - let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let index_dir = layer_dir + .parent() + .expect("layer_dir has a parent index dir"); let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?); let mat = Layer::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?; let n_cols = mat.n_cols().min(n_genomes); - let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k }); + let ctx = Arc::new(LayerCtx { + mat, + n_parts, + n_genomes, + n_cols, + k, + }); // Streamed straight from `iter_minorants_batch` (zips this layer's own // `iter_kmers()` with the annex, both in iteration order — never an @@ -209,11 +217,14 @@ pub(super) fn scan_layer_families( // (see the module docs). `.scan()` computes each batch's starting // family index lazily, mirroring what the eager `chunks()`+running // `offset` used to do. - let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| { - let start = *offset; - *offset += batch.len(); - Some((start, batch)) - }); + let batches = + ctx.mat + .iter_minorants_batch(annex, FAMILY_BATCH) + .scan(0usize, |offset, batch| { + let start = *offset; + *offset += batch.len(); + Some((start, batch)) + }); let n_workers = obisys::effective_parallelism(); let capacity = 4; @@ -288,7 +299,7 @@ pub(super) fn scan_layer_families( // — a pure MPHF lookup, no evidence check, since these are // this layer's own kmers, known members by construction. let kmers: Vec = batch.entries.iter().map(|e| e.kmer).collect(); - let slots = ctx.mat.index_batch(&kmers); + let slots = ctx.mat.hash_batch(&kmers); let mut carries: Vec> = (0..ctx.n_cols).map(|_| Vec::new()).collect(); ctx.mat.fill_sub_matrix_carries(&slots, &mut carries); for (g, col) in carries.iter().enumerate() { @@ -321,18 +332,24 @@ pub(super) fn scan_layer_families( // (one batch at a time, never several concurrently — see the // module docs), each thread owning one partition's queries // contiguously until this batch is done. - let genome_mask: Vec = batch.genome_mask.into_iter().map(AtomicU8::new).collect(); + let genome_mask: Vec = + batch.genome_mask.into_iter().map(AtomicU8::new).collect(); let fast_mode = cache.fast_mode(); - batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { - let on_hit = |i: usize, base: u8, g: usize| { - genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); - }; - if fast_mode { - cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit); - } else { - cache.find_presence_batch(dest, queries, n_genomes, on_hit); - } - }); + batch + .outgoing + .par_iter() + .enumerate() + .filter(|(_, q)| !q.is_empty()) + .for_each(|(dest, queries)| { + let on_hit = |i: usize, base: u8, g: usize| { + genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); + }; + if fast_mode { + cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit); + } else { + cache.find_presence_batch(dest, queries, n_genomes, on_hit); + } + }); for i in 0..n { let family_idx = batch.start_family_idx + i; @@ -347,7 +364,10 @@ pub(super) fn scan_layer_families( next_expected += n; } } - debug_assert!(pending.is_empty(), "every generated batch must have been replayed"); + debug_assert!( + pending.is_empty(), + "every generated batch must have been replayed" + ); Ok(()) }