Skip to content

Instantly share code, notes, and snippets.

@thomcc
Last active Jun 5, 2021
Embed
What would you like to do?
/// 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));
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment