use d4::{
task::{Task, TaskPartition},
D4MatrixReader, D4TrackReader, MultiTrackReader,
};
use std::env::args;
struct CountTaskPart {
threshold: i32,
count: u64,
}
impl<T: Iterator<Item = i32> + ExactSizeIterator> TaskPartition<T> for CountTaskPart {
type ParentType = CountTask;
type ResultType = u64;
fn new(_: u32, _: u32, parent: &Self::ParentType) -> Self {
Self {
threshold: parent.3,
count: 0,
}
}
#[inline(always)]
fn feed_range(&mut self, left: u32, right: u32, value: &mut T) -> bool {
if value.all(|x| x > self.threshold) {
self.count += (right - left) as u64;
}
true
}
fn result(&mut self) -> Self::ResultType {
self.count
}
}
struct CountTask(String, u32, u32, i32);
impl<T: Iterator<Item = i32> + ExactSizeIterator> Task<T> for CountTask {
type Partition = CountTaskPart;
type Output = u64;
fn region(&self) -> (&str, u32, u32) {
(self.0.as_str(), self.1, self.2)
}
fn combine(
&self,
parts: &[<Self::Partition as d4::task::TaskPartition<T>>::ResultType],
) -> Self::Output {
parts.iter().sum()
}
}
fn main() {
let input_path = args().nth(1).unwrap();
let tracks: Vec<D4TrackReader> = D4TrackReader::open_tracks(input_path, |_| true).unwrap();
let mut matrix_reader = D4MatrixReader::new(tracks).unwrap();
let tasks: Vec<_> = matrix_reader
.chrom_regions()
.iter()
.map(|&(chr, begin, end)| CountTask(chr.to_string(), begin, end, 10))
.collect();
let result = matrix_reader.run_tasks(tasks).unwrap();
println!("{}", result.into_iter().map(|x| x.output).sum::<u64>());
}