stacking_iterator/

lib.rs

//! Iterator utilities for manipulating stacks.
//!
//! This library is intended as an alternative to "borrowing" iterators for cases where collections
//! are formed using a series of stack operations.
//!
//! For example, when enumerating the accepted sequences in a finite-state automation, entering a
//! state takes the form of a push, whereas exiting a state takes the form of a pop. If you reach an
//! ending state at any point, the stack of transition keys can be read out as an accepted sequence.
//!
//! While "pushing" to sequences takes the form of the built-in [`Extend`], "popping" from sequences
//! is done via the [`Contract`] trait provided by this crate. If you offer an iterator over
//! [`Instruction`] items, you can then take advantage of the crate's functionality using the
//! [`StackingIteratorExt`] extension trait.
#![warn(missing_docs)]
#![warn(missing_debug_implementations)]
#![deny(unsafe_code)]
#![cfg_attr(not(any(feature = "std", test, doc)), no_std)]

#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::{
    collections::{LinkedList, VecDeque},
    string::String,
    vec::Vec,
};
use core::fmt;

/// Prelude for this crate, to be glob-imported in relevant code.
///
/// This re-exports the [`Contract`] and [`StackingIteratorExt`] traits in addition to the
/// [`Instruction`] and [`Operation`] enums.
pub mod prelude {
    pub use crate::{Contract, Instruction, Operation, StackingIteratorExt};
}

/// Control flow involving a stack.
///
/// Instructions are used to convert an iterator of [`Operation`]s into an iterator of stack contents.
#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub enum Instruction<T> {
    /// Mutate the stack with the given operation.
    Mutate(Operation<T>),

    /// Yield the contents of the stack.
    Yield,
}
impl<T> Instruction<T> {
    /// Shorthand for wrapping a [`Pop`] operation in [`Mutate`].
    ///
    /// [`Pop`]: Operation::Pop
    /// [`Mutate`]: Instruction::Mutate
    #[allow(non_upper_case_globals)]
    pub const Pop: Instruction<T> = Instruction::Mutate(Operation::Pop);

    /// Shorthand for wrapping a [`Push`] operation in [`Mutate`].
    ///
    /// [`Push`]: Operation::Push
    /// [`Mutate`]: Instruction::Mutate
    #[allow(non_snake_case)]
    pub const fn Push(val: T) -> Instruction<T> {
        Instruction::Mutate(Operation::Push(val))
    }

    /// Returns the operation, if mutating.
    ///
    /// # Examples
    ///
    /// ```
    /// use stacking_iterator::prelude::*;
    ///
    /// assert_eq!(Instruction::Push(1).into_operation(), Some(Operation::Push(1)));
    /// assert_eq!(<Instruction<()>>::Pop.into_operation(), Some(Operation::Pop));
    /// assert_eq!(<Instruction<()>>::Yield.into_operation(), None);
    /// ```
    #[inline]
    pub fn into_operation(self) -> Option<Operation<T>> {
        match self {
            Instruction::Mutate(operation) => Some(operation),
            Instruction::Yield => None,
        }
    }

    /// Returns a reference to the operation, if mutating.
    ///
    /// # Examples
    ///
    /// ```
    /// use stacking_iterator::prelude::*;
    ///
    /// assert_eq!(Instruction::Push(1).as_operation(), Some(&Operation::Push(1)));
    /// assert_eq!(<Instruction<()>>::Pop.as_operation(), Some(&Operation::Pop));
    /// assert_eq!(<Instruction<()>>::Yield.as_operation(), None);
    /// ```
    #[inline]
    pub fn as_operation(&self) -> Option<&Operation<T>> {
        match self {
            Instruction::Mutate(operation) => Some(operation),
            Instruction::Yield => None,
        }
    }
}
impl<T: fmt::Debug> fmt::Debug for Instruction<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Instruction::Mutate(operation) => fmt::Debug::fmt(operation, f),
            Instruction::Yield => f.pad("Yield"),
        }
    }
}

/// Stack operation.
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum Operation<T> {
    /// Push an item onto the stack.
    Push(T),

    /// Pop an item from the stack.
    Pop,
}
impl<T> Operation<T> {
    /// Returns the pushed item, if pushing.
    ///
    /// # Examples
    ///
    /// ```
    /// use stacking_iterator::prelude::*;
    ///
    /// assert_eq!(Operation::Push(1).into_item(), Some(1));
    /// assert_eq!(<Operation<()>>::Pop.into_item(), None);
    /// ```
    #[inline]
    pub fn into_item(self) -> Option<T> {
        match self {
            Operation::Push(item) => Some(item),
            Operation::Pop => None,
        }
    }

    /// Returns a reference to the pushed item, if pushing.
    ///
    /// # Examples
    ///
    /// ```
    /// use stacking_iterator::prelude::*;
    ///
    /// assert_eq!(Operation::Push(1).as_item(), Some(&1));
    /// assert_eq!(<Operation<()>>::Pop.as_item(), None);
    /// ```
    #[inline]
    pub fn as_item(&self) -> Option<&T> {
        match self {
            Operation::Push(item) => Some(item),
            Operation::Pop => None,
        }
    }
}

/// Opposite of [`Extend`].
pub trait Contract {
    /// Contracts the collection, removing the given number of items from its end.
    ///
    /// # Overflow Behavior
    ///
    /// This method *may* panic if `count` is larger than the size of the collection.
    /// Since this check may incur performance penalities, it is only guaranteed if
    /// debug assertions are enabled.
    ///
    /// In these cases, the collection must still correctly remove the elements, but
    /// may decide to not notify the user that less than the count of items were removed.
    fn contract(&mut self, count: usize);
}

fn stack_extend<T>(
    stack: &mut (impl Extend<T> + Contract),
    mut iter: impl Iterator<Item = Operation<T>>,
) {
    // the iterator will be a series of pushes and pops
    loop {
        // use native extend for pushes;
        // keep track of whether we end on a pop, or `None`
        let mut pop_count = 0;
        stack.extend(iter.by_ref().map_while(|op| match op {
            Operation::Push(val) => Some(val),
            Operation::Pop => {
                pop_count = 1;
                None
            }
        }));

        // if we end on None, break out of loop
        if pop_count == 0 {
            return;
        }

        // use native count for pops;
        // keep track of whether we end on a push, or `None`
        let mut next = None;

        // we have to use filter_map here to ensure we take the values by-value
        pop_count += iter
            .by_ref()
            .map_while(|op| match op {
                Operation::Pop => Some(()),
                Operation::Push(val) => {
                    next = Some(val);
                    None
                }
            })
            .count();
        stack.contract(pop_count);

        // if we end on None, break out of loop
        if let Some(val) = next {
            stack.extend(Some(val));
        } else {
            return;
        }
    }
}

/// Panic when contracting more than possible.
#[cold]
#[track_caller]
fn panic_imploded(count: usize, len: usize) -> ! {
    if len == 1 {
        panic!("cannot contract by {count} when collection has 1 element");
    } else {
        panic!("cannot contract by {count} when collection has {len} elements")
    }
}

/// Pops items off the [`Vec`].
#[cfg(feature = "alloc")]
impl<T> Contract for Vec<T> {
    #[track_caller]
    fn contract(&mut self, count: usize) {
        if count > self.len() {
            panic_imploded(count, self.len())
        }
        self.truncate(self.len() - count);
    }
}
/// Pops items off the back of the [`VecDeque`].
#[cfg(feature = "alloc")]
impl<T> Contract for VecDeque<T> {
    #[track_caller]
    fn contract(&mut self, count: usize) {
        if count > self.len() {
            panic_imploded(count, self.len())
        }
        self.truncate(self.len() - count);
    }
}
/// Pops items off the back of the [`LinkedList`].
#[cfg(feature = "alloc")]
impl<T> Contract for LinkedList<T> {
    #[track_caller]
    fn contract(&mut self, count: usize) {
        // this is still useful here, since otherwise an excessively large `count` could incur performance penalties
        if count > self.len() {
            panic_imploded(count, self.len())
        }
        // we don't optimize for count == self.len() since we need to drop all the items anyway
        for _ in 0..count {
            self.pop_back();
        }
    }
}

/// Pops some number of characters off the end of the string.
#[cfg(feature = "alloc")]
impl Contract for String {
    #[track_caller]
    fn contract(&mut self, count: usize) {
        let mut chars = self.chars();

        // it's slower, but manually counting the number of characters excess is nice to have in debug mode
        if cfg!(debug_assertions) {
            let mut excess = 0;
            let mut amount = 0;
            for _ in 0..count {
                if chars.next_back().is_none() {
                    excess += 1;
                } else {
                    amount += 1;
                }
            }
            if excess > 0 {
                panic_imploded(count, amount);
            }
        } else {
            chars.nth_back(count);
        }

        let len = chars.as_str().len();
        self.truncate(len);
    }
}

#[cfg(feature = "alloc")]
impl<T> Extend<Operation<T>> for Vec<T> {
    #[inline]
    fn extend<I: IntoIterator<Item = Operation<T>>>(&mut self, iter: I) {
        stack_extend(self, iter.into_iter())
    }
}
#[cfg(feature = "alloc")]
impl<T> Extend<Operation<T>> for VecDeque<T> {
    #[inline]
    fn extend<I: IntoIterator<Item = Operation<T>>>(&mut self, iter: I) {
        stack_extend(self, iter.into_iter())
    }
}
#[cfg(feature = "alloc")]
impl<T> Extend<Operation<T>> for LinkedList<T> {
    #[inline]
    fn extend<I: IntoIterator<Item = Operation<T>>>(&mut self, iter: I) {
        stack_extend(self, iter.into_iter())
    }
}

/// Extra methods for iterators of [`Instruction`]s.
pub trait StackingIteratorExt<T>: Iterator<Item = Instruction<T>> {
    /// Converts an iterator of [`Instruction`]s into an iterator of collections.
    ///
    /// # Example
    ///
    /// ```
    /// use stacking_iterator::prelude::*;
    ///
    /// let mut inst = vec![
    ///     Instruction::Push('h'),
    ///     Instruction::Push('e'),
    ///     Instruction::Push('l'),
    ///     Instruction::Push('l'),
    ///     Instruction::Push('o'),
    ///     Instruction::Yield,
    ///     Instruction::Pop,
    ///     Instruction::Yield,
    /// ].into_iter();
    /// assert_eq!(inst.collecting().collect::<Vec<String>>(), ["hello", "hell"]);
    /// ```
    #[inline]
    fn collecting<C: Default + Clone + Extend<T> + Contract>(self) -> Collecting<Self, C>
    where
        Self: Sized,
    {
        Collecting {
            iter: self,
            collection: Default::default(),
        }
    }

    /// Performs the operations of the iterator until the next yield.
    ///
    /// If the iterator ends before a yield is encountered, this returns `None` but still mutates
    /// the original collection.
    ///
    /// # Example
    ///
    /// ```
    /// use stacking_iterator::prelude::*;
    ///
    /// let mut inst = vec![
    ///     Instruction::Push('h'),
    ///     Instruction::Push('e'),
    ///     Instruction::Push('l'),
    ///     Instruction::Push('l'),
    ///     Instruction::Push('o'),
    ///     Instruction::Yield,
    ///     Instruction::Pop,
    ///     Instruction::Yield,
    ///     Instruction::Pop,
    ///     Instruction::Pop,
    ///     Instruction::Pop,
    ///     Instruction::Pop,
    /// ].into_iter();
    /// let mut coll = String::new();
    /// assert_eq!(inst.operate(&mut coll).map(|x| &**x), Some("hello"));
    /// assert_eq!(coll, "hello");
    /// assert_eq!(inst.operate(&mut coll).map(|x| &**x), Some("hell"));
    /// assert_eq!(coll, "hell");
    /// assert_eq!(inst.operate(&mut coll), None);
    /// assert_eq!(coll, "");
    /// ```
    #[inline]
    fn operate<'c, C: Extend<T> + Contract>(&mut self, collection: &'c mut C) -> Option<&'c C> {
        let mut yielding = false;
        stack_extend(
            collection,
            self.map_while(|inst| match inst {
                Instruction::Mutate(operation) => Some(operation),
                Instruction::Yield => {
                    yielding = true;
                    None
                }
            }),
        );
        if yielding {
            Some(collection)
        } else {
            None
        }
    }
}
impl<T, I: Iterator<Item = Instruction<T>>> StackingIteratorExt<T> for I {}

/// Iterator returned by [`StackingIteratorExt::collecting`].
#[derive(Clone, Debug)]
pub struct Collecting<I, C> {
    iter: I,
    collection: C,
}
impl<T, I: Iterator<Item = Instruction<T>>, C: Clone + Extend<T> + Contract> Iterator
    for Collecting<I, C>
{
    type Item = C;

    #[inline]
    fn next(&mut self) -> Option<C> {
        self.iter.operate(&mut self.collection).cloned()
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let (_, hi) = self.iter.size_hint();

        // we could never yield a sequence, or we could yield every time
        (0, hi)
    }
}

#[cfg(test)]
mod tests {
    use core::iter::repeat;

    use super::{Contract, Instruction, Operation, StackingIteratorExt};
    use alloc::collections::{LinkedList, VecDeque};

    macro_rules! seq {
        (@ $stack:expr; - $($item:tt)*) => {
            $stack.push(Instruction::Mutate(Operation::Pop));
            seq!(@ $stack; $($item)*);
        };
        (@ $stack:expr; + $($item:tt)*) => {
            $stack.push(Instruction::Yield);
            seq!(@ $stack; $($item)*);
        };
        (@ $stack:expr; $val:ident $($item:tt)*) => {
            $stack.push(Instruction::Mutate(Operation::Push(stringify!($val))));
            seq!(@ $stack; $($item)*);
        };
        (@ $stack:expr;) => {
        };
        ($($item:tt)*) => {
            {
                #[allow(unused_mut)]
                let mut stack: Vec<Instruction<&'static str>> = Vec::new();
                seq!(@ stack; $($item)*);
                stack.into_iter()
            }
        };
    }

    #[test]
    fn contract_vec() {
        let mut v = vec![1, 2, 3, 4, 5];
        assert_eq!(v, [1, 2, 3, 4, 5]);
        v.contract(0);
        assert_eq!(v, [1, 2, 3, 4, 5]);
        v.contract(2);
        assert_eq!(v, [1, 2, 3]);
        v.contract(3);
        assert_eq!(v, []);
        v.contract(0);
    }

    #[test]
    fn contract_vec_deque() {
        let mut v = VecDeque::from(vec![1, 2, 3, 4, 5]);
        assert_eq!(v, [1, 2, 3, 4, 5]);
        v.contract(0);
        assert_eq!(v, [1, 2, 3, 4, 5]);
        v.contract(2);
        assert_eq!(v, [1, 2, 3]);
        v.contract(3);
        assert_eq!(v, []);
        v.contract(0);
    }

    #[test]
    fn contract_linked_list() {
        let mut l = LinkedList::from([1, 2, 3, 4, 5]);
        assert_eq!(l.iter().copied().collect::<Vec<_>>(), [1, 2, 3, 4, 5]);
        l.contract(0);
        assert_eq!(l.iter().copied().collect::<Vec<_>>(), [1, 2, 3, 4, 5]);
        l.contract(2);
        assert_eq!(l.iter().copied().collect::<Vec<_>>(), [1, 2, 3]);
        l.contract(3);
        assert_eq!(l.iter().copied().collect::<Vec<_>>(), []);
        l.contract(0);
    }

    #[test]
    fn contract_string() {
        let mut s = "世界こんにちは".to_owned();
        assert_eq!(s, "世界こんにちは");
        s.contract(0);
        assert_eq!(s, "世界こんにちは");
        s.contract(5);
        assert_eq!(s, "世界");
        s.contract(1);
        assert_eq!(s, "世");
        s.contract(1);
        assert_eq!(s, "");
        s.contract(0);
    }

    #[test]
    #[should_panic = "cannot contract by 1 when collection has 0 elements"]
    fn implode_empty_vec() {
        let mut v = <Vec<()>>::new();
        v.contract(1);
    }

    #[test]
    #[should_panic = "cannot contract by 2 when collection has 1 element"]
    fn implode_vec() {
        let mut v = vec![1];
        v.contract(2);
    }

    #[test]
    #[should_panic = "cannot contract by 1 when collection has 0 elements"]
    fn implode_empty_vec_deque() {
        let mut v = <VecDeque<()>>::new();
        v.contract(1);
    }

    #[test]
    #[should_panic = "cannot contract by 2 when collection has 1 element"]
    fn implode_vec_deque() {
        let mut v = VecDeque::from([1]);
        v.contract(2);
    }

    #[test]
    #[should_panic = "cannot contract by 1 when collection has 0 elements"]
    fn implode_empty_linked_list() {
        let mut l = <LinkedList<()>>::new();
        l.contract(1);
    }

    #[test]
    #[should_panic = "cannot contract by 2 when collection has 1 element"]
    fn implode_linked_list() {
        let mut l = LinkedList::from([1]);
        l.contract(2);
    }

    #[test]
    #[should_panic = "cannot contract by 1 when collection has 0 elements"]
    fn implode_empty_string() {
        let mut s = String::new();
        s.contract(1);
    }

    #[test]
    #[should_panic = "cannot contract by 2 when collection has 1 element"]
    fn implode_string() {
        let mut l = "あ".to_owned();
        l.contract(2);
    }

    #[test]
    fn extend_push() {
        let mut v = vec![1, 2, 3];
        v.extend([4, 5, 6].into_iter().map(Operation::Push));
        assert_eq!(v, [1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn extend_pop() {
        let mut v = vec![1, 2, 3];
        v.extend(repeat(Operation::Pop).take(2));
        assert_eq!(v, [1]);
    }

    #[test]
    fn empty() {
        let seq = seq!();
        assert_eq!(seq.collecting::<String>().next(), None);
    }

    #[test]
    fn sneaky_empty() {
        let seq = seq!(h e l l o);
        let mut iter = seq.collecting::<String>();
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn hello() {
        let seq = seq!(h e l l o +);
        let mut iter = seq.collecting::<String>();
        assert_eq!(iter.next().as_deref(), Some("hello"));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn hello_hell() {
        let seq = seq!(h e l l o + - +);
        let collected: Vec<String> = seq.collecting().collect();
        assert_eq!(collected, ["hello", "hell"]);
    }

    #[test]
    fn hello_world() {
        let seq = seq!(h e l l o + - + - - - - w o r l d +);
        let collected: Vec<String> = seq.collecting().collect();
        assert_eq!(collected, ["hello", "hell", "world"]);
    }
}