Mercurial > repos > alg_tools / file revision

/*!
Iteration utilities.

This includes
 * [Stateful][StatefulIterator] and [restartable][RestartableIterator] iterators.
 * Variants of [`Iterator::map`] and [`Iterator::filter_map`], etc., that work on functions
   instead of closures, forgoing exact closure type signature specification; see [`Mappable`].

*/

use crate::types::Integer;
use std::marker::PhantomData;

/// Trait for iterators that can be queried the current state without advancing.
pub trait StatefulIterator : Iterator {
    /// Return current value of iterator. If the iterator has not yet started,
    /// `None` should be returned.
    fn current(&self) -> Option<Self::Item>;
}

/// Iterators that can be restarted and queried the current state without
/// advancing to next state.
pub trait RestartableIterator : StatefulIterator {
    /// Restart the iterator. Return the first item.
    fn restart(&mut self) -> Option<Self::Item>;
}

/// A variant of `std::iter::Map` that works for methods instead of closures,
/// so can be returned from functions without `dyn`.
/// Use [`Mappable::mapF`] to create this iterator.
pub struct MapF<I,J> where I : Iterator {
    iter : I,
    func : fn(I::Item) -> J,
}

impl<I,J> Iterator for MapF<I,J> where I : Iterator {
    type Item = J;
    fn next(&mut self) -> Option<J> {
        self.iter.next().map(|item| (self.func)(item))
    }
}

/// A variant of [`MapF`] with extra data for the function stored as a reference.
/// See also [`MapZ`] and use [`Mappable::mapX`] to create this iteartor.
pub struct MapX<'a,I,T,J> where I : Iterator {
    iter : I,
    data : &'a T,
    func : fn(&'a T, I::Item) -> J,
}

impl<'a,I,T,J> Iterator for MapX<'a,I,T,J> where I : Iterator {
    type Item = J;
    fn next(&mut self) -> Option<J> {
        self.iter.next().map(|item| (self.func)(self.data, item))
    }
}

/// A variant of [`MapF`] and [`MapX`] with extra data for the function stored in the struct.
/// Use [`Mappable::mapZ`] to create this iterator.
pub struct MapZ<I,T,J> where I : Iterator {
    iter : I,
    data : T,
    func : fn(&T, I::Item) -> J,
}

impl<'a,I,T,J> Iterator for MapZ<I,T,J> where I : Iterator {
    type Item = J;
    fn next(&mut self) -> Option<J> {
        self.iter.next().map(|item| (self.func)(&self.data, item))
    }
}

/// A variant of `std::iter::FilterMap` that works for methods instead of closures,
/// so can be returned from functions without `dyn`.
pub struct FilterMapX<'a,I,T,J> where I : Iterator {
    iter : I,
    data : &'a T,
    func : fn(&'a T, I::Item) -> Option<J>,
}

impl<'a,I,T,J> Iterator for FilterMapX<'a,I,T,J> where I : Iterator {
    type Item = J;
    fn next(&mut self) -> Option<J> {
        while let Some(item) = self.iter.next() {
            let v = (self.func)(self.data, item);
            if v.is_some() {
                return v
            }
        }
        None
    }
}

/// Helper for [`Mappable::filter_zip`].
pub struct FilterZip<I, J>
where I : Iterator,
      J : Iterator<Item=bool> {
    iter : I,
    filter : J,
}

impl<I,J> Iterator for FilterZip<I, J>
where I : Iterator,
      J : Iterator<Item=bool> {
    type Item = I::Item;
    #[inline]
    fn next(&mut self) -> Option<I::Item> {
        while let (v@Some(..), Some(filt)) = (self.iter.next(), self.filter.next()) {
            if filt {
                return v
            }
        }
        None
    }
}

/// This trait implements several mapping methods on iterators.
pub trait Mappable : Iterator + Sized {
    /// Apply `func` to the output of the iterator.
    #[allow(non_snake_case)]
    fn mapF<J>(self, func : fn(Self::Item) -> J) -> MapF<Self,J>;
    /// Apply `func` to the output of the iterator, passing also `d` to it.
    #[allow(non_snake_case)]
    fn mapX<'a,T,J>(self, d : &'a T, func : fn(&'a T, Self::Item) -> J) -> MapX<'a,Self,T,J>;
    /// Apply `func` to the output of the iterator, passing also `d` to it.
    #[allow(non_snake_case)]
    fn mapZ<T,J>(self, d : T, func : fn(&T, Self::Item) -> J) -> MapZ<Self,T,J>;
    /// Apply `func` to the output of the iterator, filtering based on the result.
    /// The data `d` is passed to `func`.
    #[allow(non_snake_case)]
    fn filter_mapX<'a,T,J>(self, d : &'a T, func : fn(&'a T, Self::Item) -> Option<J>) -> FilterMapX<'a,Self,T,J>;

    /// Filter `self` based on another iterator `iter` with `bool` items.
    fn filter_zip<J>(self, filter : J) -> FilterZip<Self, J> where J : Iterator<Item=bool>;
}

impl<I> Mappable for I where I : Iterator + Sized {
    #[inline]
    fn mapF<J>(self, func : fn(Self::Item) -> J) -> MapF<Self,J> {
        MapF{ iter : self, func : func }
    }
    #[inline]
    fn mapX<'a,T,J>(self, d : &'a T, func : fn(&'a T, Self::Item) -> J) -> MapX<'a,Self,T,J> {
        MapX{ iter : self, data : d, func : func }
    }
    #[inline]
    fn mapZ<T,J>(self, d : T, func : fn(&T, Self::Item) -> J) -> MapZ<Self,T,J> {
        MapZ{ iter : self, data : d, func : func }
    }
    #[inline]
    fn filter_mapX<'a,T,J>(self, d : &'a T, func : fn(&'a T, Self::Item) -> Option<J>)
    -> FilterMapX<'a,Self,T,J> {
        FilterMapX{ iter : self, data : d, func : func }
    }

    #[inline]
    fn filter_zip<J>(self, filter : J) -> FilterZip<Self, J>
    where J : Iterator<Item=bool> {
        FilterZip{ iter : self, filter : filter }
    }
}

/// A [`RestartableIterator`] over the range ´[start, end)`.
#[derive(Clone,Copy,Debug)]
pub struct RangeIter<T> {
    start : T,
    end : T,
    value : Option<T>,
}

pub trait StepNext : Ord + Copy {
    fn step_next(&self) -> Self;
    fn dist_to(&self, other : &Self) -> usize;
}

impl<T : Integer + Into<usize>> StepNext for T {
    #[inline]
    fn step_next(&self) -> Self { *self + Self::ONE }
    
    #[inline]
    fn dist_to(&self, other : &Self) -> usize {
        if *other >= *self {
             (*other - *self).into()
        } else {
            0
        }
    }
}

impl<T : StepNext> RangeIter<T> {
    /// Create a new [`RangeIter`].
    pub fn new(start : T, end : T) -> Self {
        RangeIter{ start : start.min(end), end : end, value : None }
    }
}

impl<T : StepNext> Iterator for RangeIter<T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.value {
            None => {
                if self.start < self.end {
                    self.value = Some(self.start);
                    self.value
                } else {
                    None
                }
            }
            Some(v) => {
                let vn = v.step_next();
                if vn < self.end {
                    self.value = Some(vn);
                    self.value
                } else {
                    None
                }
            }
        }
    }
}

impl<T : StepNext> ExactSizeIterator for RangeIter<T> {
    #[inline]
    fn len(&self) -> usize { self.start.dist_to(&self.end) }
}

impl<T : StepNext> StatefulIterator for RangeIter<T> {
    #[inline]
    fn current(&self) -> Option<Self::Item> {
        self.value
    }
}

impl<T : StepNext> RestartableIterator for RangeIter<T> {
    fn restart(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            self.value = Some(self.start);
            self.value
        } else {
            None
        }
    }
}

/// A restartable slice iterator.
pub struct RestartableSliceIter<'a, T> {
    inner : RangeIter<*const T>,
    _phantom : PhantomData<&'a[T]>
}

impl<T> StepNext for *const T {
    #[inline]
    fn step_next(&self) -> Self { unsafe { self.add(1) } }
    
    #[inline]
    fn dist_to(&self, other : &Self) -> usize {
        unsafe { other.offset_from(*self).max(0) as usize }
    }
}

impl<'a, T : 'a> RestartableSliceIter<'a, T> {
    /// Converts `Some` pointer to `Some` reference
    fn map_result(result : Option<*const T>) -> Option<&'a T> {
        match result {
            None        => { None }
            Some(ptr)   => { Some(unsafe{ &*ptr }) }
        }
    }

    /// Creates a restartable iterator over `slice`.
    pub fn new(slice : &'a [T]) -> Self {
        let ptr_range = slice.as_ptr_range();
        let inner = RangeIter::new(ptr_range.start, ptr_range.end);
        RestartableSliceIter{ inner : inner, _phantom : PhantomData }
    }
}

impl<'a, T : 'a> Iterator for RestartableSliceIter<'a, T> {
    type Item = &'a T;
    #[inline]
    fn next(&mut self) -> Option<&'a T> {
        Self::map_result(self.inner.next())
    }
}

impl<'a, T : 'a> ExactSizeIterator for RestartableSliceIter<'a, T> {
    #[inline]
    fn len(&self) -> usize { self.inner.len() }
}

impl<'a, T : 'a> StatefulIterator for RestartableSliceIter<'a, T> {
    #[inline]
    fn current(&self) -> Option<Self::Item> {
        Self::map_result(self.inner.current())
    }
}

impl<'a, T : 'a> RestartableIterator for RestartableSliceIter<'a, T> {
    fn restart(&mut self) -> Option<Self::Item> {
        Self::map_result(self.inner.current())
    }
}

/// A restartable variant of [`std::iter::Cloned`].
pub struct RestartableCloned<I : Iterator> {
    inner : I,
}

impl<'a, T : Clone + 'a , I : Iterator<Item=&'a T>> RestartableCloned<I> {
    fn map_result(result : Option<I::Item>) -> Option<T> {
        match result {
            None    => { None }
            Some(v) => { Some(v.clone()) }
        }
    }

    pub fn new(inner : I) -> Self {
        RestartableCloned{ inner : inner }
    }
}

impl<'a, T : Clone + 'a , I : Iterator<Item=&'a T>>  Iterator
for RestartableCloned<I> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        Self::map_result(self.inner.next())
    }
}

impl<'a, T : Clone + 'a , I : ExactSizeIterator<Item=&'a T>> ExactSizeIterator
for RestartableCloned<I> {
    fn len(&self) -> usize {
        self.inner.len()
    }
}

impl<'a, T : Clone + 'a , I : StatefulIterator<Item=&'a T>>  StatefulIterator
for RestartableCloned<I> {
    fn current(&self)  -> Option<Self::Item> {
        Self::map_result(self.inner.current())
    }
}

impl<'a, T : Clone + 'a , I : RestartableIterator<Item=&'a T>> RestartableIterator
for RestartableCloned<I> {
    fn restart(&mut self) -> Option<Self::Item> {
        Self::map_result(self.inner.restart())
    }
}