iter_log/

lib.rs

#[cfg(feature = "rayon")]
use rayon::iter::plumbing::{Consumer, Folder, UnindexedConsumer};
#[cfg(feature = "rayon")]
use rayon::prelude::*;
use std::env::var;
use std::io::{self, Write};
#[cfg(feature = "rayon")]
use std::sync::Arc;
use std::sync::Mutex;
#[cfg(feature = "rayon")]
use std::sync::atomic::{AtomicUsize, Ordering};
#[cfg(feature = "time")]
use std::time::Instant;

static ENABLE_LOG: Mutex<Option<bool>> = Mutex::new(None);

/// Gets the current logging enabled state.
///
/// On first call, initializes from the `ENABLE_ITER_LOG` environment variable.
/// Subsequently, returns the value set via `set_log_enabled()`.
fn enable_log() -> bool {
    let mut state = ENABLE_LOG.lock().unwrap();
    match *state {
        Some(enabled) => enabled,
        None => {
            let enabled = var("ENABLE_ITER_LOG").is_ok();
            *state = Some(enabled);
            enabled
        }
    }
}

/// Sets the logging enabled state at runtime.
///
/// This allows you to enable or disable logging dynamically without relying on
/// the `ENABLE_ITER_LOG` environment variable. Once set, this value will be used
/// for all subsequent logging decisions.
///
/// # Arguments
///
/// * `enabled` - `true` to enable logging, `false` to disable it.
///
/// # Examples
///
/// ```no_run
/// use iter_log::set_log_enabled;
///
/// // Enable logging at runtime
/// set_log_enabled(true);
///
/// // Disable logging at runtime
/// set_log_enabled(false);
/// ```
pub fn set_log_enabled(enabled: bool) {
    let mut state = ENABLE_LOG.lock().unwrap();
    *state = Some(enabled);
}

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(),
            #[cfg(feature = "time")]
            t0: Instant::now(),
        }
    }
}

/// A struct that wraps an iterator and tracks progress.
pub struct LogProgress<I> {
    iter: I,
    progress: usize,
    total: usize,
    name: String,
    #[cfg(feature = "time")]
    t0: 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));
            #[cfg(feature = "time")]
            let elapsed = self.t0.elapsed();
            #[cfg(feature = "time")]
            let time_str = format!("{:?}", elapsed);
            #[cfg(feature = "time")]
            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.
#[cfg(feature = "rayon")]
struct OrderedLogger {
    last_logged: AtomicUsize,
    pending_logs: Mutex<Vec<usize>>,
    #[cfg(feature = "time")]
    t0: Instant,
}

#[cfg(feature = "rayon")]
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`.
    #[cfg(feature = "time")]
    fn new(t0: Instant) -> Arc<Self> {
        Arc::new(Self {
            last_logged: AtomicUsize::new(0),
            pending_logs: Mutex::new(Vec::new()),
            t0,
        })
    }

    /// 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`.
    #[cfg(not(feature = "time"))]
    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.
    ///
    /// # 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));
            #[cfg(feature = "time")]
            let elapsed = self.t0.elapsed();
            #[cfg(feature = "time")]
            let time_str = format!("{:?}", elapsed);
            #[cfg(feature = "time")]
            print!(
                "| Took {} to complete{}|\n+{}+\n",
                time_str,
                " ".repeat(32 - time_str.chars().count()),
                BARS
            );
        }
    }
}

#[cfg(feature = "rayon")]
impl Drop for OrderedLogger {
    fn drop(&mut self) {
        self.finish();
    }
}

/// A struct that wraps a parallel iterator and tracks progress.
#[cfg(feature = "rayon")]
pub struct LogProgressPar<I> {
    iter: I,
    progress: Arc<AtomicUsize>,
    total: usize,
    logger: Arc<OrderedLogger>,
}

#[cfg(feature = "rayon")]
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.
#[cfg(feature = "rayon")]
struct LogProgressConsumer<C> {
    base: C,
    progress: Arc<AtomicUsize>,
    total: usize,
    logger: Arc<OrderedLogger>,
}

#[cfg(feature = "rayon")]
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()
    }
}

#[cfg(feature = "rayon")]
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.
#[cfg(feature = "rayon")]
struct LogProgressFolder<F> {
    base: F,
    progress: Arc<AtomicUsize>,
    total: usize,
    logger: Arc<OrderedLogger>,
}

#[cfg(feature = "rayon")]
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.
#[cfg(feature = "rayon")]
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>;
}

#[cfg(feature = "rayon")]
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
        #[cfg(feature = "time")]
        let t0 = Instant::now();
        #[cfg(feature = "time")]
        let logger = OrderedLogger::new(t0); // Create the logger
        #[cfg(not(feature = "time"))]
        let logger = OrderedLogger::new(); // 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
    }

    #[test]
    fn test_set_log_enabled_true() {
        // Reset to initial state by setting to false
        set_log_enabled(false);
        assert!(!enable_log());

        // Enable logging
        set_log_enabled(true);
        assert!(enable_log());
    }

    #[test]
    fn test_set_log_enabled_false() {
        // Enable logging first
        set_log_enabled(true);
        assert!(enable_log());

        // Disable logging
        set_log_enabled(false);
        assert!(!enable_log());
    }

    #[test]
    fn test_set_log_enabled_multiple_calls() {
        // Test multiple toggles
        set_log_enabled(true);
        assert!(enable_log());

        set_log_enabled(false);
        assert!(!enable_log());

        set_log_enabled(true);
        assert!(enable_log());

        set_log_enabled(false);
        assert!(!enable_log());
    }

    #[test]
    fn test_log_progress_with_logging_disabled() {
        set_log_enabled(false);

        let data = vec![1, 2, 3, 4, 5];
        let result: Vec<_> = data.iter().log_progress("test").map(|&x| x * 2).collect();

        assert_eq!(result, vec![2, 4, 6, 8, 10]);
    }

    #[test]
    fn test_log_progress_with_logging_enabled() {
        set_log_enabled(true);

        let data = vec![1, 2, 3, 4, 5];
        let result: Vec<_> = data.iter().log_progress("test").map(|&x| x * 2).collect();

        assert_eq!(result, vec![2, 4, 6, 8, 10]);
    }

    #[cfg(feature = "rayon")]
    #[test]
    fn test_log_progress_par_with_logging_disabled() {
        use rayon::prelude::*;

        set_log_enabled(false);

        let data = vec![1, 2, 3, 4, 5];
        let result: Vec<_> = data
            .par_iter()
            .log_progress("test_par")
            .map(|&x| x * 2)
            .collect();

        result.iter().for_each(|&x| {
            assert!(x % 2 == 0); // All results should be even
        });
    }

    #[cfg(feature = "rayon")]
    #[test]
    fn test_log_progress_par_with_logging_enabled() {
        use rayon::prelude::*;

        set_log_enabled(true);

        let data = vec![1, 2, 3, 4, 5];
        let result: Vec<_> = data
            .par_iter()
            .log_progress("test_par")
            .map(|&x| x * 2)
            .collect();

        result.iter().for_each(|&x| {
            assert!(x % 2 == 0); // All results should be even
        });
    }

    #[test]
    fn test_toggle_logging_within_single_run() {
        let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];

        // First iteration with logging disabled
        set_log_enabled(false);
        assert!(!enable_log());
        let result1: Vec<_> = data
            .iter()
            .log_progress("iteration_1_disabled")
            .map(|&x| x * 2)
            .collect();
        assert_eq!(result1.len(), 10);

        // Second iteration with logging enabled
        set_log_enabled(true);
        assert!(enable_log());
        let result2: Vec<_> = data
            .iter()
            .log_progress("iteration_2_enabled")
            .map(|&x| x * 2)
            .collect();
        assert_eq!(result2.len(), 10);

        // Third iteration with logging disabled again
        set_log_enabled(false);
        assert!(!enable_log());
        let result3: Vec<_> = data
            .iter()
            .log_progress("iteration_3_disabled")
            .map(|&x| x * 2)
            .collect();
        assert_eq!(result3.len(), 10);

        // Fourth iteration with logging enabled again
        set_log_enabled(true);
        assert!(enable_log());
        let result4: Vec<_> = data
            .iter()
            .log_progress("iteration_4_enabled")
            .map(|&x| x * 2)
            .collect();
        assert_eq!(result4.len(), 10);

        // Verify all results are identical
        assert_eq!(result1, result2);
        assert_eq!(result2, result3);
        assert_eq!(result3, result4);
    }

    #[cfg(feature = "rayon")]
    #[test]
    fn test_toggle_logging_parallel_within_single_run() {
        use rayon::prelude::*;

        let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];

        // First parallel iteration with logging disabled
        set_log_enabled(false);
        assert!(!enable_log());
        let result1: Vec<_> = data
            .par_iter()
            .log_progress("par_iteration_1_disabled")
            .map(|&x| x * 2)
            .collect();
        result1.iter().for_each(|&x| {
            assert!(x % 2 == 0);
        });

        // Second parallel iteration with logging enabled
        set_log_enabled(true);
        assert!(enable_log());
        let result2: Vec<_> = data
            .par_iter()
            .log_progress("par_iteration_2_enabled")
            .map(|&x| x * 2)
            .collect();
        result2.iter().for_each(|&x| {
            assert!(x % 2 == 0);
        });

        // Third parallel iteration with logging disabled again
        set_log_enabled(false);
        assert!(!enable_log());
        let result3: Vec<_> = data
            .par_iter()
            .log_progress("par_iteration_3_disabled")
            .map(|&x| x * 2)
            .collect();
        result3.iter().for_each(|&x| {
            assert!(x % 2 == 0);
        });

        // Fourth parallel iteration with logging enabled again
        set_log_enabled(true);
        assert!(enable_log());
        let result4: Vec<_> = data
            .par_iter()
            .log_progress("par_iteration_4_enabled")
            .map(|&x| x * 2)
            .collect();
        result4.iter().for_each(|&x| {
            assert!(x % 2 == 0);
        });

        // Verify state toggles were properly applied
        assert!(enable_log());
    }
}