iter_log/

lib.rs

use rayon::iter::plumbing::{Consumer, Folder, UnindexedConsumer};
use rayon::prelude::*;
use std::env::var;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::OnceLock;
use std::sync::{Arc, Mutex};

fn enable_log() -> &'static bool {
    static ENABLE_LOG: OnceLock<bool> = OnceLock::new();
    ENABLE_LOG.get_or_init(|| var("ENABLE_ITER_LOG").is_ok())
}

/// Extension trait to add a `log_progress` method to regular iterators.
pub trait LogProgressExt: Iterator + Sized {
    /// Wraps the iterator with progress logging.
    ///
    /// This method will print progress updates every `step_percent`% of the way through the iteration.
    ///
    /// # Arguments
    ///
    /// * `step_percent` - The percentage of items to be processed before logging a progress update.
    ///
    /// # Returns
    ///
    /// Returns a new `LogProgress` iterator which tracks and logs progress.
    fn log_progress(self, step_percent: usize) -> LogProgress<Self>;
}

impl<I> LogProgressExt for I
where
    I: Iterator, // Ensure `I` is an iterator
{
    fn log_progress(self, step_percent: usize) -> LogProgress<Self> {
        let total = self.size_hint().1.unwrap_or(0); // Get the total number of items (if possible)

        LogProgress {
            iter: self,
            progress: 0,
            total,
            step_percent,
        }
    }
}

/// A struct that wraps an iterator and tracks progress.
pub struct LogProgress<I> {
    iter: I,
    progress: usize,
    total: usize,
    step_percent: usize,
}

impl<I: Iterator> Iterator for LogProgress<I> {
    type Item = I::Item;

    /// Returns the next item in the iteration, while logging progress at intervals.
    ///
    /// The method increments the progress count and checks if the progress reaches a new milestone
    /// (i.e., a multiple of `step_percent`). If so, it logs the progress to the console.
    ///
    /// # Returns
    ///
    /// Returns the next item in the iterator, or `None` if the iterator is finished.
    fn next(&mut self) -> Option<Self::Item> {
        let item = self.iter.next()?;

        if *enable_log() {
            self.progress += 1;
            let old_percent = (self.progress * 100) / self.total;
            let new_percent = ((self.progress + 1) * 100) / self.total;

            // Log the progress if we hit a new milestone
            if new_percent / self.step_percent > old_percent / self.step_percent {
                let rounded_percent = new_percent - (new_percent % self.step_percent);
                println!("Progress: {}%", rounded_percent);
            }
        }

        Some(item)
    }
}

/// A struct to handle ordered progress logging during parallel iteration.
struct OrderedLogger {
    last_logged: AtomicUsize,
    pending_logs: Mutex<Vec<usize>>,
}

impl OrderedLogger {
    /// Creates a new `OrderedLogger` instance.
    ///
    /// This function initializes the logger with a fresh `AtomicUsize` and an empty vector for
    /// pending progress updates.
    ///
    /// # Returns
    ///
    /// Returns an `Arc` (thread-safe reference) of the `OrderedLogger`.
    fn new() -> Arc<Self> {
        Arc::new(Self {
            last_logged: AtomicUsize::new(0),
            pending_logs: Mutex::new(Vec::new()),
        })
    }

    /// Logs the progress if it matches the expected step and ensures ordered output.
    ///
    /// # Arguments
    ///
    /// * `progress` - The current progress percentage.
    /// * `step_percent` - The percentage step at which to log progress.
    ///
    /// # Notes
    ///
    /// This method ensures that progress logs are printed in the correct order, even when the
    /// parallel tasks report progress asynchronously.
    fn log_progress(&self, progress: usize, step_percent: usize) {
        if *enable_log() {
            let mut pending = self.pending_logs.lock().unwrap();

            if progress == self.last_logged.load(Ordering::Relaxed) + step_percent {
                // Print the progress immediately if it's the next expected one
                println!("Progress: {}%", progress);
                self.last_logged.fetch_add(step_percent, Ordering::Relaxed);

                // Print any pending logs that can now be processed in order
                while let Some(&next) = pending.first() {
                    if next == self.last_logged.load(Ordering::Relaxed) + step_percent {
                        println!("Progress: {}%", next);
                        self.last_logged.fetch_add(step_percent, Ordering::Relaxed);
                        pending.remove(0);
                    } else {
                        break;
                    }
                }
            } else {
                // If progress is not expected yet, store it for later
                pending.push(progress);
                pending.sort_unstable(); // Sort pending progress updates
            }
        }
    }
}

/// A struct that wraps a parallel iterator and tracks progress.
pub struct LogProgressPar<I> {
    iter: I,
    progress: Arc<AtomicUsize>,
    total: usize,
    step_percent: usize,
    logger: Arc<OrderedLogger>,
}

impl<I> ParallelIterator for LogProgressPar<I>
where
    I: ParallelIterator,
{
    type Item = I::Item;

    /// Drives the parallel iteration using a custom consumer.
    ///
    /// This method wraps the original consumer with the `LogProgressConsumer`, which tracks the
    /// progress and logs it at the specified intervals.
    ///
    /// # Arguments
    ///
    /// * `consumer` - The base consumer that will process the items of the parallel iterator.
    ///
    /// # Returns
    ///
    /// Returns the result of the parallel iteration after it is consumed.
    fn drive_unindexed<C>(self, consumer: C) -> C::Result
    where
        C: UnindexedConsumer<Self::Item>,
    {
        let wrapped_consumer = LogProgressConsumer {
            base: consumer,
            progress: self.progress,
            total: self.total,
            step_percent: self.step_percent,
            logger: self.logger,
        };
        self.iter.drive_unindexed(wrapped_consumer)
    }
}

/// A consumer for parallel iterations that tracks progress and logs it.
struct LogProgressConsumer<C> {
    base: C,
    progress: Arc<AtomicUsize>,
    total: usize,
    step_percent: usize,
    logger: Arc<OrderedLogger>,
}

impl<C, T> Consumer<T> for LogProgressConsumer<C>
where
    C: Consumer<T>,
{
    type Folder = LogProgressFolder<C::Folder>;
    type Reducer = C::Reducer;
    type Result = C::Result;

    /// Splits the consumer at the specified index and returns two new consumers.
    ///
    /// This method ensures that progress tracking is correctly propagated through both consumers.
    ///
    /// # Arguments
    ///
    /// * `index` - The index at which to split the consumer.
    ///
    /// # Returns
    ///
    /// Returns two new consumers and a reducer for parallel reduction.
    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
        let (left, right, reducer) = self.base.split_at(index);
        (
            LogProgressConsumer {
                base: left,
                progress: Arc::clone(&self.progress),
                total: self.total,
                step_percent: self.step_percent,
                logger: Arc::clone(&self.logger),
            },
            LogProgressConsumer {
                base: right,
                progress: Arc::clone(&self.progress),
                total: self.total,
                step_percent: self.step_percent,
                logger: Arc::clone(&self.logger),
            },
            reducer,
        )
    }

    fn into_folder(self) -> Self::Folder {
        LogProgressFolder {
            base: self.base.into_folder(),
            progress: self.progress,
            total: self.total,
            step_percent: self.step_percent,
            logger: Arc::clone(&self.logger),
        }
    }

    fn full(&self) -> bool {
        self.base.full()
    }
}

impl<C, T> UnindexedConsumer<T> for LogProgressConsumer<C>
where
    C: UnindexedConsumer<T>,
{
    fn split_off_left(&self) -> Self {
        LogProgressConsumer {
            base: self.base.split_off_left(),
            progress: Arc::clone(&self.progress),
            total: self.total,
            step_percent: self.step_percent,
            logger: Arc::clone(&self.logger),
        }
    }

    fn to_reducer(&self) -> Self::Reducer {
        self.base.to_reducer()
    }
}

/// A folder for processing items in parallel while tracking progress.
struct LogProgressFolder<F> {
    base: F,
    progress: Arc<AtomicUsize>,
    total: usize,
    step_percent: usize,
    logger: Arc<OrderedLogger>,
}

impl<F, T> Folder<T> for LogProgressFolder<F>
where
    F: Folder<T>,
{
    type Result = F::Result;

    /// Consumes an item and tracks progress.
    ///
    /// This method updates the progress counter and logs progress at specified intervals.
    ///
    /// # Arguments
    ///
    /// * `item` - The item to consume and process.
    ///
    /// # Returns
    ///
    /// Returns a new `LogProgressFolder` with the updated state.
    fn consume(self, item: T) -> Self {
        let old_count = self.progress.fetch_add(1, Ordering::Relaxed);
        let old_percent = (old_count * 100) / self.total;
        let new_percent = ((old_count + 1) * 100) / self.total;

        if new_percent / self.step_percent > old_percent / self.step_percent {
            let rounded_percent = new_percent - (new_percent % self.step_percent);
            self.logger.log_progress(rounded_percent, self.step_percent);
        }

        LogProgressFolder {
            base: self.base.consume(item),
            progress: self.progress,
            total: self.total,
            step_percent: self.step_percent,
            logger: self.logger,
        }
    }

    fn complete(self) -> Self::Result {
        self.base.complete()
    }

    fn full(&self) -> bool {
        self.base.full()
    }
}

/// Extension trait to add a `log_progress` method to parallel iterators.
pub trait LogProgressParExt: Sized + ParallelIterator {
    /// Wraps the parallel iterator with progress logging.
    ///
    /// This method will print progress updates every `step_percent`% of the way through the iteration.
    ///
    /// # Arguments
    ///
    /// * `step_percent` - The percentage of items to be processed before logging a progress update.
    ///
    /// # Returns
    ///
    /// Returns a new `LogProgressPar` iterator which tracks and logs progress.
    fn log_progress(self, step_percent: usize) -> LogProgressPar<Self>;
}

impl<I> LogProgressParExt for I
where
    I: ParallelIterator + IndexedParallelIterator,
{
    fn log_progress(self, step_percent: usize) -> LogProgressPar<Self> {
        let total = self.len(); // Get the total number of items
        let logger = OrderedLogger::new(); // Create the logger

        LogProgressPar {
            iter: self,
            progress: Arc::new(AtomicUsize::new(0)),
            total,
            step_percent,
            logger,
        }
    }
}

/// A long computation function to simulate more expensive work per item.
pub fn long_computation(x: u32) -> u32 {
    // Simulate a heavy calculation by introducing a delay
    let mut result = x;
    for _ in 0..1_000_000 {
        result = result.saturating_mul(2);
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_long_computation() {
        assert_eq!(long_computation(2), 4294967295); // Expected result
    }
}