use core::ops::ControlFlow;
use crate::{
FallibleLend, FallibleLender, FallibleLending, FusedFallibleLender, TakeWhile,
try_trait_v2::Try,
};
impl<L: FallibleLender, P> TakeWhile<L, P> {
#[inline]
pub(crate) fn new_fallible(lender: L, predicate: P) -> TakeWhile<L, P> {
crate::__check_fallible_lender_covariance::<L>();
TakeWhile {
lender,
predicate,
flag: false,
}
}
}
impl<'lend, L, P> FallibleLending<'lend> for TakeWhile<L, P>
where
P: FnMut(&FallibleLend<'lend, L>) -> Result<bool, L::Error>,
L: FallibleLender,
{
type Lend = FallibleLend<'lend, L>;
}
impl<L, P> FallibleLender for TakeWhile<L, P>
where
P: FnMut(&FallibleLend<'_, L>) -> Result<bool, L::Error>,
L: FallibleLender,
{
type Error = L::Error;
crate::unsafe_assume_covariance_fallible!();
#[inline]
fn next(&mut self) -> Result<Option<FallibleLend<'_, Self>>, Self::Error> {
if !self.flag {
let x = match self.lender.next()? {
Some(x) => x,
None => {
self.flag = true;
return Ok(None);
}
};
if (self.predicate)(&x)? {
return Ok(Some(x));
}
self.flag = true;
}
Ok(None)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
if self.flag {
(0, Some(0))
} else {
let (_, upper) = self.lender.size_hint();
(0, upper)
}
}
#[inline]
fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> Result<R, Self::Error>
where
Self: Sized,
F: FnMut(B, FallibleLend<'_, Self>) -> Result<R, Self::Error>,
R: Try<Output = B>,
{
if self.flag {
return Ok(R::from_output(init));
}
let flag = &mut self.flag;
let predicate = &mut self.predicate;
match self.lender.try_fold(init, move |acc, x| {
if (predicate)(&x)? {
match f(acc, x)?.branch() {
ControlFlow::Continue(b) => Ok(ControlFlow::Continue(b)),
ControlFlow::Break(residual) => {
Ok(ControlFlow::Break(R::from_residual(residual)))
}
}
} else {
*flag = true;
Ok(ControlFlow::Break(R::from_output(acc)))
}
})? {
ControlFlow::Continue(acc) => Ok(R::from_output(acc)),
ControlFlow::Break(r) => Ok(r),
}
}
#[inline]
fn fold<B, F>(mut self, init: B, mut f: F) -> Result<B, Self::Error>
where
Self: Sized,
F: FnMut(B, FallibleLend<'_, Self>) -> Result<B, Self::Error>,
{
if self.flag {
return Ok(init);
}
let flag = &mut self.flag;
let predicate = &mut self.predicate;
match self.lender.try_fold(init, move |acc, x| {
if (predicate)(&x)? {
match f(acc, x) {
Ok(b) => Ok(ControlFlow::Continue(b)),
Err(e) => Err(e),
}
} else {
*flag = true;
Ok(ControlFlow::Break(acc))
}
})? {
ControlFlow::Continue(acc) | ControlFlow::Break(acc) => Ok(acc),
}
}
}
impl<L, P> FusedFallibleLender for TakeWhile<L, P>
where
P: FnMut(&FallibleLend<'_, L>) -> Result<bool, L::Error>,
L: FallibleLender,
{
}
#[cfg(test)]
mod test {
use core::convert::Infallible;
use crate::prelude::*;
struct Resurrect {
step: usize,
}
impl<'l> crate::FallibleLending<'l> for Resurrect {
type Lend = i32;
}
impl crate::FallibleLender for Resurrect {
type Error = Infallible;
crate::check_covariance_fallible!();
fn next(&mut self) -> Result<Option<crate::FallibleLend<'_, Self>>, Self::Error> {
let s = self.step;
self.step += 1;
Ok(match s {
0 => Some(10),
2 => Some(30),
_ => None,
})
}
}
#[test]
fn test_none_is_terminal() -> Result<(), Infallible> {
let mut tw = (Resurrect { step: 0 }).take_while(|_| Ok(true));
assert_eq!(tw.next()?, Some(10));
assert_eq!(tw.next()?, None); assert_eq!(tw.next()?, None); Ok(())
}
}