#![cfg_attr(not(feature = "std"), no_std)]
#[cfg(feature = "alloc")]
extern crate alloc;
#[cfg(feature = "alloc")]
use alloc::vec::Vec;
pub trait TryNextWithContext<C, S = ()>
where
S: Default + Copy,
{
type Item;
type Error;
fn try_next_with_context(&mut self, context: &mut C)
-> Result<Option<Self::Item>, Self::Error>;
#[cfg(feature = "alloc")]
#[inline]
fn try_collect_with_context(
&mut self,
context: &mut C,
) -> Result<Vec<Self::Item>, Self::Error> {
let mut vs = Vec::new();
while let Some(v) = self.try_next_with_context(context)? {
vs.push(v);
}
Ok(vs)
}
fn stats(&self) -> S {
S::default()
}
}
pub trait TryNext<S = ()>
where
S: Default + Copy,
{
type Item;
type Error;
fn try_next(&mut self) -> Result<Option<Self::Item>, Self::Error>;
#[cfg(feature = "alloc")]
#[inline]
fn try_collect(&mut self) -> Result<Vec<Self::Item>, Self::Error> {
let mut vs = Vec::new();
while let Some(v) = self.try_next()? {
vs.push(v);
}
Ok(vs)
}
fn stats(&self) -> S {
S::default()
}
}
pub struct IterInput<I>
where
I: Iterator,
{
iter: core::iter::Fuse<I>,
}
impl<I> IterInput<I>
where
I: Iterator,
{
pub fn from(iter: I) -> Self {
Self { iter: iter.fuse() }
}
}
impl<I> TryNext for IterInput<I>
where
I: Iterator,
{
type Item = I::Item;
type Error = core::convert::Infallible;
#[inline]
fn try_next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
Ok(self.iter.next())
}
}
impl<I, C> TryNextWithContext<C> for IterInput<I>
where
I: Iterator,
{
type Item = I::Item;
type Error = core::convert::Infallible;
#[inline]
fn try_next_with_context(
&mut self,
_context: &mut C,
) -> Result<Option<Self::Item>, Self::Error> {
Ok(self.iter.next())
}
}
#[cfg(feature = "std")]
pub mod io;
#[cfg(test)]
mod tests {
use super::{IterInput, TryNext, TryNextWithContext};
use core::convert::Infallible;
struct Counter {
current: usize,
limit: usize,
}
impl TryNext for Counter {
type Item = usize;
type Error = Infallible;
fn try_next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
if self.current < self.limit {
let v = self.current;
self.current += 1;
Ok(Some(v))
} else {
Ok(None)
}
}
}
struct FailableCounter {
current: usize,
fail_at: usize,
failed: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct UnitErr;
impl TryNext for FailableCounter {
type Item = usize;
type Error = UnitErr;
fn try_next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
if self.failed {
return Err(UnitErr);
}
if self.current == self.fail_at {
self.failed = true;
return Err(UnitErr);
}
let v = self.current;
self.current += 1;
Ok(Some(v))
}
}
#[cfg(feature = "alloc")]
fn drain<S: TryNext>(mut src: S) -> Result<Vec<S::Item>, S::Error> {
src.try_collect()
}
#[test]
fn counter_yields_then_none() {
let mut c = Counter {
current: 0,
limit: 3,
};
assert_eq!(c.try_next().unwrap(), Some(0));
assert_eq!(c.try_next().unwrap(), Some(1));
assert_eq!(c.try_next().unwrap(), Some(2));
assert_eq!(c.try_next().unwrap(), None);
assert_eq!(c.try_next().unwrap(), None);
}
#[test]
#[cfg(feature = "alloc")]
fn drain_collects_all_items() {
let c = Counter {
current: 0,
limit: 5,
};
let items = drain(c).unwrap();
assert_eq!(items, vec![0, 1, 2, 3, 4]);
}
#[test]
fn error_propagates() {
let mut s = FailableCounter {
current: 0,
fail_at: 2,
failed: false,
};
assert_eq!(s.try_next(), Ok(Some(0)));
assert_eq!(s.try_next(), Ok(Some(1)));
assert_eq!(s.try_next(), Err(UnitErr));
assert_eq!(s.try_next(), Err(UnitErr));
}
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
struct Ctx {
calls: usize,
}
impl TryNextWithContext<Ctx> for Counter {
type Item = usize;
type Error = Infallible;
fn try_next_with_context(
&mut self,
ctx: &mut Ctx,
) -> Result<Option<Self::Item>, Self::Error> {
ctx.calls += 1;
if self.current < self.limit {
let v = self.current;
self.current += 1;
Ok(Some(v))
} else {
Ok(None)
}
}
}
#[cfg(feature = "alloc")]
fn drain_with_ctx<C, S: TryNextWithContext<C>>(
mut src: S,
mut ctx: C,
) -> Result<(Vec<S::Item>, C), S::Error> {
let mut out = Vec::new();
while let Some(item) = src.try_next_with_context(&mut ctx)? {
out.push(item);
}
Ok((out, ctx))
}
#[test]
#[cfg(feature = "alloc")]
fn context_counter_yields_and_updates_context() {
let src = Counter {
current: 0,
limit: 3,
};
let (items, ctx) = drain_with_ctx(src, Ctx::default()).unwrap();
assert_eq!(items, vec![0, 1, 2]);
assert_eq!(ctx.calls, 4);
}
#[derive(Default)]
struct DummyCtx;
#[test]
fn iter_input_yields_all_items_from_array() {
let mut input = IterInput::from([1, 2, 3].into_iter());
let mut ctx = DummyCtx;
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), Some(1));
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), Some(2));
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), Some(3));
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
}
#[test]
fn iter_input_is_fused_after_exhaustion() {
let mut input = IterInput::from(0..1);
let mut ctx = DummyCtx;
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), Some(0));
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
}
#[test]
fn iter_input_with_empty_range_returns_none() {
let mut input = IterInput::from(0..0);
let mut ctx = DummyCtx;
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
}
#[test]
fn iter_input_over_string_bytes() {
let mut input = IterInput::from("hello, world!".bytes());
let mut ctx = DummyCtx;
let mut collected = [0u8; 13];
let mut count = 0;
while let Some(byte) = input.try_next_with_context(&mut ctx).unwrap() {
collected[count] = byte;
count += 1;
}
assert_eq!(&collected[..count], b"hello, world!");
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
assert_eq!(input.try_next_with_context(&mut ctx).unwrap(), None);
}
#[test]
fn stats_default_is_unit() {
let c = Counter {
current: 0,
limit: 1,
};
let unit_stats: () = TryNext::stats(&c);
assert_eq!(unit_stats, ());
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
struct MyStats {
calls: usize,
}
struct StatsSource {
remaining: usize,
emitted: usize,
call_count: usize,
}
impl TryNext<MyStats> for StatsSource {
type Item = usize;
type Error = Infallible;
fn try_next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
self.call_count += 1;
if self.remaining == 0 {
return Ok(None);
}
let v = self.emitted;
self.emitted += 1;
self.remaining -= 1;
Ok(Some(v))
}
fn stats(&self) -> MyStats {
MyStats {
calls: self.call_count,
}
}
}
#[test]
fn custom_stats_are_reported() {
let mut s = StatsSource {
remaining: 3,
emitted: 0,
call_count: 0,
};
assert_eq!(s.stats(), MyStats { calls: 0 });
assert_eq!(s.try_next().unwrap(), Some(0));
assert_eq!(s.try_next().unwrap(), Some(1));
assert_eq!(s.try_next().unwrap(), Some(2));
assert_eq!(s.try_next().unwrap(), None);
assert_eq!(s.stats(), MyStats { calls: 4 });
}
#[test]
fn iter_input_try_next_works_without_context() {
let mut input = IterInput::from([10, 20, 30].into_iter());
assert_eq!(TryNext::try_next(&mut input).unwrap(), Some(10));
assert_eq!(TryNext::try_next(&mut input).unwrap(), Some(20));
assert_eq!(TryNext::try_next(&mut input).unwrap(), Some(30));
assert_eq!(TryNext::try_next(&mut input).unwrap(), None);
assert_eq!(TryNext::try_next(&mut input).unwrap(), None); }
#[test]
#[cfg(feature = "alloc")]
fn iter_input_try_collect_collects_all() {
let mut input = IterInput::from([7, 8, 9].into_iter());
let items = input.try_collect().unwrap();
assert_eq!(items, vec![7, 8, 9]);
assert_eq!(TryNext::try_next(&mut input).unwrap(), None);
}
}