try_flat_map.rs

/// Something like a `flat_map` over a fallible iterator.
///
/// Note that the `map` itself is not fallible.
///
/// See test at end of file for usage.
//
// TODO this should probably take `impl IntoIterator<IntoIter = I>` but the
// generics here are already so fucking complicated, and i'm worried that
// will cause inference errors
#[inline]
pub fn try_flat_map<I, U, F, S, T, E>(iter: I, map: F) -> TryFlatMap<I, U, F>
where
    I: Iterator<Item = Result<S, E>>,
    U: IntoIterator<Item = T>,
    F: FnMut(S) -> U,
{
    TryFlatMap {
        outer: iter,
        map,
        inner: None,
    }
}

/// Iterator implementing flat mapping over a fallible iterator.
///
/// Constructed using [`try_flat_map`]
#[derive(Debug, Clone)]
pub struct TryFlatMap<I, U: IntoIterator, F> {
    outer: I,
    map: F,
    // `Some` iff `outer` returned Ok(v) and we haven't
    // finished going over it yet.
    inner: Option<U::IntoIter>,
}

impl<I, U, F, S, T, E> Iterator for TryFlatMap<I, U, F>
where
    I: Iterator<Item = Result<S, E>>,
    U: IntoIterator<Item = T>,
    F: FnMut(S) -> U,
{
    type Item = Result<T, E>;
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(ref mut inner) = self.inner {
                match inner.next() {
                    Some(item) => return Some(Ok(item)),
                    None => self.inner = None,
                }
            }
            match self.outer.next() {
                Some(Ok(v)) => self.inner = Some((self.map)(v).into_iter()),
                Some(Err(e)) => return Some(Err(e)),
                _ => return None,
            }
        }
    }
    // TODO: we should be able to compute at least a lower bound
    // for size_hint
}

impl<I, U, F, S, T, E> core::iter::FusedIterator for TryFlatMap<I, U, F>
where
    I: Iterator<Item = Result<S, E>> + core::iter::FusedIterator,
    U: IntoIterator<Item = T>,
    F: FnMut(S) -> U,
{
}

// TODO: implement DoubleEndedIterator. Note: requires storing another
// copy `Option<U::IntoIter>` for the back, see FlattenCompat in the
// stdlib.

#[test]
fn test_try_flat_map() {
    let v: Vec<Result<&'static str, ()>> =
        vec![Ok(""), Ok("abc"), Ok(""), Ok("def"), Ok("123"), Ok("")];
    let s: Result<String, ()> = try_flat_map(v.iter().copied(), |s| s.chars()).collect();
    assert_eq!(s, Ok("abcdef123".to_string()));

    type AnyError = Box<dyn std::error::Error + 'static + Send + Sync>;
    let example = || -> Vec<Result<Vec<i32>, AnyError>> {
        vec![
            Ok(vec![1, 2, 3]),
            Ok(vec![4, 5]),
            Ok(vec![]),
            Ok(vec![6, 7]),
            Ok(vec![]),
            Ok(vec![]),
            Ok(vec![8]),
            Err("oh no".to_string().into()),
            Ok(vec![9, 10, 11]),
        ]
    };

    let res = try_flat_map(example().into_iter(), |items| items.into_iter())
        .collect::<Result<Vec<i32>, _>>();
    assert_eq!(format!("{}", res.unwrap_err()), "oh no");

    let mut it = try_flat_map(example().into_iter(), |items| items.into_iter());
    assert!(matches!(it.next(), Some(Ok(1))));
    assert!(matches!(it.next(), Some(Ok(2))));
    assert!(matches!(it.next(), Some(Ok(3))));
    assert!(matches!(it.next(), Some(Ok(4))));
    assert!(matches!(it.next(), Some(Ok(5))));
    assert!(matches!(it.next(), Some(Ok(6))));
    assert!(matches!(it.next(), Some(Ok(7))));
    assert!(matches!(it.next(), Some(Ok(8))));
    assert!(matches!(it.next(), Some(Err(_))));
    assert!(matches!(it.next(), Some(Ok(9))));
    assert!(matches!(it.next(), Some(Ok(10))));
    assert!(matches!(it.next(), Some(Ok(11))));
    assert!(matches!(it.next(), None));
}