feat: Add parallel execution and thread-safe graph operations

Integrate rayon to enable parallel processing of k-mer partitions and degree computation. Replace Cell with AtomicU8 to ensure thread-safe node state management, and add a merge method for combining disjoint graphs. Additionally, introduce progress tracking utilities and a test-utils feature flag for development dependencies.
This commit is contained in:
Eric Coissac
2026-06-04 22:40:21 +02:00
parent edd5e3f8ee
commit 2f29ee2240
5 changed files with 85 additions and 39 deletions
+5 -1
View File
@@ -7,8 +7,12 @@ edition = "2021"
obikseq = { path = "../obikseq" }
obifastwrite = { path = "../obifastwrite" }
ahash = "0.8"
hashbrown = "0.14"
hashbrown = { version = "0.14", features = ["rayon"] }
rayon = "1"
xxhash-rust = { version = "0.8.15", features = ["xxh3", "const_xxh3"] }
[features]
test-utils = []
[dev-dependencies]
obikseq = { path = "../obikseq", features = ["test-utils"] }
+45 -23
View File
@@ -3,8 +3,9 @@ use hashbrown::HashMap;
use obikseq::k;
use obikseq::unitig::Unitig;
use obikseq::{CanonicalKmer, Kmer, Sequence};
use std::cell::Cell;
use rayon::prelude::*;
use std::fmt;
use std::sync::atomic::{AtomicU8, Ordering};
use xxhash_rust::xxh3::Xxh3Builder;
// ── Types ─────────────────────────────────────────────────────────────────────
@@ -140,7 +141,7 @@ impl fmt::Display for Node {
// ── GraphDeBruijn ─────────────────────────────────────────────────────────────
pub struct GraphDeBruijn {
nodes: FastHashMap<CanonicalKmer, Cell<Node>>,
nodes: FastHashMap<CanonicalKmer, AtomicU8>,
}
impl GraphDeBruijn {
@@ -158,24 +159,34 @@ impl GraphDeBruijn {
/// Insert a canonical kmer into the graph. No-op if already present.
pub fn push(&mut self, kmer: CanonicalKmer) {
self.nodes
.entry(kmer)
.or_insert_with(|| Cell::new(Node::default()));
self.nodes.entry(kmer).or_insert_with(|| AtomicU8::new(0));
}
/// For every node, find its unique right/left canonical neighbour (if any)
/// and store the nucleotide index in the Node flags.
///
/// Single pass thanks to Cell interior mutability.
/// In production builds, runs in parallel across all nodes (each entry is
/// written by exactly one thread). In test builds, runs sequentially to
/// avoid propagating thread-local k/m values to rayon worker threads.
pub fn compute_degrees(&self) {
for (&kmer, cell) in &self.nodes {
#[cfg(not(any(test, feature = "test-utils")))]
self.nodes.par_iter().for_each(|(&kmer, atomic)| {
let (rc, rn) = count_neighbors(kmer.right_canonical_neighbors(), &self.nodes);
let (lc, ln) = count_neighbors(kmer.left_canonical_neighbors(), &self.nodes);
let mut node = cell.get();
let mut node = Node(atomic.load(Ordering::Relaxed));
node.set_right(rc, rn);
node.set_left(lc, ln);
cell.set(node);
atomic.store(node.0, Ordering::Relaxed);
});
#[cfg(any(test, feature = "test-utils"))]
for (&kmer, atomic) in &self.nodes {
let (rc, rn) = count_neighbors(kmer.right_canonical_neighbors(), &self.nodes);
let (lc, ln) = count_neighbors(kmer.left_canonical_neighbors(), &self.nodes);
let mut node = Node(atomic.load(Ordering::Relaxed));
node.set_right(rc, rn);
node.set_left(lc, ln);
atomic.store(node.0, Ordering::Relaxed);
}
}
@@ -206,20 +217,20 @@ impl GraphDeBruijn {
}
pub fn is_visited(&self, kmer: &CanonicalKmer) -> Option<bool> {
self.nodes.get(kmer).map(|cell| cell.get().is_visited())
self.nodes.get(kmer).map(|a| Node(a.load(Ordering::Relaxed)).is_visited())
}
pub fn set_visited(&self, kmer: CanonicalKmer) {
if let Some(cell) = self.nodes.get(&kmer) {
let mut node = cell.get();
if let Some(a) = self.nodes.get(&kmer) {
let mut node = Node(a.load(Ordering::Relaxed));
node.set_visited();
cell.set(node);
a.store(node.0, Ordering::Relaxed);
}
}
/// Returns the single right neighbor of `kmer`, if it exists.
pub fn the_single_right_neighbor(&self, kmer: CanonicalKmer) -> Option<CanonicalKmer> {
let node = self.nodes.get(&kmer)?.get();
let node = Node(self.nodes.get(&kmer)?.load(Ordering::Relaxed));
if !node.can_extend_right() {
return None;
}
@@ -229,7 +240,7 @@ impl GraphDeBruijn {
/// Returns the single left neighbor of `kmer`, if it exists.
pub fn the_single_left_neighbor(&self, kmer: CanonicalKmer) -> Option<CanonicalKmer> {
let node = self.nodes.get(&kmer)?.get();
let node = Node(self.nodes.get(&kmer)?.load(Ordering::Relaxed));
if !node.can_extend_left() {
return None;
}
@@ -253,7 +264,7 @@ impl GraphDeBruijn {
fn next_unitig_kmer(&self, kmer: Kmer) -> Option<Kmer> {
let canonical = kmer.canonical();
let node = self.nodes.get(&canonical)?.get();
let node = Node(self.nodes.get(&canonical)?.load(Ordering::Relaxed));
let direct = kmer.raw() == canonical.raw();
@@ -270,8 +281,8 @@ impl GraphDeBruijn {
canonical.into_kmer().push_left(node.left_nuc()).canonical()
};
let cell = self.nodes.get(&next_c)?;
let next_node = cell.get();
let atomic = self.nodes.get(&next_c)?;
let next_node = Node(atomic.load(Ordering::Relaxed));
if next_node.is_visited() {
return None;
}
@@ -287,7 +298,7 @@ impl GraphDeBruijn {
let mut updated = next_node;
updated.set_visited();
cell.set(updated);
atomic.store(updated.0, Ordering::Relaxed);
Some(oriented)
}
@@ -311,6 +322,17 @@ impl GraphDeBruijn {
})
}
/// Merge `other` into `self`.
///
/// The caller guarantees that the two graphs have **disjoint** kmer sets —
/// this is structurally true when each graph was built from a distinct
/// minimizer partition. No conflict resolution is needed: entries are moved
/// directly without checking for duplicates.
pub fn merge(&mut self, other: Self) {
self.nodes.reserve(other.nodes.len());
self.nodes.extend(other.nodes);
}
pub fn len(&self) -> usize {
self.nodes.len()
}
@@ -323,7 +345,7 @@ impl GraphDeBruijn {
// --- StartIter -----------------------------------------------------------------
struct StartIter<'a> {
graph: &'a GraphDeBruijn,
nodes: hashbrown::hash_map::Iter<'a, CanonicalKmer, Cell<Node>>,
nodes: hashbrown::hash_map::Iter<'a, CanonicalKmer, AtomicU8>,
suspended: Vec<CanonicalKmer>,
in_cycle_pass: bool,
}
@@ -364,7 +386,7 @@ impl<'a> Iterator for StartIter<'a> {
};
let node = match self.graph.nodes.get(&current) {
Some(c) => c.get(),
Some(a) => Node(a.load(Ordering::Relaxed)),
None => continue,
};
if node.is_visited() {
@@ -439,7 +461,7 @@ fn oriented_next(from: Kmer, to: CanonicalKmer) -> Kmer {
/// zero or ≥2 existing neighbours.
fn count_neighbors(
neighbors: [CanonicalKmer; 4],
nodes: &FastHashMap<CanonicalKmer, Cell<Node>>,
nodes: &FastHashMap<CanonicalKmer, AtomicU8>,
) -> (u8, Option<u8>) {
let mut count = 0u8;
let mut first = None;
+1 -1
View File
@@ -115,7 +115,7 @@ fn unitig_roundtrip_linear() {
g.compute_degrees();
println!("Les kmers:");
for (kmer, v) in g.nodes.iter() {
println!("{}: {}", String::from_utf8_lossy(&kmer.to_ascii()), v.get());
println!("{}: {}", String::from_utf8_lossy(&kmer.to_ascii()), v.load(std::sync::atomic::Ordering::Relaxed));
}
println!("Les unitig:");