diff --git a/crates/sum_tree/src/sum_tree.rs b/crates/sum_tree/src/sum_tree.rs index c9e1ae5bcc..3909ef1b39 100644 --- a/crates/sum_tree/src/sum_tree.rs +++ b/crates/sum_tree/src/sum_tree.rs @@ -179,7 +179,7 @@ impl SumTree { ) -> Self { let mut nodes = Vec::new(); - let mut iter = iter.into_iter().peekable(); + let mut iter = iter.into_iter().fuse().peekable(); while iter.peek().is_some() { let items: ArrayVec = iter.by_ref().take(2 * TREE_BASE).collect(); let item_summaries: ArrayVec = @@ -1244,6 +1244,27 @@ mod tests { assert_eq!(tree.get(&4, &()), Some(&4)); } + #[test] + fn test_from_iter() { + assert_eq!( + SumTree::from_iter(0..100, &()).items(&()), + (0..100).collect::>() + ); + + // Ensure `from_iter` works correctly when the given iterator restarts + // after calling `next` if `None` was already returned. + let mut ix = 0; + let iterator = std::iter::from_fn(|| { + ix = (ix + 1) % 2; + if ix == 1 { + Some(1) + } else { + None + } + }); + assert_eq!(SumTree::from_iter(iterator, &()).items(&()), vec![1]); + } + #[derive(Clone, Default, Debug)] pub struct IntegersSummary { count: usize,