rune/src/iter.rs

//! An iterator adapter for arbitrary lookahead functionality.
//!
//! This module provides [`Lookahead`], an adapter for any iterator that allows
//! you to peek ahead by any number of elements, without consuming them.
//!
//! ## Example
//! ```
//! use rune::LookaheadExt;
//!
//! let mut it = vec![10, 20, 30].into_iter().lookahead();
//!
//! assert_eq!(it.peek(0), Some(&10));
//! assert_eq!(it.peek(1), Some(&20));
//! assert_eq!(it.next(), Some(10));
//! assert_eq!(it.peek(0), Some(&20));
//! ```

use std::collections::VecDeque;
use std::fmt;
use std::iter::{Fuse, FusedIterator};



/// An iterator adapter that allows arbitrary lookahead peeking.
///
/// This struct wraps an iterator and buffers items so that any future
/// item can be accessed by index without consuming them. Similar to
/// [`std::iter::Peekable`], but supports peeking any number of steps ahead.
pub struct Lookahead<I>
   where I: Iterator
{
   iter: Fuse<I>,
   buffer: VecDeque<I::Item>
}



impl<I> Lookahead<I> where I: Iterator
{
   /// Creates a new [`Lookahead`] from the given iterator.
   ///
   /// This constructor is typically used indirectly via the
   /// [`LookaheadExt::lookahead()`] method or [`lookahead()`] free function.
   #[must_use]
   pub fn new(iter: I) -> Self
   {
      Lookahead { iter: iter.fuse(),
                  buffer: VecDeque::new() }
   }

   /// Returns a reference to the `n`th upcoming item, if it exists.
   ///
   /// `peek(0)` is the same as peeking at the next item.
   ///
   /// This does **not consume** any items from the iterator.
   ///
   /// # Examples
   /// ```
   /// use rune::LookaheadExt;
   ///
   /// let mut it = vec![1, 2, 3].into_iter().lookahead();
   ///
   /// assert_eq!(it.peek(1), Some(&2));
   /// assert_eq!(it.next(), Some(1));
   /// ```
   pub fn peek(&mut self, n: usize) -> Option<&I::Item>
   {
      while self.buffer.len() <= n
      {
         if let Some(item) = self.iter.next()
         {
            self.buffer.push_back(item);
         }
         else
         {
            break;
         }
      }
      self.buffer.get(n)
   }

   /// Returns a mutable reference to the `n`th upcoming item, if it exists.
   ///
   /// This allows in-place modification of peeked items before consumption.
   ///
   /// # Examples
   /// ```
   /// use rune::LookaheadExt;
   ///
   /// let mut it = vec![1, 2, 3].into_iter().lookahead();
   /// if let Some(x) = it.peek_mut(1)
   /// {
   ///    *x *= 10;
   /// }
   /// assert_eq!(it.next(), Some(1));
   /// assert_eq!(it.next(), Some(20));
   /// ```
   pub fn peek_mut(&mut self, n: usize) -> Option<&mut I::Item>
   {
      while self.buffer.len() <= n
      {
         if let Some(item) = self.iter.next()
         {
            self.buffer.push_back(item);
         }
         else
         {
            break;
         }
      }
      self.buffer.get_mut(n)
   }
}

impl<I> Iterator for Lookahead<I> where I: Iterator
{
   type Item = I::Item;

   /// Retrieves the next item, consuming it.
   ///
   /// If any items were previously peeked and buffered, they are returned
   /// first before accessing the underlying iterator.
   fn next(&mut self) -> Option<Self::Item>
   {
      if let Some(front) = self.buffer.pop_front()
      {
         Some(front)
      }
      else
      {
         self.iter.next()
      }
   }

   /// Provides a size hint accounting for both buffered and remaining elements.
   fn size_hint(&self) -> (usize, Option<usize>)
   {
      let (low, high) = self.iter.size_hint();
      let buffered = self.buffer.len();
      (low.saturating_add(buffered), high.and_then(|h| h.checked_add(buffered)))
   }
}

impl<I> Clone for Lookahead<I>
   where I: Iterator + Clone,
         I::Item: Clone
{
   fn clone(&self) -> Self
   {
      Lookahead { iter: self.iter.clone(),
                  buffer: self.buffer.clone() }
   }
}

impl<I> fmt::Debug for Lookahead<I>
   where I: Iterator + fmt::Debug,
         I::Item: fmt::Debug
{
   fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
   {
      f.debug_struct("Lookahead")
       .field("iter", &self.iter)
       .field("buffer", &self.buffer)
       .finish()
   }
}

impl<I> FusedIterator for Lookahead<I> where I: Iterator + FusedIterator {}



/// Extension trait to provide `.lookahead()` on all iterators.
///
/// This lets you easily call `.lookahead()` on any iterator to
/// create a [`Lookahead`] instance.
pub trait LookaheadExt: Iterator + Sized
{
   /// Wraps the iterator in a [`Lookahead`] adapter.
   fn lookahead(self) -> Lookahead<Self>;
}

impl<I: Iterator> LookaheadExt for I
{
   fn lookahead(self) -> Lookahead<Self>
   {
      Lookahead::new(self)
   }
}



/// Creates a [`Lookahead`] from any iterable.
///
/// This is a convenience function for use in functional-style code or
/// when not using the extension trait.
///
/// # Example
/// ```
/// use rune::lookahead;
///
/// let mut it = lookahead(vec![1, 2, 3]);
///
/// assert_eq!(it.peek(2), Some(&3));
/// ```
pub fn lookahead<I>(iterable: I) -> Lookahead<I::IntoIter>
   where I: IntoIterator
{
   Lookahead::new(iterable.into_iter())
}