iter_log/

lib.rs

use rayon::iter::plumbing::{Consumer, Folder, UnindexedConsumer};
use rayon::prelude::*;
use std::env::var;
use std::io::{self, Write};
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())
}

const BARS: &str = "--------------------------------------------------";
static STEP_PERCENT: usize = 2;

/// 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 2% of the way through the iteration.
    ///
    /// # Returns
    ///
    /// Returns a new `LogProgress` iterator which tracks and logs progress.
    fn log_progress(self, name: &str) -> LogProgress<Self>;
}

impl<I> LogProgressExt for I
where
    I: Iterator, // Ensure `I` is an iterator
{
    fn log_progress(self, name: &str) -> 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,
            name: name.to_string(),
            t0: std::time::Instant::now(),
        }
    }
}

/// A struct that wraps an iterator and tracks progress.
pub struct LogProgress<I> {
    iter: I,
    progress: usize,
    total: usize,
    name: String,
    t0: std::time::Instant,
}

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 2). 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()?;
        let step = STEP_PERCENT / 2;

        if *enable_log() {
            if self.progress == 0 {
                print!(
                    "+{}+\n| {}{}|\n+{}+\n|",
                    BARS,
                    self.name,
                    " ".repeat(49 - self.name.len()),
                    BARS,
                );
            }
            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 / STEP_PERCENT > old_percent / STEP_PERCENT {
                print!("{}", "=".repeat(step));
                io::stdout().flush().unwrap();
            }
        }

        Some(item)
    }
}

impl<I> LogProgress<I> {
    fn finish(&self) {
        if *enable_log() {
            println!("|\n+{}+", "-".repeat(50));
            let elapsed = self.t0.elapsed();
            let time_str = format!("{:?}", elapsed);
            print!(
                "| Took {} to complete{}|\n+{}+\n",
                time_str,
                " ".repeat(32 - time_str.chars().count()),
                BARS
            );
        }
    }
}

impl<I> Drop for LogProgress<I> {
    fn drop(&mut self) {
        self.finish();
    }
}

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

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(t0: std::time::Instant) -> Arc<Self> {
        Arc::new(Self {
            last_logged: AtomicUsize::new(0),
            pending_logs: Mutex::new(Vec::new()),
            t0,
        })
    }

    /// Logs the progress if it matches the expected step and ensures ordered output.
    ///
    /// # Arguments
    ///
    /// * `progress` - The current progress percentage.
    ///
    /// # 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) {
        if *enable_log() {
            let mut pending = self.pending_logs.lock().unwrap();
            let step = STEP_PERCENT / 2;

            if progress == self.last_logged.load(Ordering::Relaxed) + STEP_PERCENT {
                // Print the progress immediately if it's the next expected one
                print!("{}", "=".repeat(step));
                io::stdout().flush().unwrap();
                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 {
                        print!("{}", "=".repeat(step));
                        io::stdout().flush().unwrap();
                        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
            }
        }
    }

    fn finish(&self) {
        if *enable_log() {
            println!("|\n+{}+", "-".repeat(50));
            let elapsed = self.t0.elapsed();
            let time_str = format!("{:?}", elapsed);
            print!(
                "| Took {} to complete{}|\n+{}+\n",
                time_str,
                " ".repeat(32 - time_str.chars().count()),
                BARS
            );
        }
    }
}

impl Drop for OrderedLogger {
    fn drop(&mut self) {
        self.finish();
    }
}

/// A struct that wraps a parallel iterator and tracks progress.
pub struct LogProgressPar<I> {
    iter: I,
    progress: Arc<AtomicUsize>,
    total: 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,
            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,
    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,
                logger: Arc::clone(&self.logger),
            },
            LogProgressConsumer {
                base: right,
                progress: Arc::clone(&self.progress),
                total: self.total,
                logger: Arc::clone(&self.logger),
            },
            reducer,
        )
    }

    fn into_folder(self) -> Self::Folder {
        LogProgressFolder {
            base: self.base.into_folder(),
            progress: self.progress,
            total: self.total,
            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,
            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,
    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 / STEP_PERCENT > old_percent / STEP_PERCENT {
            let rounded_percent = new_percent - (new_percent % STEP_PERCENT);
            self.logger.log_progress(rounded_percent);
        }

        LogProgressFolder {
            base: self.base.consume(item),
            progress: self.progress,
            total: self.total,
            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 2% of the way through the iteration.
    ///
    /// # Returns
    ///
    /// Returns a new `LogProgressPar` iterator which tracks and logs progress.
    fn log_progress(self, name: &str) -> LogProgressPar<Self>;
}

impl<I> LogProgressParExt for I
where
    I: ParallelIterator + IndexedParallelIterator,
{
    fn log_progress(self, name: &str) -> LogProgressPar<Self> {
        let total = self.len(); // Get the total number of items
        let t0 = std::time::Instant::now();
        let logger = OrderedLogger::new(t0); // Create the logger

        if *enable_log() {
            print!(
                "+{}+\n| {}{}|\n+{}+\n|",
                BARS,
                name,
                " ".repeat(49 - name.len()),
                BARS,
            );
        }

        LogProgressPar {
            iter: self,
            progress: Arc::new(AtomicUsize::new(0)),
            total,
            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
    }
}