Extract k-mer counting logic into a dedicated counter module

Decouple k-mer counting from the partitioner by introducing a new `Counter` struct. The module exposes a fluent builder API with optional partial file retention, executes partition processing in parallel via Rayon with memory-aware chunk sizing, and integrates thread-safe progress callbacks. Update all callers to use the new counter, simplify test pipelines by removing serialization overhead, and clarify algorithm separation in module documentation.

src/obikindexer/src/algorithms/counter/mod.rs

@@ -0,0 +1,133 @@
+//! Kmer counting — the sibling algorithm to
+//! [`crate::algorithms::dereplicator`]: runs after `dereplicator::
+//! Dereplicator::run` has deduplicated each partition's raw superkmer
+//! file, the last step before `obikindex::KmerIndex::build_layers` turns
+//! the provisional per-partition MPHF/counts this produces into the real
+//! layer 0 — see `DevDocMD/implementation/partition_layer_cache.md`.
+
+mod count;
+mod kmer_sort;
+
+use std::collections::BTreeMap;
+use std::fs;
+use std::io;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use obikindex::KmerIndex;
+use obiskio::SKResult;
+use obisys::Progress;
+use rayon::prelude::*;
+use sysinfo::System;
+
+use count::count_partition;
+use kmer_sort::chunk_size_from_ram;
+
+pub struct KmerSpectrum {
+    pub f0: u64,
+    pub f1: u64,
+    pub counts: BTreeMap<u32, u64>,
+}
+
+/// For each partition with a dereplicated superkmer file:
+/// 1. Enumerates all unique canonical kmers (two passes over the file).
+/// 2. Builds a provisional MPHF (FMPHGO) over those kmers.
+/// 3. Writes a flat binary count file (`counts1.bin`, one `u32` per slot,
+///    memory-mapped) accumulating kmer abundances from the superkmer counts.
+/// 4. Persists the MPHF to `mphf1.bin` for downstream use.
+///
+/// Two-phase construction, same shape as `PartitionRouter`/`Dereplicator`:
+/// `new` (no disk access), optional `keep_partial` setter, then `run`,
+/// which returns the aggregated [`KmerSpectrum`].
+pub struct Counter<'a> {
+    index: &'a KmerIndex,
+    n_partitions: usize,
+    keep_partial: bool,
+}
+
+impl<'a> Counter<'a> {
+    pub fn new(index: &'a KmerIndex) -> Self {
+        Self {
+            index,
+            n_partitions: index.n_partitions(),
+            keep_partial: false,
+        }
+    }
+
+    /// Keep each partition's own `kmer_spectrum_raw.json` after aggregation
+    /// instead of deleting it (default: `false`).
+    pub fn keep_partial(mut self, v: bool) -> Self {
+        self.keep_partial = v;
+        self
+    }
+
+    /// Count every partition in parallel, then aggregate their spectra.
+    ///
+    /// `on_progress`, when set, is called once per completed partition, from
+    /// whichever rayon worker thread finished it — `Fn(...) + Sync`, not
+    /// `FnMut`, same reason as `Dereplicator::run`: the counting pass itself
+    /// is a parallel `par_iter`, so the callback must tolerate concurrent
+    /// calls. `total: Some(n_partitions)` — known up front. This algorithm
+    /// never renders anything itself.
+    pub fn run(&self, on_progress: Option<impl Fn(Progress) + Sync>) -> SKResult<KmerSpectrum> {
+        let sys = System::new_all();
+        // available_memory() can return 0 on macOS when the compressor page count exceeds
+        // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total.
+        let available = match sys.available_memory() {
+            0 => sys.total_memory() / 2,
+            n => n,
+        };
+        let n_threads = rayon::current_num_threads().max(1) as u64;
+        let chunk_kmers = chunk_size_from_ram(available / n_threads);
+        let done = AtomicU64::new(0);
+
+        let results: Vec<SKResult<()>> = (0..self.n_partitions)
+            .into_par_iter()
+            .map(|i| {
+                let dir = self.index.layer_dir(i, 0);
+                let dedup_path = obikindex::layer::dereplicated_superkmers_path(&dir);
+                let result = if dedup_path.exists() {
+                    count_partition(&dir, &dedup_path, chunk_kmers)
+                } else {
+                    Ok(())
+                };
+                if let Some(cb) = &on_progress {
+                    let pos = done.fetch_add(1, Ordering::Relaxed) + 1;
+                    cb(Progress { position: pos, total: Some(self.n_partitions as u64) });
+                }
+                result
+            })
+            .collect();
+
+        for r in results {
+            r?;
+        }
+
+        // Aggregate per-partition spectra.
+        let mut counts: BTreeMap<u32, u64> = BTreeMap::new();
+        let mut f0: u64 = 0;
+        let mut f1: u64 = 0;
+
+        for i in 0..self.n_partitions {
+            let path = self.index.layer_dir(i, 0).join("kmer_spectrum_raw.json");
+            if !path.exists() {
+                continue;
+            }
+            let v: serde_json::Value =
+                serde_json::from_str(&fs::read_to_string(&path)?).map_err(io::Error::other)?;
+            f0 += v["f0"].as_u64().unwrap_or(0);
+            f1 += v["f1"].as_u64().unwrap_or(0);
+            if let Some(obj) = v["spectrum"].as_object() {
+                for (c_str, freq) in obj {
+                    if let (Ok(c), Some(f)) = (c_str.parse::<u32>(), freq.as_u64()) {
+                        *counts.entry(c).or_insert(0) += f;
+                    }
+                }
+            }
+            if !self.keep_partial {
+                let _ = fs::remove_file(&path);
+            }
+        }
+
+        Ok(KmerSpectrum { f0, f1, counts })
+    }
+}

src/obikindexer/src/algorithms/mod.rs

@@ -2,8 +2,10 @@
 //! [`obikindex::KmerIndex`] to build or transform its content — `obikindex`
 //! itself is the `index`/`partition`/`layer` data model, not this. Each
 //! algorithm is its own submodule: [`partitionner`] (routing raw super-kmers
-//! into partitions, then counting), [`dereplicator`] (deduplicating a
-//! partition's raw super-kmers before counting).
+//! into partitions), [`dereplicator`] (deduplicating a partition's raw
+//! super-kmers), [`counter`] (counting unique canonical kmers from a
+//! partition's dereplicated super-kmers) — the three run in that order.
 
+pub mod counter;
 pub mod dereplicator;
 pub mod partitionner;

src/obikindexer/src/algorithms/partitionner/mod.rs

@@ -1,20 +1,16 @@
 //! K-mer partitioning: routing super-kmers into per-partition, layer-0
-//! files, and counting unique canonical k-mers. Dereplication of the raw
-//! super-kmers in between is [`crate::algorithms::dereplicator`], a
-//! sibling algorithm, not a step of this one — see
+//! files. Dereplication and counting of the raw super-kmers are separate
+//! sibling algorithms ([`crate::algorithms::dereplicator`],
+//! [`crate::algorithms::counter`]), not steps of this one — see
 //! `DevDocMD/implementation/partition_layer_cache.md`.
 //!
-//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
-//! routing/counting lifecycle API — partition/layer path naming itself
-//! lives on `obikindex::KmerIndex`, not here), [`count`] (unique-kmer
-//! enumeration, MPHF, abundance counting), `kmer_sort` (external sort
-//! support for `count`).
+//! Submodule: [`router`] (`PartitionRouter`, the routing lifecycle API —
+//! partition/layer path naming itself lives on `obikindex::KmerIndex`,
+//! not here).
 
-mod count;
-mod kmer_sort;
 mod router;
 
 #[cfg(test)]
 mod tests;
 
-pub use router::{KmerSpectrum, PartitionRouter};
+pub use router::PartitionRouter;

src/obikindexer/src/algorithms/partitionner/router.rs

@@ -1,5 +1,3 @@
-use std::collections::BTreeMap;
-use std::fs;
 use std::io;
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
@@ -10,9 +8,7 @@ use obikindex::KmerIndex;
 use obikseq::RoutableSuperKmer;
 use obikindex::layer::Layer;
 use obiskio::SKResult;
-use obisys::{progress_bar, Progress};
-use rayon::prelude::*;
-use sysinfo::System;
+use obisys::Progress;
 use tracing::info;
 
 use niffler::Level;
@@ -22,16 +18,6 @@ use obiskio::SKFileWriter;
 use obipipeline::{throttle, ThrottleGuard, Throttled};
 use obiread::NucPage;
 
-use super::kmer_sort::chunk_size_from_ram;
-
-use super::count::count_partition;
-
-pub struct KmerSpectrum {
-    pub f0: u64,
-    pub f1: u64,
-    pub counts: BTreeMap<u32, u64>,
-}
-
 // ── Pipeline plumbing, private to `run` ─────────────────────────────────────
 
 /// Carrier enum for `obipipeline::make_pipe!`'s two-stage transform — local
@@ -72,17 +58,17 @@ impl Drop for GuardedIter {
 
 // ── PartitionRouter ──────────────────────────────────────────────────────────
 
-/// Routes raw super-kmers into per-partition, layer-0 files, then
-/// dereplicates (via the sibling [`crate::algorithms::dereplicator`]
-/// algorithm) and counts them — the entry point of the indexing pipeline,
-/// operating on layer-0 content that doesn't exist yet at this stage (see
+/// Routes raw super-kmers into per-partition, layer-0 files — the entry
+/// point of the indexing pipeline. Dereplication
+/// ([`crate::algorithms::dereplicator`]) and counting
+/// ([`crate::algorithms::counter`]) are separate sibling algorithms that
+/// run after this one, on the layer-0 content it produces (see
 /// `DevDocMD/implementation/partition_layer_cache.md`).
 ///
 /// Holds `&mut KmerIndex` — this is an algorithm operating on an index, not
 /// a data structure of its own; it owns no path-naming knowledge (every
 /// path comes from `index.index_dir`/`obikindex::layer::layer_dir`/
-/// `Layer::create`), only the transient routing/dereplication/counting
-/// state a run needs.
+/// `Layer::create`), only the transient routing state a run needs.
 ///
 /// Two-phase construction: `new` + optional setters (`level_max`/`theta`/
 /// `workers`/`max_open`) configure the run, `run` executes it. `run` takes
@@ -92,7 +78,6 @@ impl Drop for GuardedIter {
 /// the caller's decision, not this crate's.
 pub struct PartitionRouter<'a> {
     index: &'a mut KmerIndex,
-    n_partitions: usize,
     partitions_mask: u64,
     writers: Vec<Option<SKFileWriter>>,
     level: Level,
@@ -113,7 +98,6 @@ impl<'a> PartitionRouter<'a> {
         let workers = obisys::effective_parallelism();
         Self {
             index,
-            n_partitions,
             partitions_mask: (1u64 << n_bits) - 1,
             writers: (0..n_partitions).map(|_| None).collect(),
             level: Level::One,
@@ -267,77 +251,6 @@ impl<'a> PartitionRouter<'a> {
         self.close()
     }
 
-    /// For each partition that has a `dereplicated.{ext}` file:
-    /// 1. Enumerates all unique canonical kmers (two passes over the file).
-    /// 2. Builds a provisional MPHF (FMPHGO) over those kmers.
-    /// 3. Writes a flat binary count file (`counts1.bin`, one `u32` per slot,
-    ///    memory-mapped) accumulating kmer abundances from the superkmer counts.
-    /// 4. Persists the MPHF to `mphf1.bin` for downstream use.
-    ///
-    /// Returns the aggregated `KmerSpectrum`. Per-partition spectrum files are
-    /// deleted after aggregation unless `keep_partial` is true.
-    pub fn count_kmer(&self, keep_partial: bool) -> SKResult<KmerSpectrum> {
-        let sys = System::new_all();
-        let available = match sys.available_memory() {
-            0 => sys.total_memory() / 2,
-            n => n,
-        };
-        let n_threads = rayon::current_num_threads().max(1) as u64;
-        let chunk_kmers = chunk_size_from_ram(available / n_threads);
-
-        let pb = progress_bar("counting", self.n_partitions as u64, "partitions");
-
-        let results: Vec<SKResult<()>> = (0..self.n_partitions)
-            .into_par_iter()
-            .map(|i| {
-                let dir = self.layer0_dir(i);
-                let dedup_path = obikindex::layer::dereplicated_superkmers_path(&dir);
-                if !dedup_path.exists() {
-                    pb.inc(1);
-                    return Ok(());
-                }
-                let t = Instant::now();
-                let result = count_partition(&dir, &dedup_path, chunk_kmers);
-                pb.set_message(format!("last {:.0}ms", t.elapsed().as_millis()));
-                pb.inc(1);
-                result
-            })
-            .collect();
-
-        pb.finish_and_clear();
-        for r in results {
-            r?;
-        }
-
-        // Aggregate per-partition spectra.
-        let mut counts: BTreeMap<u32, u64> = BTreeMap::new();
-        let mut f0: u64 = 0;
-        let mut f1: u64 = 0;
-
-        for i in 0..self.n_partitions {
-            let path = self.layer0_dir(i).join("kmer_spectrum_raw.json");
-            if !path.exists() {
-                continue;
-            }
-            let v: serde_json::Value =
-                serde_json::from_str(&fs::read_to_string(&path)?).map_err(io::Error::other)?;
-            f0 += v["f0"].as_u64().unwrap_or(0);
-            f1 += v["f1"].as_u64().unwrap_or(0);
-            if let Some(obj) = v["spectrum"].as_object() {
-                for (c_str, freq) in obj {
-                    if let (Ok(c), Some(f)) = (c_str.parse::<u32>(), freq.as_u64()) {
-                        *counts.entry(c).or_insert(0) += f;
-                    }
-                }
-            }
-            if !keep_partial {
-                let _ = fs::remove_file(&path);
-            }
-        }
-
-        Ok(KmerSpectrum { f0, f1, counts })
-    }
-
     // ── private ───────────────────────────────────────────────────────────────
 
     /// Directory of partition `i`'s layer 0 — every raw/dereplicated

src/obikindexer/src/algorithms/partitionner/tests.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
-use std::fs;
 
+use crate::algorithms::counter::Counter;
 use crate::algorithms::dereplicator::Dereplicator;
 use obikindex::{IndexConfig, KmerIndex};
 use obikindex::layer::IndexMode;
@@ -8,7 +8,6 @@ use obikrope::Rope;
 use obikseq::SuperKmer;
 use obiskbuilder::build_superkmers;
 
-use super::count::count_partition;
 use super::PartitionRouter;
 
 const K: usize = 11;
@@ -46,7 +45,7 @@ fn direct_counts(seqs: &[&[u8]]) -> (u64, u64) {
 }
 
 /// Run the full pipeline on a list of sequences and return (f0, f1) from
-/// the `kmer_spectrum_raw.json` produced by `count_partition`.
+/// the aggregated `KmerSpectrum` produced by `Counter::run`.
 fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
     setup();
 
@@ -73,14 +72,8 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
     if !dedup_path.exists() {
         return (0, 0);
     }
-    count_partition(&part_dir, &dedup_path, 1 << 20).unwrap();
-
-    let spec: serde_json::Value =
-        serde_json::from_reader(fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap())
-            .unwrap();
-    let f0 = spec["f0"].as_u64().unwrap_or(0);
-    let f1 = spec["f1"].as_u64().unwrap_or(0);
-    (f0, f1)
+    let spectrum = Counter::new(&index).run(None::<fn(obisys::Progress)>).unwrap();
+    (spectrum.f0, spectrum.f1)
 }
 
 #[test]

src/obikmer/src/cmd/index/mod.rs

@@ -2,6 +2,7 @@ use std::path::PathBuf;
 use std::time::Instant;
 
 use clap::Args;
+use obikindexer::algorithms::counter::Counter;
 use obikindexer::algorithms::dereplicator::Dereplicator;
 use obikindexer::algorithms::partitionner::PartitionRouter;
 use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
@@ -301,15 +302,17 @@ pub fn run(args: IndexArgs) {
         pb.finish_and_clear();
         rep.push(t.stop());
 
-        let router = PartitionRouter::new(&mut idx);
-
         let t = Stage::start("count_kmer");
-        let spectrum = router.count_kmer(args.keep_intermediate).unwrap_or_else(|e| {
-            eprintln!("error: {e}");
-            std::process::exit(1);
-        });
+        let pb = progress_bar("counting", idx.n_partitions() as u64, "partitions");
+        let spectrum = Counter::new(&idx)
+            .keep_partial(args.keep_intermediate)
+            .run(Some(|_: Progress| pb.inc(1)))
+            .unwrap_or_else(|e| {
+                eprintln!("error: {e}");
+                std::process::exit(1);
+            });
+        pb.finish_and_clear();
         rep.push(t.stop());
-        drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
 
         idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).unwrap_or_else(|e| {
             eprintln!("error: {e}");

src/obikphylo/src/siblings/tests.rs

@@ -6,6 +6,7 @@ use obikindex::layer::MphfLayer;
 use obisys::Reporter;
 use tempfile::tempdir;
 
+use obikindexer::algorithms::counter::Counter;
 use obikindexer::algorithms::dereplicator::Dereplicator;
 use obikindexer::algorithms::partitionner::PartitionRouter;
 use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
@@ -80,9 +81,7 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
     drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
 
     Dereplicator::new(&idx).run(None::<fn(obisys::Progress)>).expect("dereplicate");
-    let router = PartitionRouter::new(&mut idx);
-    let spectrum = router.count_kmer(false).expect("count_kmer");
-    drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
+    let spectrum = Counter::new(&idx).run(None::<fn(obisys::Progress)>).expect("count_kmer");
 
     idx.mark_scattered().expect("mark_scattered");
     idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum");