src/obidebruinj/src/debruijn.rs

@@ -3,6 +3,7 @@ use hashbrown::HashMap;
 use obikseq::k;
 use obikseq::{CanonicalKmer, Sequence};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use std::cell::RefCell;
 use std::fmt;
 use std::sync::atomic::{AtomicU8, Ordering};
 use xxhash_rust::xxh3::Xxh3Builder;
@@ -450,17 +451,25 @@ impl GraphDeBruijn {
     pub fn try_for_each_unitig<E, F>(&self, f: F) -> Result<(), E>
     where
         E: Send,
-        F: FnMut(UnitigNucIter<'_>) -> Result<(), E> + Send,
+        F: FnMut(&[u8]) -> Result<(), E> + Send,
     {
+        thread_local! {
+            static BUF: std::cell::RefCell<Vec<u8>> = RefCell::new(Vec::with_capacity(4096));
+        }
         let error = std::sync::Mutex::new(None::<E>);
         let f = std::sync::Mutex::new(f);
         self.for_each_unitig(|iter| {
             if error.lock().unwrap().is_some() {
                 return;
             }
-            if let Err(e) = f.lock().unwrap()(iter) {
-                *error.lock().unwrap() = Some(e);
-            }
+            BUF.with(|buf| {
+                let mut buf = buf.borrow_mut();
+                buf.clear();
+                buf.extend(iter);
+                if let Err(e) = f.lock().unwrap()(&buf) {
+                    *error.lock().unwrap() = Some(e);
+                }
+            });
         });
         error.into_inner().unwrap().map_or(Ok(()), Err)
     }

src/obidebruinj/src/tests/debruijn.rs

@@ -24,8 +24,8 @@ fn canonical_kmers(seq: &[u8]) -> Vec<CanonicalKmer> {
 
 fn collect_unitigs(g: &GraphDeBruijn) -> Vec<Unitig> {
     let mut unitigs = Vec::new();
-    g.try_for_each_unitig(|nuc_iter| -> Result<(), std::convert::Infallible> {
-        unitigs.push(nuc_iter.collect());
+    g.try_for_each_unitig(|nucs| -> Result<(), std::convert::Infallible> {
+        unitigs.push(Unitig::from_nucleotides(nucs));
         Ok(())
     })
     .unwrap();

src/obikpartitionner/src/index_layer.rs

@@ -107,9 +107,8 @@ impl KmerPartition {
         fs::create_dir_all(&layer_dir)?;
 
         let mut uw = Layer::<()>::unitig_writer(&layer_dir).map_err(olm_to_sk)?;
-        g.try_for_each_unitig(|nuc_iter| {
-            let unitig: obikseq::unitig::Unitig = nuc_iter.collect();
-            uw.write(&unitig)
+        g.try_for_each_unitig(|nucs| {
+            uw.write(&obikseq::unitig::Unitig::from_nucleotides(nucs))
         })?;
         uw.close()?;
 

src/obikpartitionner/src/merge_layer.rs

@@ -307,9 +307,8 @@ impl KmerPartition {
             g.compute_degrees_and_mark_starts();
             fs::create_dir_all(&new_layer_dir)?;
             let mut uw = Layer::<()>::unitig_writer(&new_layer_dir).map_err(olm_to_sk)?;
-            g.try_for_each_unitig(|nuc_iter| {
-                let unitig: obikseq::unitig::Unitig = nuc_iter.collect();
-                uw.write(&unitig)
+            g.try_for_each_unitig(|nucs| {
+                uw.write(&obikseq::unitig::Unitig::from_nucleotides(nucs))
             })?;
             uw.close()?;
             let n = g.len();

src/obikpartitionner/src/rebuild_layer.rs

@@ -168,9 +168,8 @@ impl KmerPartition {
         fs::create_dir_all(&dst_layer_dir)?;
 
         let mut uw = Layer::<()>::unitig_writer(&dst_layer_dir).map_err(olm_to_sk)?;
-        g.try_for_each_unitig(|nuc_iter| {
-            let unitig: obikseq::unitig::Unitig = nuc_iter.collect();
-            uw.write(&unitig)
+        g.try_for_each_unitig(|nucs| {
+            uw.write(&obikseq::unitig::Unitig::from_nucleotides(nucs))
         })?;
         uw.close()?;
         drop(g);