Rust backend: parallelism fixes for bigWig writer, alignmentSieve, computeMatrix

src/alignmentsieve.rs

@@ -4,6 +4,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyList;
 use rust_htslib::bam::record::CigarString;
 use rust_htslib::bam::{self, Header, Read, Reader, Writer};
+use rust_htslib::tpool::ThreadPool;
 use std::collections::HashMap;
 use std::fs::File;
 use std::io::{BufWriter, Write};
@@ -108,7 +109,13 @@ pub fn r_alignmentsieve(
         None,
     );
 
-    // Open output writers
+    // Open output writers. When both the main and filtered-out BAM writers are active,
+    // give them a single shared htslib thread pool instead of calling `set_threads` on
+    // each independently: `set_threads` spins up its own private pool per call, so two
+    // independent calls with `writerthreads` each would request `2 * writerthreads` BGZF
+    // compression threads total, oversubscribing the `writerthreads` budget we intended.
+    let writer_pool = ThreadPool::new(writerthreads as u32).ok();
+
     let mut obam = if !bed {
         Some(
             Writer::from_path(ofile, &header, bam::Format::Bam)
@@ -118,7 +125,14 @@ pub fn r_alignmentsieve(
         None
     };
     if let Some(ref mut w) = obam {
-        let _ = w.set_threads(writerthreads);
+        match &writer_pool {
+            Some(pool) => {
+                let _ = w.set_thread_pool(pool);
+            }
+            None => {
+                let _ = w.set_threads(writerthreads);
+            }
+        }
     }
     let mut obed = if bed {
         Some(BufWriter::new(File::create(ofile).unwrap_or_else(|e| {
@@ -143,7 +157,14 @@ pub fn r_alignmentsieve(
         None
     };
     if let Some(ref mut w) = ofilterbam {
-        let _ = w.set_threads(writerthreads);
+        match &writer_pool {
+            Some(pool) => {
+                let _ = w.set_thread_pool(pool);
+            }
+            None => {
+                let _ = w.set_threads(writerthreads);
+            }
+        }
     }
     let mut ofilterbed = if write_filters && bed {
         Some(BufWriter::new(

src/bamcompare.rs

@@ -395,7 +395,7 @@ pub fn r_bamcompare(
     } else {
         Box::new(raw_lines)
     };
-    write_covfile(lines, ofile, ofiletype, chromsizes);
+    write_covfile(lines, ofile, ofiletype, chromsizes, nproc);
     Ok(())
 }
 

src/bamcoverage.rs

@@ -289,6 +289,6 @@ pub fn r_bamcoverage(
         println!("Writing output to: {}", ofile);
     }
 
-    write_covfile(lines, ofile, ofiletype, chromsizes);
+    write_covfile(lines, ofile, ofiletype, chromsizes, nproc);
     Ok(())
 }

src/computematrix.rs

@@ -273,28 +273,45 @@ pub fn r_computematrix(
 
     // Discriminate between reference-point and scale-regions mode.
 
-    let matrix: Vec<Vec<f32>> = pool.install(|| {
-        bw_files
+    // bwintervals is single-threaded per bigwig file (BigWigRead::get_interval takes
+    // &mut self, so one reader can't be shared across threads). Parallelizing only
+    // across bw_files.par_iter() therefore leaves most of a reserved node idle whenever
+    // there are fewer tracks than threads, which is the common case. Add a second
+    // parallelism axis by also chunking the region list per file, mirroring the
+    // "open one reader per worker, reuse across a chunk of regions" pattern bam_pileup
+    // already uses for BAM reading. Chunk count in one dimension is intentionally
+    // rounded up so busy still finishes at nproc-ish task count even with 1 bigwig file.
+    let chunks_per_file = (nproc / bw_files.len().max(1)).max(1);
+    let region_chunks = chunk_ranges(regions.len(), chunks_per_file);
+    let tasks: Vec<(usize, usize, usize)> = (0..bw_files.len())
+        .flat_map(|fi| region_chunks.iter().map(move |&(s, e)| (fi, s, e)))
+        .collect();
+
+    // rayon's par_iter().collect() on an indexed source (a Vec, here) is
+    // order-preserving: `results[k]` always corresponds to `tasks[k]`, regardless of
+    // which task finishes first. That lets the merge below stay a plain sequential
+    // walk instead of needing its own bookkeeping to figure out where each chunk goes.
+    let results: Vec<Vec<Vec<f32>>> = pool.install(|| {
+        tasks
             .par_iter()
-            .map(|i| {
+            .map(|&(fi, s, e)| {
                 bwintervals(
-                    &i,
-                    &regions,
-                    &slopregions,
+                    &bw_files[fi],
+                    &regions[s..e],
+                    &slopregions[s..e],
                     &scale_regions,
                     blacklist_index.as_deref(),
                 )
             })
-            .reduce(
-                || vec![vec![]; regions.len()],
-                |mut acc, vec_of_vecs| {
-                    for (i, inner_vec) in vec_of_vecs.into_iter().enumerate() {
-                        acc[i].extend(inner_vec);
-                    }
-                    acc
-                },
-            )
+            .collect()
     });
+
+    let mut matrix: Vec<Vec<f32>> = vec![Vec::new(); regions.len()];
+    for (&(_fi, s, e), chunk_result) in tasks.iter().zip(results.into_iter()) {
+        for (local_ix, region_ix) in (s..e).enumerate() {
+            matrix[region_ix].extend(chunk_result[local_ix].iter().copied());
+        }
+    }
     matrix_dump(
         sortregions,
         sortusing,
@@ -312,6 +329,27 @@ pub fn r_computematrix(
     Ok(())
 }
 
+/// Split `[0, n)` into up to `k` contiguous, roughly-equal ranges. Returns a single
+/// `(0, n)` range if `n == 0` or `k <= 1`, so callers never need to special-case
+/// "no chunking" separately from "one chunk".
+fn chunk_ranges(n: usize, k: usize) -> Vec<(usize, usize)> {
+    if n == 0 || k <= 1 {
+        return vec![(0, n)];
+    }
+    let k = k.min(n);
+    let base = n / k;
+    let rem = n % k;
+    let mut ranges = Vec::with_capacity(k);
+    let mut start = 0;
+    for i in 0..k {
+        let len = base + if i < rem { 1 } else { 0 };
+        let end = start + len;
+        ranges.push((start, end));
+        start = end;
+    }
+    ranges
+}
+
 fn slop_region(
     region: &Region,
     scale_regions: &Scalingregions,

src/filehandler.rs

@@ -31,7 +31,13 @@ pub fn bam_ispaired(bam_ifile: &str) -> bool {
     return false;
 }
 
-pub fn write_covfile<LI>(lines: LI, ofile: &str, filetype: &str, chromsizes: HashMap<String, u32>)
+pub fn write_covfile<LI>(
+    lines: LI,
+    ofile: &str,
+    filetype: &str,
+    chromsizes: HashMap<String, u32>,
+    nproc: usize,
+)
 where
     LI: Iterator<Item = (String, Value)>,
 {
@@ -51,12 +57,26 @@ where
         }
     } else {
         let vals = BedParserStreamingIterator::wrap_infallible_iter(lines, false);
-        let runtime = tokio::runtime::Builder::new_multi_thread()
-            .worker_threads(1)
-            .build()
-            .expect("Unable to create tokio runtime for bw writing.");
-        let writer = BigWigWrite::create_file(ofile, chromsizes)
+        let mut writer = BigWigWrite::create_file(ofile, chromsizes)
             .unwrap_or_else(|e| panic!("Failed to create output bigwig file '{}': {}", ofile, e));
+        // bigtools spawns one async task per chromosome (zoom aggregation + section
+        // encoding) and funnels their output through a single sequential writer task
+        // that owns the actual file handle; the crate's own channels serialize that
+        // hand-off, so raising worker count only widens the per-chromosome encode
+        // stage, it does not add any concurrent writers to `ofile`.
+        // Mirrors bigtools' own bedgraphtobigwig CLI: skip the multi-thread runtime
+        // entirely at nproc == 1, since it would only ever have one consumer.
+        let runtime = if nproc <= 1 {
+            writer.options.channel_size = 0;
+            tokio::runtime::Builder::new_current_thread()
+                .build()
+                .expect("Unable to create tokio runtime for bw writing.")
+        } else {
+            tokio::runtime::Builder::new_multi_thread()
+                .worker_threads(nproc)
+                .build()
+                .expect("Unable to create tokio runtime for bw writing.")
+        };
         let _ = writer.write(vals, runtime);
     }
 }
@@ -657,8 +677,8 @@ pub fn chrombounds_from_bam(bamfiles: Vec<&str>) -> HashMap<String, u32> {
 
 pub fn bwintervals(
     bwfile: &str,
-    regions: &Vec<Region>,
-    slopregions: &Vec<Vec<Bin>>,
+    regions: &[Region],
+    slopregions: &[Vec<Bin>],
     scale_regions: &Scalingregions,
     blacklist_index: Option<&crate::filtering::BlacklistIndex>,
 ) -> Vec<Vec<f32>> {