Mercurial > repos > alg_tools / file revision

/// Iteration utilities; [stateful][StatefulIterator] and
/// [restartable][RestartableIterator] iterators.

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 }
    }
}

/*
/// This is a helper struct returning iterators with a lifetime from trait functions.
/// `T` will typically be `&'a Self` or some other reference that can only be
/// specified in the trait function declaration without having to parametrise the entire
/// trait with `'a`. The parameter `ImplData` is supposed to be a unique type that idenfiers
/// the iterator for the implementation of the trait. Then `Iterator` is to be implemented
/// for `ProxyIterator<T, ImplState>`. For example:
/// TODO: not updated.
/// ```
/// trait Test {
///    type ImplData;
///    fn returns_iterator<'a>(&'a self) -> ProxyIterator<&'a Self, Self::ImplData>;
/// }
/// struct MyIndex{
///     index : usize,
/// }
/// impl<const N : usize> Test for [f64; N] {
///     type ImplData = MyIndex;
///    fn returns_iterator<'a>(&'a self) -> ProxyIterator<&'a Self, MyIndex> {
///        TraitIterator::new(self, MyIndex{ index : 0 })
///    }
/// }
/// impl<'a, const N : usize> Iterator for `TraitIterator<&'a [f64; N], MyIndex> {
///     type Item = f64;
///     fn next(&mut self) -> Option<f64> {
///         if self.state.index < N {
///             let v = self.data[self.state.index];
///             self.state.index += 1;
///             Some(v)
///         } else {
///             None
///         }
///     }
/// }
/// ```
*/

macro_rules! impl_trait_iterator {
    ($name:ident, $proxyname:ident, &'a $($mut:ident)?, $($itemref:tt)*) => {
        pub trait $name {
            type Item : 'static;
            type Data : ?Sized;
            fn next_with<'a>(&mut self, data : &'a $($mut)? Self::Data) -> Option<$($itemref)* Self::Item>;
        }

        pub struct $proxyname<'a, Impl>
        where Impl : $name {
            pub data :  &'a $($mut)? Impl::Data,
            pub iter : Impl,
        }

        impl<'a, Impl> $proxyname<'a, Impl>
        where Impl : $name {
            #[inline]
            pub fn new(data : &'a $($mut)? Impl::Data, iter : Impl) -> $proxyname<'a, Impl> {
                $proxyname{ data : data, iter : iter}
            }
        }

        impl<'a, Impl> Iterator for $proxyname<'a, Impl>
        where Impl : $name {
            type Item = $($itemref)* Impl::Item;
            #[inline]
            fn next(&mut self) -> Option<Self::Item> {
                self.iter.next_with(self.data)
            }
        }
    }
}

impl_trait_iterator!(TraitIterator,    ProxyIterator,    &'a,     );
impl_trait_iterator!(TraitIteratorRef, ProxyIteratorRef, &'a,     &'a);
// TODO: Lifetime errors. Probably due to potential overlapping mutable borrows.
// So should use some unsafe pointer hacks or so.
//impl_trait_iterator!(TraitIteratorMut, ProxyIteratorMut, &'a mut, &'a mut);

/*
pub struct FlattenTI<SI>
where SI : TraitIterator,
      SI::Item : Iterator {
    iter_iter : SI,
    element_iter : Option<SI::Item>,
}

impl<SI> FlattenTI<SI>
where SI : TraitIterator,
      SI::Item : Iterator {
    fn new(iter : SI) -> Self {
        FlattenTI{ iter_iter : iter, element_iter : None  }
    }
}

impl<SI : TraitIterator> TraitIterator for FlattenTI<SI>
where SI::Item : Iterator {
    type Item = <SI::Item as Iterator>::Item;
    type Data = SI::Data;

    fn next_with<'a>(&mut self, data : &'a $($mut)? SI::Data) -> Option<Self::Item> {
        loop {
            if let Some(ei) = self.element_iter {
                if let n @ Some(_) = ei.next() {
                    return n
                }
            } else {
                self.element_iter = self.iter_iter.next_with(data);
                if let None = self.element_iter {
                    return None
                }
            }
        }
    }
}
*/

/*
pub enum EitherIterator<T, L : Iterator<Item=T>, R : Iterator<Item=T>> {
    Left(L),
    Right(R)
}

impl<T, L : Iterator<Item=T>, R : Iterator<Item=T>> Iterator for EitherIterator<T, L, R> {
    type Item = T;
    #[inline]
    fn next(&mut self) -> Option<T> {
        match self {
            EitherIterator::Left(ref l)  => { l.next() }
            EitherIterator::Right(ref r) => { r.next() }
        }
    }
}
*/

/// 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
        }
    }
}


pub struct RestartableIter<'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> RestartableIter<'a, T> {
    fn map_result(result : Option<*const T>) -> Option<&'a T> {
        match result {
            None        => { None }
            Some(ptr)   => { Some(unsafe{ &*ptr }) }
        }
    }

    pub fn new(slice : &'a [T]) -> Self {
        let ptr_range = slice.as_ptr_range();
        let inner = RangeIter::new(ptr_range.start, ptr_range.end);
        RestartableIter{ inner : inner, _phantom : PhantomData }
    }
}

impl<'a, T : 'a> Iterator for RestartableIter<'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 RestartableIter<'a, T> {
    #[inline]
    fn len(&self) -> usize { self.inner.len() }
}

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

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

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())
    }
}