Add modular data structures, parallel pipelines, and system profiling

Establishes foundational infrastructure across multiple crates by introducing unified persistent bit matrix storage with columnar, packed, and implicit variants, alongside De Bruijn graph node encoding and unitig iteration logic. Adds a macro-driven parallel pipeline scheduler featuring NUMA-aware runners, bounded channels, and memory budgets to enforce concurrency limits. Implements streaming nucleotide parsers with pooled page buffers for FASTA, FASTQ, and Genbank formats, complemented by system resource monitoring, progress tracking, and stage profiling utilities. Collectively, these changes provide the core data models, execution frameworks, and I/O pipelines required for downstream k-mer indexing and analysis workloads.
This commit is contained in:
Eric Coissac
2026-08-14 14:20:08 +02:00
parent cc67023e2c
commit 79346c0c86
69 changed files with 7487 additions and 7060 deletions
@@ -1,250 +1,24 @@
//use ahash::RandomState;
use crossbeam_channel;
use hashbrown::HashMap;
use obikseq::k;
use obikseq::{CanonicalKmer, Sequence, Unitig};
use obikseq::{CanonicalKmer, Unitig};
#[cfg(not(any(test, feature = "test-utils")))]
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use std::cell::RefCell;
use std::fmt;
use std::sync::atomic::{AtomicU8, Ordering};
use xxhash_rust::xxh3::Xxh3Builder;
use super::node::{IS_VISITED_MASK, Node};
use super::unitig_iter::UnitigNucIter;
use super::walk::WalkState;
// ── Types ─────────────────────────────────────────────────────────────────────
type FastHashMap<K, V> = HashMap<K, V, Xxh3Builder>;
// ── Node ──────────────────────────────────────────────────────────────────────
//
// bit layout (LSB first):
// bit 0 : can_extend_right — exactly one right canonical neighbour exists
// bit 1 : can_extend_left — exactly one left canonical neighbour exists
// bit 2 : visited
// bits 34 : right_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 0 = 1
// bits 56 : left_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 1 = 1
// bit 7 : marked as start node (1)
//
// "can_extend" = false covers both 0 neighbours and ≥2 neighbours; the only
// information needed for traversal is "exactly one".
#[repr(transparent)]
#[derive(Debug, Clone, Copy, Default)]
pub struct Node(u8);
const CAN_EXTEND_RIGHT_MASK: u8 = 0b0000_0001; // bit 0: can_extend_right — exactly one right canonical neighbour exists
const CAN_EXTEND_LEFT_MASK: u8 = 0b0000_0010; // bit 1: can_extend_left — exactly one left canonical neighbour exists
const IS_VISITED_MASK: u8 = 0b0000_0100; // bit 2: visited
const RIGHT_NUC_MASK: u8 = 0b0001_1000; // bits 34: right_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 0 = 1
const LEFT_NUC_MASK: u8 = 0b0110_0000; // bits 56: left_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 1 = 1
const IS_START_MASK: u8 = 0b1000_0000; // bit 7: marked as start node
impl Node {
/// Returns `true` if the node can be extended to the right.
///
/// A single right neighbour exists.
#[inline]
pub fn can_extend_right(self) -> bool {
self.0 & CAN_EXTEND_RIGHT_MASK != 0
}
/// Returns `true` if the node can be extended to the left.
///
/// A single left neighbour exists.
#[inline]
pub fn can_extend_left(self) -> bool {
self.0 & CAN_EXTEND_LEFT_MASK != 0
}
/// Returns `true` if the node has been visited.
#[inline]
pub fn is_visited(self) -> bool {
self.0 & IS_VISITED_MASK != 0
}
/// Returns `true` if the node is a start node.
#[inline]
pub fn is_start(self) -> bool {
self.0 & IS_START_MASK != 0
}
#[inline]
pub fn set_start(&mut self) {
self.0 |= IS_START_MASK;
}
pub fn unset_start(&mut self) {
self.0 &= !IS_START_MASK;
}
/// Index of the unique right neighbour (0=A, 1=C, 2=G, 3=T).
/// Only meaningful when `can_extend_right()` is true.
#[inline]
pub fn right_nuc(self) -> u8 {
debug_assert!(
self.can_extend_right(),
"from: right_nuc -> The node cannot be extended to the right"
);
(self.0 >> 3) & 0b11
}
/// Index of the unique left neighbour (0=A, 1=C, 2=G, 3=T).
/// Only meaningful when `can_extend_left()` is true.
#[inline]
pub fn left_nuc(self) -> u8 {
debug_assert!(
self.can_extend_left(),
"from: left_nuc -> The node cannot be extended to the left"
);
(self.0 >> 5) & 0b11
}
/// Marks the node as visited.
#[inline]
pub fn set_visited(&mut self) {
debug_assert!(
!self.is_visited(),
"from: is_visited -> The node has already been visited"
);
self.0 |= IS_VISITED_MASK;
}
/// `nuc` = Some(i) → exactly one neighbour (bit 0 set, bits 34 = nucleotide index).
/// `nuc` = None → 0 or ≥2 neighbours; `count` encoded in bits 34 as count.sat_sub(1).
pub fn set_right(&mut self, count: u8, nuc: Option<u8>) {
self.0 &= !(CAN_EXTEND_RIGHT_MASK | RIGHT_NUC_MASK);
if count == 1 {
self.0 |= CAN_EXTEND_RIGHT_MASK;
if let Some(n) = nuc {
self.0 |= (n & 0b11) << 3;
return;
}
unreachable!("nuc must be Some when count is 1");
}
self.0 |= (count.saturating_sub(1).min(3)) << 3;
}
/// `nuc` = Some(i) → exactly one neighbour (bit 0 set, bits 34 = nucleotide index).
/// `nuc` = None → 0 or ≥2 neighbours; `count` encoded in bits 34 as count.sat_sub(1).
pub fn set_left(&mut self, count: u8, nuc: Option<u8>) {
self.0 &= !(CAN_EXTEND_LEFT_MASK | LEFT_NUC_MASK);
if count == 1 {
self.0 |= CAN_EXTEND_LEFT_MASK;
if let Some(n) = nuc {
self.0 |= (n & 0b11) << 5;
return;
}
unreachable!("nuc must be Some when count is 1");
}
self.0 |= (count.saturating_sub(1).min(3)) << 5;
}
}
impl fmt::Display for Node {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
const NUC: [char; 4] = ['A', 'C', 'G', 'T'];
let r = if self.can_extend_right() {
format!("{}", NUC[self.right_nuc() as usize])
} else if (self.0 >> 3) & 0b11 == 0 {
"→0".to_string()
} else {
"→≥2".to_string()
};
let l = if self.can_extend_left() {
format!("{}", NUC[self.left_nuc() as usize])
} else if (self.0 >> 5) & 0b11 == 0 {
"←0".to_string()
} else {
"←≥2".to_string()
};
let v = if self.is_visited() { "V" } else { "." };
write!(f, "Node({r} {l} {v})")
}
}
pub struct WalkState {
kmer: CanonicalKmer,
node: Node,
direct: bool,
}
impl WalkState {
pub fn new(kmer: CanonicalKmer, node: Node, direct: bool) -> Self {
debug_assert!(!node.is_visited(), "Cannot walk over a visited node");
Self { kmer, node, direct }
}
pub fn leavable(&self, graph: &GraphDeBruijn) -> bool {
self.walk(graph).is_some()
}
pub fn reachable(&self, graph: &GraphDeBruijn) -> bool {
WalkState {
kmer: self.kmer,
node: self.node,
direct: !self.direct,
}
.leavable(graph)
}
pub fn walk(&self, graph: &GraphDeBruijn) -> Option<(WalkState, u8)> {
if self.direct {
if !self.node.can_extend_right() {
return None;
}
let nuc = self.node.right_nuc();
let next = self.kmer.into_kmer().push_right(nuc);
let cnext = next.canonical();
let dnext = next.raw() == cnext.raw();
let next_node = Node(graph.nodes.get(&cnext).unwrap().load(Ordering::Relaxed));
if next_node.is_visited() {
return None;
}
let reachable = if dnext {
next_node.can_extend_left()
} else {
next_node.can_extend_right()
};
reachable.then_some((
WalkState {
kmer: cnext,
node: next_node,
direct: dnext,
},
nuc,
))
} else {
if !self.node.can_extend_left() {
return None;
}
let nuc = self.node.left_nuc();
let next = self.kmer.into_kmer().push_left(nuc);
let cnext = next.canonical();
let dnext = next.raw() != cnext.raw();
let next_node = Node(graph.nodes.get(&cnext).unwrap().load(Ordering::Relaxed));
if next_node.is_visited() {
return None;
}
let reachable = if dnext {
next_node.can_extend_right()
} else {
next_node.can_extend_left()
};
reachable.then_some((
WalkState {
kmer: cnext,
node: next_node,
direct: dnext,
},
3 - nuc,
))
}
}
}
pub(super) type FastHashMap<K, V> = HashMap<K, V, Xxh3Builder>;
// ── GraphDeBruijn ─────────────────────────────────────────────────────────────
pub struct GraphDeBruijn {
nodes: FastHashMap<CanonicalKmer, AtomicU8>,
pub(super) nodes: FastHashMap<CanonicalKmer, AtomicU8>,
}
impl GraphDeBruijn {
@@ -346,7 +120,7 @@ impl GraphDeBruijn {
Some(WalkState::new(kmer, node, true))
}
fn unitig_nucleotides(&self, kmer: CanonicalKmer, k: usize) -> Option<UnitigNucIter<'_>> {
pub(super) fn unitig_nucleotides(&self, kmer: CanonicalKmer, k: usize) -> Option<UnitigNucIter<'_>> {
let old = self
.nodes
.get(&kmer)?
@@ -362,13 +136,7 @@ impl GraphDeBruijn {
.fetch_or(IS_VISITED_MASK, Ordering::AcqRel);
(ext_old & IS_VISITED_MASK == 0).then_some((next_state, nuc))
});
Some(UnitigNucIter {
graph: self,
start: kmer,
pos: 0,
k,
next_step,
})
Some(UnitigNucIter::new(self, kmer, k, next_step))
}
pub fn for_each_unitig(&self, f: impl Fn(UnitigNucIter<'_>) + Sync) {
@@ -467,12 +235,7 @@ impl GraphDeBruijn {
}
fn is_start(&self, query: CanonicalKmer, node: Node) -> bool {
!WalkState {
kmer: query,
node,
direct: true,
}
.reachable(self)
!WalkState::new(query, node, true).reachable(self)
}
pub fn try_for_each_unitig<E, F>(&self, f: F) -> Result<(), E>
@@ -514,44 +277,6 @@ impl GraphDeBruijn {
}
}
// ── UnitigNucIter ─────────────────────────────────────────────────────────────
pub struct UnitigNucIter<'a> {
graph: &'a GraphDeBruijn,
start: CanonicalKmer,
pos: usize,
k: usize,
next_step: Option<(WalkState, u8)>,
}
impl Iterator for UnitigNucIter<'_> {
type Item = u8;
fn next(&mut self) -> Option<u8> {
if self.pos < self.k {
let nuc = self.start.nucleotide(self.pos);
self.pos += 1;
Some(nuc)
} else if let Some((state, nuc)) = self.next_step.take() {
self.next_step = state.walk(self.graph).and_then(|(next_state, next_nuc)| {
let old = self
.graph
.nodes
.get(&next_state.kmer)?
.fetch_or(IS_VISITED_MASK, Ordering::AcqRel);
(old & IS_VISITED_MASK == 0).then_some((next_state, next_nuc))
});
Some(nuc)
} else {
None
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.k - self.pos.min(self.k), None)
}
}
/// Returns the count of neighbors and the index of the first
/// neighbor if exactly one of the four canonical neighbours exists in
/// the graph, where `i` is its index (0=A, 1=C, 2=G, 3=T).
@@ -580,8 +305,3 @@ fn count_neighbors(
(0, None)
}
}
// ── tests ─────────────────────────────────────────────────────────────────────
#[cfg(test)]
#[path = "tests/debruijn.rs"]
mod tests;
+20
View File
@@ -0,0 +1,20 @@
//! De Bruijn graph over canonical k-mers, built for unitig extraction.
//!
//! Submodules: [`node`] (packed per-kmer neighbour/visited/start flags),
//! [`walk`] (single-step traversal), [`graph`] ([`GraphDeBruijn`] itself),
//! [`unitig_iter`] (nucleotide-by-nucleotide unitig walk iterator).
mod graph;
mod node;
mod unitig_iter;
mod walk;
pub use graph::GraphDeBruijn;
// Only used by `tests/debruijn.rs` (`use super::*`) below.
#[cfg(test)]
use obikseq::{CanonicalKmer, Sequence};
#[cfg(test)]
#[path = "../tests/debruijn.rs"]
mod tests;
+148
View File
@@ -0,0 +1,148 @@
use std::fmt;
// ── Node ──────────────────────────────────────────────────────────────────────
//
// bit layout (LSB first):
// bit 0 : can_extend_right — exactly one right canonical neighbour exists
// bit 1 : can_extend_left — exactly one left canonical neighbour exists
// bit 2 : visited
// bits 34 : right_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 0 = 1
// bits 56 : left_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 1 = 1
// bit 7 : marked as start node (1)
//
// "can_extend" = false covers both 0 neighbours and ≥2 neighbours; the only
// information needed for traversal is "exactly one".
#[repr(transparent)]
#[derive(Debug, Clone, Copy, Default)]
pub struct Node(pub(super) u8);
const CAN_EXTEND_RIGHT_MASK: u8 = 0b0000_0001; // bit 0: can_extend_right — exactly one right canonical neighbour exists
const CAN_EXTEND_LEFT_MASK: u8 = 0b0000_0010; // bit 1: can_extend_left — exactly one left canonical neighbour exists
pub(super) const IS_VISITED_MASK: u8 = 0b0000_0100; // bit 2: visited
const RIGHT_NUC_MASK: u8 = 0b0001_1000; // bits 34: right_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 0 = 1
const LEFT_NUC_MASK: u8 = 0b0110_0000; // bits 56: left_nuc — index 03 (A/C/G/T) of that neighbour; valid iff bit 1 = 1
const IS_START_MASK: u8 = 0b1000_0000; // bit 7: marked as start node
impl Node {
/// Returns `true` if the node can be extended to the right.
///
/// A single right neighbour exists.
#[inline]
pub fn can_extend_right(self) -> bool {
self.0 & CAN_EXTEND_RIGHT_MASK != 0
}
/// Returns `true` if the node can be extended to the left.
///
/// A single left neighbour exists.
#[inline]
pub fn can_extend_left(self) -> bool {
self.0 & CAN_EXTEND_LEFT_MASK != 0
}
/// Returns `true` if the node has been visited.
#[inline]
pub fn is_visited(self) -> bool {
self.0 & IS_VISITED_MASK != 0
}
/// Returns `true` if the node is a start node.
#[inline]
pub fn is_start(self) -> bool {
self.0 & IS_START_MASK != 0
}
#[inline]
pub fn set_start(&mut self) {
self.0 |= IS_START_MASK;
}
pub fn unset_start(&mut self) {
self.0 &= !IS_START_MASK;
}
/// Index of the unique right neighbour (0=A, 1=C, 2=G, 3=T).
/// Only meaningful when `can_extend_right()` is true.
#[inline]
pub fn right_nuc(self) -> u8 {
debug_assert!(
self.can_extend_right(),
"from: right_nuc -> The node cannot be extended to the right"
);
(self.0 >> 3) & 0b11
}
/// Index of the unique left neighbour (0=A, 1=C, 2=G, 3=T).
/// Only meaningful when `can_extend_left()` is true.
#[inline]
pub fn left_nuc(self) -> u8 {
debug_assert!(
self.can_extend_left(),
"from: left_nuc -> The node cannot be extended to the left"
);
(self.0 >> 5) & 0b11
}
/// Marks the node as visited.
#[inline]
pub fn set_visited(&mut self) {
debug_assert!(
!self.is_visited(),
"from: is_visited -> The node has already been visited"
);
self.0 |= IS_VISITED_MASK;
}
/// `nuc` = Some(i) → exactly one neighbour (bit 0 set, bits 34 = nucleotide index).
/// `nuc` = None → 0 or ≥2 neighbours; `count` encoded in bits 34 as count.sat_sub(1).
pub fn set_right(&mut self, count: u8, nuc: Option<u8>) {
self.0 &= !(CAN_EXTEND_RIGHT_MASK | RIGHT_NUC_MASK);
if count == 1 {
self.0 |= CAN_EXTEND_RIGHT_MASK;
if let Some(n) = nuc {
self.0 |= (n & 0b11) << 3;
return;
}
unreachable!("nuc must be Some when count is 1");
}
self.0 |= (count.saturating_sub(1).min(3)) << 3;
}
/// `nuc` = Some(i) → exactly one neighbour (bit 0 set, bits 34 = nucleotide index).
/// `nuc` = None → 0 or ≥2 neighbours; `count` encoded in bits 34 as count.sat_sub(1).
pub fn set_left(&mut self, count: u8, nuc: Option<u8>) {
self.0 &= !(CAN_EXTEND_LEFT_MASK | LEFT_NUC_MASK);
if count == 1 {
self.0 |= CAN_EXTEND_LEFT_MASK;
if let Some(n) = nuc {
self.0 |= (n & 0b11) << 5;
return;
}
unreachable!("nuc must be Some when count is 1");
}
self.0 |= (count.saturating_sub(1).min(3)) << 5;
}
}
impl fmt::Display for Node {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
const NUC: [char; 4] = ['A', 'C', 'G', 'T'];
let r = if self.can_extend_right() {
format!("{}", NUC[self.right_nuc() as usize])
} else if (self.0 >> 3) & 0b11 == 0 {
"→0".to_string()
} else {
"→≥2".to_string()
};
let l = if self.can_extend_left() {
format!("{}", NUC[self.left_nuc() as usize])
} else if (self.0 >> 5) & 0b11 == 0 {
"←0".to_string()
} else {
"←≥2".to_string()
};
let v = if self.is_visited() { "V" } else { "." };
write!(f, "Node({r} {l} {v})")
}
}
@@ -0,0 +1,61 @@
use obikseq::CanonicalKmer;
use std::sync::atomic::Ordering;
use super::graph::GraphDeBruijn;
use super::node::IS_VISITED_MASK;
use super::walk::WalkState;
// ── UnitigNucIter ─────────────────────────────────────────────────────────────
pub struct UnitigNucIter<'a> {
graph: &'a GraphDeBruijn,
start: CanonicalKmer,
pos: usize,
k: usize,
next_step: Option<(WalkState, u8)>,
}
impl<'a> UnitigNucIter<'a> {
pub(super) fn new(
graph: &'a GraphDeBruijn,
start: CanonicalKmer,
k: usize,
next_step: Option<(WalkState, u8)>,
) -> Self {
Self {
graph,
start,
pos: 0,
k,
next_step,
}
}
}
impl Iterator for UnitigNucIter<'_> {
type Item = u8;
fn next(&mut self) -> Option<u8> {
if self.pos < self.k {
let nuc = self.start.nucleotide(self.pos);
self.pos += 1;
Some(nuc)
} else if let Some((state, nuc)) = self.next_step.take() {
self.next_step = state.walk(self.graph).and_then(|(next_state, next_nuc)| {
let old = self
.graph
.nodes
.get(&next_state.kmer)?
.fetch_or(IS_VISITED_MASK, Ordering::AcqRel);
(old & IS_VISITED_MASK == 0).then_some((next_state, next_nuc))
});
Some(nuc)
} else {
None
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.k - self.pos.min(self.k), None)
}
}
+85
View File
@@ -0,0 +1,85 @@
use obikseq::{CanonicalKmer, Sequence};
use std::sync::atomic::Ordering;
use super::graph::GraphDeBruijn;
use super::node::Node;
pub struct WalkState {
pub(super) kmer: CanonicalKmer,
pub(super) node: Node,
pub(super) direct: bool,
}
impl WalkState {
pub fn new(kmer: CanonicalKmer, node: Node, direct: bool) -> Self {
debug_assert!(!node.is_visited(), "Cannot walk over a visited node");
Self { kmer, node, direct }
}
pub fn leavable(&self, graph: &GraphDeBruijn) -> bool {
self.walk(graph).is_some()
}
pub fn reachable(&self, graph: &GraphDeBruijn) -> bool {
WalkState {
kmer: self.kmer,
node: self.node,
direct: !self.direct,
}
.leavable(graph)
}
pub fn walk(&self, graph: &GraphDeBruijn) -> Option<(WalkState, u8)> {
if self.direct {
if !self.node.can_extend_right() {
return None;
}
let nuc = self.node.right_nuc();
let next = self.kmer.into_kmer().push_right(nuc);
let cnext = next.canonical();
let dnext = next.raw() == cnext.raw();
let next_node = Node(graph.nodes.get(&cnext).unwrap().load(Ordering::Relaxed));
if next_node.is_visited() {
return None;
}
let reachable = if dnext {
next_node.can_extend_left()
} else {
next_node.can_extend_right()
};
reachable.then_some((
WalkState {
kmer: cnext,
node: next_node,
direct: dnext,
},
nuc,
))
} else {
if !self.node.can_extend_left() {
return None;
}
let nuc = self.node.left_nuc();
let next = self.kmer.into_kmer().push_left(nuc);
let cnext = next.canonical();
let dnext = next.raw() != cnext.raw();
let next_node = Node(graph.nodes.get(&cnext).unwrap().load(Ordering::Relaxed));
if next_node.is_visited() {
return None;
}
let reachable = if dnext {
next_node.can_extend_right()
} else {
next_node.can_extend_left()
};
reachable.then_some((
WalkState {
kmer: cnext,
node: next_node,
direct: dnext,
},
3 - nuc,
))
}
}
}