paraseq 0.4.14

use smallvec::SmallVec;

use crate::{Record, MAX_ARITY};

use super::error::Result;

/// Trait implemented for a type that processes records in parallel
pub trait GenericProcessor<Rf>: Send + Clone {
    /// Called on a batch of records.
    fn process_record_batch(&mut self, records: impl Iterator<Item = Rf>) -> Result<()>;

    /// Called when a batch of records is complete
    fn on_batch_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Called when the processing for a thread is complete
    fn on_thread_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Sets the thread id for the processor
    #[allow(unused_variables)]
    fn set_thread_id(&mut self, thread_id: usize) {
        // Default implementation does nothing
    }
}

/// Implement either `process_record_batch` or `process_record`.
/// By default, `process_record_batch` calls `process_record` for each record.
pub trait ParallelProcessor<Rf: Record>: Send + Clone {
    /// Called on a batch of records.
    fn process_record_batch(&mut self, records: impl Iterator<Item = Rf>) -> Result<()> {
        for record in records {
            self.process_record(record)?;
        }
        Ok(())
    }

    /// Called on an individual record
    fn process_record(&mut self, _record: Rf) -> Result<()> {
        unimplemented!("Either ParallelProcessor::process_record or ParallelProcessor::process_record_batch must be implemented!");
    }

    /// Called when a batch of records is complete
    fn on_batch_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Called when the processing for a thread is complete
    fn on_thread_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Sets the thread id for the processor
    #[allow(unused_variables)]
    fn set_thread_id(&mut self, thread_id: usize) {
        // Default implementation does nothing
    }

    /// Gets the thread id for the processor
    fn get_thread_id(&self) -> usize {
        unimplemented!("Must be implemented by the processor to be used")
    }
}

impl<Rf: Record, P: ParallelProcessor<Rf>> GenericProcessor<Rf> for P {
    fn process_record_batch(&mut self, records: impl Iterator<Item = Rf>) -> Result<()> {
        self.process_record_batch(records)
    }
    fn on_batch_complete(&mut self) -> Result<()> {
        self.on_batch_complete()
    }
    fn on_thread_complete(&mut self) -> Result<()> {
        self.on_thread_complete()
    }
    fn set_thread_id(&mut self, thread_id: usize) {
        self.set_thread_id(thread_id);
    }
}

pub trait PairedParallelProcessor<Rf: Record>: Send + Clone {
    /// Called on a batch of record pairs.
    fn process_record_pair_batch(
        &mut self,
        record_pairs: impl Iterator<Item = (Rf, Rf)>,
    ) -> Result<()> {
        for record_pair in record_pairs {
            self.process_record_pair(record_pair.0, record_pair.1)?;
        }
        Ok(())
    }

    /// Called on an individual record pair.
    fn process_record_pair(&mut self, _record1: Rf, _record2: Rf) -> Result<()> {
        unimplemented!("Either PairedParallelProcessor::process_record_pair or PairedParallelProcessor::process_record_pair_batch must be implemented!");
    }

    /// Called when a batch of records is complete
    fn on_batch_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Called when the processing for a thread is complete
    fn on_thread_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Sets the thread id for the processor
    #[allow(unused_variables)]
    fn set_thread_id(&mut self, thread_id: usize) {
        // Default implementation does nothing
    }

    /// Gets the thread id for the processor
    fn get_thread_id(&self) -> usize {
        unimplemented!("Must be implemented by the processor to be used")
    }
}

impl<Rf: Record, P: PairedParallelProcessor<Rf>> GenericProcessor<(Rf, Rf)> for P {
    fn process_record_batch(&mut self, records: impl Iterator<Item = (Rf, Rf)>) -> Result<()> {
        self.process_record_pair_batch(records)
    }
    fn on_batch_complete(&mut self) -> Result<()> {
        self.on_batch_complete()
    }
    fn on_thread_complete(&mut self) -> Result<()> {
        self.on_thread_complete()
    }
    fn set_thread_id(&mut self, thread_id: usize) {
        self.set_thread_id(thread_id);
    }
}

pub trait MultiParallelProcessor<Rf: Record>: Send + Clone {
    /// Called on a batch of multi-records.
    fn process_multi_record_batch(
        &mut self,
        multi_records: impl Iterator<Item = SmallVec<[Rf; MAX_ARITY]>>,
    ) -> Result<()> {
        for multi_record in multi_records {
            self.process_multi_record(&multi_record)?;
        }
        Ok(())
    }

    /// Called on an individual set of record
    fn process_multi_record(&mut self, _records: &[Rf]) -> Result<()> {
        unimplemented!("Either MultiParallelProcessor::process_multi_record or MultiParallelProcessor::process_multi_record_batch must be implemented!");
    }

    /// Called when a batch of records is complete
    fn on_batch_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Called when the processing for a thread is complete
    fn on_thread_complete(&mut self) -> Result<()> {
        Ok(())
    }

    /// Sets the thread id for the processor
    #[allow(unused_variables)]
    fn set_thread_id(&mut self, thread_id: usize) {
        // Default implementation does nothing
    }

    /// Gets the thread id for the processor
    fn get_thread_id(&self) -> usize {
        unimplemented!("Must be implemented by the processor to be used")
    }
}

impl<Rf: Record, P: MultiParallelProcessor<Rf>> GenericProcessor<SmallVec<[Rf; MAX_ARITY]>> for P {
    fn process_record_batch(
        &mut self,
        multi_records: impl Iterator<Item = SmallVec<[Rf; MAX_ARITY]>>,
    ) -> Result<()> {
        self.process_multi_record_batch(multi_records)
    }
    fn on_batch_complete(&mut self) -> Result<()> {
        self.on_batch_complete()
    }
    fn on_thread_complete(&mut self) -> Result<()> {
        self.on_thread_complete()
    }
    fn set_thread_id(&mut self, thread_id: usize) {
        self.set_thread_id(thread_id);
    }
}

/// A convenience wrapper that can create a [`ParallelProcessor`] from a closure
/// that takes an `&dyn Iterator<Item=Rf>` iterator over records.
///
/// (To avoid the `&dyn`, we would have to lift the exact type of the iterator
/// as a generic of `ParallelProcessor`, but that does not seem worth it.)
impl<Rf: Record, F> ParallelProcessor<Rf> for F
where
    F: for<'a> FnMut(&'a mut dyn Iterator<Item = Rf>) -> Result<()> + Send + Clone,
{
    fn process_record_batch(&mut self, mut records: impl Iterator<Item = Rf>) -> Result<()> {
        self(&mut records)
    }
}

impl<Rf: Record, F> PairedParallelProcessor<Rf> for F
where
    F: FnMut(&mut dyn Iterator<Item = (Rf, Rf)>) -> Result<()> + Send + Clone,
{
    fn process_record_pair_batch(
        &mut self,
        mut record_pairs: impl Iterator<Item = (Rf, Rf)>,
    ) -> Result<()> {
        self(&mut record_pairs)
    }
}

impl<Rf: Record, F> MultiParallelProcessor<Rf> for F
where
    F: FnMut(&mut dyn Iterator<Item = SmallVec<[Rf; MAX_ARITY]>>) -> Result<()> + Send + Clone,
{
    fn process_multi_record_batch(
        &mut self,
        mut multi_records: impl Iterator<Item = SmallVec<[Rf; MAX_ARITY]>>,
    ) -> Result<()> {
        self(&mut multi_records)
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;

    use smallvec::SmallVec;

    use crate::fastq;
    use crate::parallel::ParallelReader;
    use crate::MAX_ARITY;

    fn make_fastq(n: usize) -> Vec<u8> {
        (0..n)
            .flat_map(|i| format!("@seq{i}\nACGT\n+\nIIII\n").into_bytes())
            .collect()
    }

    // These closures rely entirely on the trait defaults (on_batch_complete,
    // on_thread_complete, set_thread_id) rather than overriding them, unlike
    // every struct-based processor elsewhere in the test suite.
    #[test]
    fn test_closure_parallel_processor_defaults() {
        let count = Arc::new(AtomicUsize::new(0));
        let count_ref = Arc::clone(&count);
        let mut processor = move |batch: &mut dyn Iterator<Item = fastq::RefRecord>| {
            for _ in batch {
                count_ref.fetch_add(1, Ordering::Relaxed);
            }
            Ok(())
        };

        let reader = fastq::Reader::new(Cursor::new(make_fastq(20)));
        reader.process_parallel(&mut processor, 1).unwrap();

        assert_eq!(count.load(Ordering::Relaxed), 20);
    }

    #[test]
    fn test_closure_paired_parallel_processor_defaults() {
        let count = Arc::new(AtomicUsize::new(0));
        let count_ref = Arc::clone(&count);
        let mut processor =
            move |batch: &mut dyn Iterator<Item = (fastq::RefRecord, fastq::RefRecord)>| {
                for _ in batch {
                    count_ref.fetch_add(1, Ordering::Relaxed);
                }
                Ok(())
            };

        let reader = fastq::Reader::new(Cursor::new(make_fastq(20)));
        reader
            .process_parallel_interleaved(&mut processor, 1)
            .unwrap();

        assert_eq!(count.load(Ordering::Relaxed), 10);
    }

    #[test]
    fn test_closure_multi_parallel_processor_defaults() {
        let count = Arc::new(AtomicUsize::new(0));
        let count_ref = Arc::clone(&count);
        let mut processor =
            move |batch: &mut dyn Iterator<Item = SmallVec<[fastq::RefRecord; MAX_ARITY]>>| {
                for _ in batch {
                    count_ref.fetch_add(1, Ordering::Relaxed);
                }
                Ok(())
            };

        let reader = fastq::Reader::new(Cursor::new(make_fastq(20)));
        reader
            .process_parallel_multi_interleaved(2, &mut processor, 1)
            .unwrap();

        assert_eq!(count.load(Ordering::Relaxed), 10);
    }
}