shadow_counted/

lib.rs

#![doc = include_str!("../README.md")]

use std::cell::Cell;

/// Trait for types that can be committed from a child to a parent during nesting.
///
/// When a nested iterator commits, this trait controls how the supplemental data
/// flows from the child to the parent. The default implementations for `()`, `Rc<T>`,
/// and `Arc<T>` are no-ops, meaning the data doesn't propagate back.
///
/// The parent is immutable (`&self`) and the child is consumed by value. The implementor
/// must use interior mutability (e.g., `Cell`, `RefCell`) if they need to mutate the parent.
///
/// # Examples
///
/// ```
/// use shadow_counted::Commit;
/// use std::cell::Cell;
///
/// #[derive(Clone, Default)]
/// struct Counter {
///     count: Cell<u32>,
/// }
///
/// impl Commit for Counter {
///     fn commit(&self, from: Self) {
///         self.count.set(self.count.get() + from.count.get());
///     }
/// }
/// ```
pub trait Commit {
    /// Called when a nested iterator commits to its parent.
    /// The implementation should merge/update the parent's state using interior mutability.
    ///
    /// # Arguments
    ///
    /// * `from` - The child's supplemental data (consumed by value)
    fn commit(&self, from: Self);
}

// No-op implementation for unit type (default)
impl Commit for () {
    fn commit(&self, _from: Self) {}
}

// No-op implementation for Rc (sharing via clone is sufficient)
impl<T> Commit for std::rc::Rc<T> {
    fn commit(&self, _from: Self) {}
}

// No-op implementation for Arc (sharing via clone is sufficient)
impl<T> Commit for std::sync::Arc<T> {
    fn commit(&self, _from: Self) {}
}

/// An iterator that counts every iteration and optionally commits the count to a parent iterator.
/// Can carry supplemental data of type S through the iteration process.
///
/// Note that the `ShadowCountedIter` implements methods on its own. These can only be accessed
/// while having the iterator itself available. Thus using a `for` loop is often not possible,
/// instead a `while let Some(item) = iterator.next() {...}` loop is required to iterate over
/// its elements.
#[derive(Debug, Clone)]
pub struct ShadowCountedIter<'a, I: Iterator, S = ()> {
    iter: I,
    counter: ShadowCounter<'a, S>,
}

impl<'a, I: Iterator, S: Default> ShadowCountedIter<'a, I, S> {
    /// Creates a new ShadowCountedIter from an iterator.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let mut iter = vec![1, 2, 3].into_iter().shadow_counted();
    /// while let Some(_) = iter.next() {}
    /// assert_eq!(iter.counter(), 3);
    /// ```
    pub fn new(iter: I) -> Self {
        Self {
            iter,
            counter: ShadowCounter::new(),
        }
    }
}

impl<'a, I: Iterator, S> ShadowCountedIter<'a, I, S> {
    /// Creates a new ShadowCountedIter from an iterator with specific supplemental data.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// #[derive(Clone, Default)]
    /// struct MyData { value: u32 }
    /// impl Commit for MyData {
    ///     fn commit(&self, _from: Self) {}
    /// }
    ///
    /// let mut iter = vec![1, 2, 3].into_iter();
    /// let mut iter = ShadowCountedIter::new_with(iter, MyData { value: 42 });
    /// assert_eq!(iter.supplement().value, 42);
    /// while let Some(_) = iter.next() {}
    /// assert_eq!(iter.counter(), 3);
    /// ```
    pub fn new_with(iter: I, supplement: S) -> Self {
        Self {
            iter,
            counter: ShadowCounter::new_with_data(supplement),
        }
    }

    /// Get a reference to the supplemental data.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// #[derive(Clone, Default)]
    /// struct MyData { value: u32 }
    /// impl Commit for MyData {
    ///     fn commit(&self, _from: Self) {}
    /// }
    ///
    /// let mut iter = ShadowCountedIter::new_with(vec![1, 2].into_iter(), MyData { value: 42 });
    /// assert_eq!(iter.supplement().value, 42);
    /// ```
    pub fn supplement(&self) -> &S {
        &self.counter.data
    }

    /// Get a mutable reference to the supplemental data.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// #[derive(Clone, Default)]
    /// struct MyData { value: u32 }
    /// impl Commit for MyData {
    ///     fn commit(&self, _from: Self) {}
    /// }
    ///
    /// let mut iter = ShadowCountedIter::new_with(vec![1, 2].into_iter(), MyData { value: 42 });
    /// iter.supplement_mut().value = 100;
    /// assert_eq!(iter.supplement().value, 100);
    /// ```
    pub fn supplement_mut(&mut self) -> &mut S {
        &mut self.counter.data
    }
}

impl<'a, I: Iterator, S: Clone> ShadowCountedIter<'a, I, S> {
    /// Creates a new nested `ShadowCountedIter` with a reference to a parent iterator.
    /// Nested iterators must be either committed to propagate their count to the parent.
    /// The supplemental data is cloned from the parent.
    ///
    /// # Arguments
    ///
    /// * `parent`    - The parent iterator to commit the count to.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let mut parent_iter = vec![1, 2, 3].into_iter().shadow_counted();
    /// parent_iter.next();
    /// # assert_eq!(parent_iter.counter(), 1);
    ///
    /// let mut nested_iter = vec![4,5].into_iter().nested_shadow_counted(&mut parent_iter);
    /// assert_eq!(nested_iter.counter(), 1);
    /// while let Some(_) = nested_iter.next() {}
    /// assert_eq!(nested_iter.counter(), 3);
    /// // destroy the nested iter while committing the count to the parent
    /// nested_iter.commit();
    /// # assert_eq!(parent_iter.counter(), 3);
    ///
    /// while let Some(_) = parent_iter.next() {}
    /// assert_eq!(parent_iter.counter(), 5);
    /// ```
    pub fn new_nested<'b, T: Iterator>(
        iter: I,
        parent: &'a mut ShadowCountedIter<'b, T, S>,
    ) -> Self {
        Self {
            iter,
            counter: parent.counter.nest(),
        }
    }
}

impl<'a, I: Iterator, S> ShadowCountedIter<'a, I, S> {
    /// Creates a new nested `ShadowCountedIter` with a reference to a parent iterator
    /// and custom supplemental data.
    /// Nested iterators must be either committed to propagate their count to the parent.
    ///
    /// # Arguments
    ///
    /// * `iter`       - The iterator to wrap
    /// * `parent`     - The parent iterator to commit the count to
    /// * `supplement` - Custom supplemental data for this nested iterator
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// use std::cell::Cell;
    ///
    /// #[derive(Clone, Default, Debug)]
    /// struct Counter { count: Cell<u32> }
    /// impl Commit for Counter {
    ///     fn commit(&self, from: Self) {
    ///         self.count.set(self.count.get() + from.count.get());
    ///     }
    /// }
    ///
    /// let mut parent = ShadowCountedIter::new_with(
    ///     vec![1, 2].into_iter(),
    ///     Counter { count: Cell::new(10) }
    /// );
    /// parent.next();
    ///
    /// let mut nested = ShadowCountedIter::new_nested_with(
    ///     vec![3, 4].into_iter(),
    ///     &mut parent,
    ///     Counter { count: Cell::new(100) }
    /// );
    /// assert_eq!(nested.supplement().count.get(), 100); // Custom data, not cloned from parent
    /// while let Some(_) = nested.next() {}
    /// nested.commit().unwrap();
    /// ```
    pub fn new_nested_with<'b, T: Iterator>(
        iter: I,
        parent: &'a mut ShadowCountedIter<'b, T, S>,
        supplement: S,
    ) -> Self
    where
        S: Clone,
    {
        Self {
            iter,
            counter: parent.counter.nest_with(supplement),
        }
    }
}

impl<'a, I: Iterator, S: Commit> ShadowCountedIter<'a, I, S> {
    /// Commits the count of a nested iterator to the parent iterator.
    /// This destroys `self` while committing the count to the parent.
    /// Only one child can commit to a parent in case these children got cloned.
    /// Any further commit will result in an error. Committing to a non nested (top level)
    /// iterator will error too. On error the iterator is returned.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let mut parent = vec![1, 2].into_iter().shadow_counted();
    /// parent.next();
    /// assert_eq!(parent.counter(), 1);
    ///
    /// let mut nested = vec![3, 4].into_iter().nested_shadow_counted(&mut parent);
    /// nested.next();
    /// nested.next();
    /// assert_eq!(nested.counter(), 3);
    ///
    /// // Commit propagates count to parent
    /// nested.commit().unwrap();
    /// assert_eq!(parent.counter(), 3);
    /// ```
    pub fn commit(self) -> Result<(), Self> {
        let Self { iter, counter } = self;
        counter.commit().map_err(|counter| Self { iter, counter })
    }

    /// This destructures the `ShadowCountedIter` and returns the iterator it wrapped.
    /// This is useful when the counter is no longer needed but the inner iterator is and one wants
    /// to shed the lifetime that comes with the `ShadowCountedIter`.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let mut iter = vec![1, 2, 3, 4, 5].into_iter().shadow_counted();
    /// iter.next(); // consume one
    ///
    /// let inner = iter.into_inner_iter();
    /// assert_eq!(inner.collect::<Vec<_>>(), vec![2, 3, 4, 5]);
    /// ```
    pub fn into_inner_iter(self) -> I {
        self.iter
    }
}

impl<I: Iterator, S> ShadowCountedIter<'_, I, S> {
    /// Allows to adjust the counter. This is required when in recursive structures only
    /// leafs shall be counted. Takes an isize so that it can be used with negative values.
    ///
    /// # Arguments
    ///
    /// * `delta` - The delta to adjust the counter by.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let mut iter = vec![1, 2, 3].into_iter().shadow_counted();
    /// iter.add(2);
    /// assert_eq!(iter.counter(), 2);
    /// iter.add(-1);
    /// assert_eq!(iter.counter(), 1);
    /// ```
    #[inline]
    pub fn add(&mut self, delta: isize) {
        self.counter.add(delta)
    }

    /// Returns the iterators current count.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let vec = vec![1, 2, 3];
    /// let mut iter = vec.into_iter().shadow_counted();
    /// while let Some(_) = iter.next() {}
    /// assert_eq!(iter.counter(), 3);
    /// ```
    #[inline]
    pub fn counter(&self) -> usize {
        self.counter.get()
    }
}

impl<I: Iterator, S> Iterator for ShadowCountedIter<'_, I, S> {
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        match self.iter.next() {
            Some(item) => {
                *self.counter.counter.get_mut() += 1;
                Some(item)
            }
            None => None,
        }
    }
}

impl<I: Iterator, S> AsRef<I> for ShadowCountedIter<'_, I, S> {
    fn as_ref(&self) -> &I {
        &self.iter
    }
}

/// Implements `ExactSizeIterator` when the inner iterator does.
/// The length reflects the remaining items in the iterator.
///
/// # Example
///
/// ```
/// # use shadow_counted::*;
/// let mut iter = vec![1, 2, 3, 4, 5].into_iter().shadow_counted();
/// assert_eq!(iter.len(), 5);
/// iter.next();
/// assert_eq!(iter.len(), 4);
/// iter.next();
/// assert_eq!(iter.len(), 3);
/// ```
impl<I: Iterator + ExactSizeIterator, S> ExactSizeIterator for ShadowCountedIter<'_, I, S> {
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<I: Iterator + std::iter::FusedIterator, S> std::iter::FusedIterator
    for ShadowCountedIter<'_, I, S>
{
}

/// Implements `DoubleEndedIterator` when the inner iterator does.
/// Both forward and backward iteration increment the counter.
///
/// # Example
///
/// ```
/// # use shadow_counted::*;
/// let mut iter = vec![1, 2, 3, 4, 5].into_iter().shadow_counted();
/// assert_eq!(iter.next(), Some(1));
/// assert_eq!(iter.counter(), 1);
/// assert_eq!(iter.next_back(), Some(5));
/// assert_eq!(iter.counter(), 2);
/// assert_eq!(iter.next_back(), Some(4));
/// assert_eq!(iter.counter(), 3);
/// ```
impl<I: DoubleEndedIterator, S> DoubleEndedIterator for ShadowCountedIter<'_, I, S> {
    fn next_back(&mut self) -> Option<Self::Item> {
        match self.iter.next_back() {
            Some(item) => {
                *self.counter.counter.get_mut() += 1;
                Some(item)
            }
            None => None,
        }
    }
}

/// A extension trait to convert any iterator into a ShadowCountedIter.
/// This must be in scope to use the `shadow_counted` method on an iterator.
/// When possible the `From` and `Into` may be more convenient to use.
///
/// # Example
///
/// ```
/// # use shadow_counted::*;
/// let vec = vec![1, 2, 3];
/// let mut iter = vec.into_iter().shadow_counted();
/// while let Some(_) = iter.next() {}
/// assert_eq!(iter.counter(), 3);
/// ```
pub trait IntoShadowCounted {
    /// Converts the iterator into a ShadowCountedIter with default (unit) supplemental data.
    fn shadow_counted<'a>(self) -> ShadowCountedIter<'a, Self, ()>
    where
        Self: Sized + Iterator,
    {
        ShadowCountedIter::new(self)
    }

    /// Converts the iterator into a nested ShadowCountedIter.
    ///
    /// # Arguments
    ///
    /// * `parent` - The parent iterator to commit the count to.
    ///
    /// # Example
    ///
    /// ```
    /// # use shadow_counted::*;
    /// let mut parent = vec![1, 2, 3].into_iter().shadow_counted();
    /// parent.next();
    ///
    /// let mut nested = vec![4, 5].into_iter().nested_shadow_counted(&mut parent);
    /// assert_eq!(nested.counter(), 1); // Inherits parent's count
    /// for _ in nested.by_ref() {}
    /// nested.commit().unwrap();
    /// assert_eq!(parent.counter(), 3); // Parent now has nested count
    /// ```
    fn nested_shadow_counted<'a, I: Iterator, S: Clone>(
        self,
        parent: &'a mut ShadowCountedIter<'_, I, S>,
    ) -> ShadowCountedIter<'a, Self, S>
    where
        Self: Sized + Iterator,
    {
        ShadowCountedIter::new_nested(self, parent)
    }
}

impl<T: Iterator> IntoShadowCounted for T {}

/// When types can be inferred then `From` and `Into` can be used to convert an iterator into a
/// `ShadowCountedIter` with default (unit) supplemental data.
impl<T: Iterator> From<T> for ShadowCountedIter<'_, T, ()> {
    fn from(iter: T) -> Self {
        ShadowCountedIter::new(iter)
    }
}

/// A shadow counter that can commit its count to a parent counter.
/// Carries supplemental data of type S.
#[derive(Debug, Clone)]
struct ShadowCounter<'a, S = ()> {
    /// holds the current count including counts from children who committed here
    counter: Cell<usize>,
    /// the parents counter when this child was created
    created_at: usize,
    /// parent for nested counters
    parent: Option<&'a ShadowCounter<'a, S>>,
    /// supplemental data
    data: S,
}

impl<'a, S: Clone> ShadowCounter<'a, S> {
    fn nest(&'a self) -> Self {
        Self {
            counter: self.counter.clone(),
            created_at: self.counter.get(),
            parent: Some(self),
            data: self.data.clone(),
        }
    }

    fn nest_with(&'a self, data: S) -> Self {
        Self {
            counter: self.counter.clone(),
            created_at: self.counter.get(),
            parent: Some(self),
            data,
        }
    }
}

impl<S: Default> ShadowCounter<'_, S> {
    fn new() -> Self {
        Self {
            counter: 0.into(),
            created_at: 0,
            parent: None,
            data: S::default(),
        }
    }
}

impl<S> ShadowCounter<'_, S> {
    fn new_with_data(data: S) -> Self {
        Self {
            counter: 0.into(),
            created_at: 0,
            parent: None,
            data,
        }
    }

    #[inline]
    fn get(&self) -> usize {
        self.counter.get()
    }

    #[inline]
    fn add(&mut self, delta: isize) {
        self.counter
            .set(self.counter.get().wrapping_add(delta as usize));
    }
}

impl<S: Commit> ShadowCounter<'_, S> {
    fn commit(self) -> Result<(), Self> {
        if let Some(p) = self.parent {
            if p.counter.get() == self.created_at {
                p.counter.set(self.counter.get());
                // Commit supplemental data to parent
                // Parent is responsible for interior mutability
                p.data.commit(self.data);
                return Ok(());
            }
        }
        Err(self)
    }
}

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

    #[test]
    fn test_empty() {
        let vec: Vec<i32> = vec![];
        let mut iter: ShadowCountedIter<'_, _, ()> = ShadowCountedIter::new(vec.into_iter());
        while iter.next().is_some() {}
        assert_eq!(iter.counter(), 0);
    }

    #[test]
    fn commit_nonnested() {
        let vec: Vec<i32> = vec![];
        let mut iter: ShadowCountedIter<'_, _, ()> = ShadowCountedIter::new(vec.into_iter());
        while iter.next().is_some() {}
        assert!(iter.commit().is_err());
    }

    #[test]
    fn test_basic_counting() {
        let vec = vec![1, 2, 3];
        let mut iter: ShadowCountedIter<'_, _, ()> = ShadowCountedIter::new(vec.into_iter());
        while iter.next().is_some() {}
        assert_eq!(iter.counter(), 3);
    }

    #[test]
    fn test_empty_iterator() {
        let vec: Vec<i32> = vec![];
        let mut iter: ShadowCountedIter<'_, _, ()> = ShadowCountedIter::new(vec.into_iter());
        while iter.next().is_some() {}
        assert_eq!(iter.counter(), 0);
    }

    #[test]
    fn test_commit() {
        let mut parent_iter = vec![1, 2, 3].into_iter().shadow_counted();
        parent_iter.next();
        let mut nested_iter = vec![4, 5]
            .into_iter()
            .nested_shadow_counted(&mut parent_iter);
        nested_iter.next();
        nested_iter.commit().unwrap();
        assert_eq!(parent_iter.counter(), 2);
    }

    #[test]
    fn test_commit_twice() {
        let mut parent_iter = vec![1, 2, 3].into_iter().shadow_counted();
        parent_iter.next();
        let mut nested_iter = vec![4, 5]
            .into_iter()
            .nested_shadow_counted(&mut parent_iter);
        let nested_iter2 = nested_iter.clone();
        nested_iter.next();
        nested_iter.commit().unwrap();
        assert!(nested_iter2.commit().is_err());
        assert_eq!(parent_iter.counter(), 2);
    }

    #[test]
    fn test_two_commits() {
        let mut parent_iter = vec![1, 2, 3].into_iter().shadow_counted();
        parent_iter.next();
        let mut nested_iter = vec![4].into_iter().nested_shadow_counted(&mut parent_iter);
        nested_iter.next();
        nested_iter.commit().unwrap();
        parent_iter.next();

        // next nested iter
        let mut nested_iter = vec![5].into_iter().nested_shadow_counted(&mut parent_iter);
        nested_iter.next();
        nested_iter.commit().unwrap();

        parent_iter.next();
        assert_eq!(parent_iter.counter(), 5);
    }

    #[test]
    fn test_shadow_count_iter() {
        #[derive(Debug, PartialEq)]
        enum Nodes<'a, T> {
            Leaf(T),
            Nested(&'a [Nodes<'a, T>]),
        }

        let items = &[
            Nodes::Leaf(1),
            Nodes::Leaf(2),
            Nodes::Nested(&[Nodes::Leaf(3), Nodes::Leaf(4), Nodes::Leaf(5)]),
            Nodes::Leaf(6),
        ];

        let mut sc_iter = items.iter().shadow_counted();

        assert!((sc_iter.counter()) == (0));
        assert!((sc_iter.next()) == (Some(&Nodes::Leaf(1))));
        assert!((sc_iter.counter()) == (1));
        assert!((sc_iter.next()) == (Some(&Nodes::Leaf(2))));
        assert!((sc_iter.counter()) == (2));

        let nested = sc_iter.next().unwrap();
        assert!((sc_iter.counter()) == (3));
        assert!((nested) == (&Nodes::Nested(&[Nodes::Leaf(3), Nodes::Leaf(4), Nodes::Leaf(5)])));
        let Nodes::Nested(nested) = nested else {
            panic!()
        };
        let mut nested_iter = nested.iter().nested_shadow_counted(&mut sc_iter);
        assert_eq!(nested_iter.counter(), 3);
        assert_eq!(nested_iter.next(), Some(&Nodes::Leaf(3)));
        assert_eq!(nested_iter.counter(), 4);
        assert_eq!(nested_iter.next(), Some(&Nodes::Leaf(4)));
        assert_eq!(nested_iter.counter(), 5);
        assert_eq!(nested_iter.next(), Some(&Nodes::Leaf(5)));
        assert_eq!(nested_iter.counter(), 6);
        assert_eq!(nested_iter.next(), None);
        assert_eq!(nested_iter.counter(), 6);
        nested_iter.commit().unwrap();
        assert_eq!(sc_iter.counter(), 6);
        assert_eq!(sc_iter.next(), Some(&Nodes::Leaf(6)));
        assert_eq!(sc_iter.counter(), 7);
    }

    #[test]
    fn test_supplement_with_custom_data() {
        #[derive(Debug, Clone, Default, PartialEq)]
        struct MyData {
            value: u32,
        }
        impl Commit for MyData {
            fn commit(&self, _from: Self) {}
        }

        let vec = vec![1, 2, 3];
        let data = MyData { value: 42 };
        let mut iter = ShadowCountedIter::new_with(vec.into_iter(), data);
        assert_eq!(iter.supplement().value, 42);
        while iter.next().is_some() {}
        assert_eq!(iter.counter(), 3);
        assert_eq!(iter.supplement().value, 42);
    }

    #[test]
    fn test_supplement_mut() {
        #[derive(Debug, Clone, Default, PartialEq)]
        struct MyData {
            value: u32,
        }
        impl Commit for MyData {
            fn commit(&self, _from: Self) {}
        }

        let mut iter = ShadowCountedIter::new_with(vec![1, 2].into_iter(), MyData { value: 10 });
        assert_eq!(iter.supplement().value, 10);
        iter.supplement_mut().value = 20;
        assert_eq!(iter.supplement().value, 20);
    }

    #[test]
    fn test_nested_data_cloned() {
        #[derive(Debug, Clone, Default, PartialEq)]
        struct MyData {
            value: u32,
        }
        impl Commit for MyData {
            fn commit(&self, _from: Self) {}
        }

        let mut parent = ShadowCountedIter::new_with(vec![1].into_iter(), MyData { value: 42 });

        let nested = vec![2].into_iter().nested_shadow_counted(&mut parent);
        assert_eq!(nested.supplement().value, 42); // Cloned from parent
    }

    #[test]
    fn test_data_committed() {
        use std::cell::Cell;

        #[derive(Debug, Clone, Default)]
        struct Counter {
            count: Cell<u32>,
        }

        impl Commit for Counter {
            fn commit(&self, from: Self) {
                self.count.set(self.count.get() + from.count.get());
            }
        }

        let mut parent = ShadowCountedIter::new_with(vec![1, 2].into_iter(), Counter::default());
        parent.next();

        let mut nested = vec![3, 4].into_iter().nested_shadow_counted(&mut parent);
        nested.supplement_mut().count.set(5);
        nested.next();
        nested.next();
        nested.commit().unwrap();

        assert_eq!(parent.supplement().count.get(), 5); // Committed from nested
    }
}