Fuse iterator supplied to `SumTree::from_iter` (#10571)

Antonio Scandurra and Kyle created 2 years ago

This fixes an issue that could cause `from_iter` to never finish if the
underlying iterator restarted after returning `None` for the first time.

We only saw this in development but I wanna cherry-pick it to stable and
preview, just in case.

Release Notes:

- N/A

Co-authored-by: Kyle <[e-mail]>

Change summary

crates/sum_tree/src/sum_tree.rs | 23 ++++++++++++++++++++++-
1 file changed, 22 insertions(+), 1 deletion(-)

Detailed changes

crates/sum_tree/src/sum_tree.rs

@@ -179,7 +179,7 @@ impl<T: Item> SumTree<T> {
     ) -> 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<T, { 2 * TREE_BASE }> = iter.by_ref().take(2 * TREE_BASE).collect();
             let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> =
@@ -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::<Vec<_>>()
+        );
+
+        // 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,