From 815cf9bf0740436a3d40e5a39bccec4987857feb Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Mon, 25 Sep 2023 11:07:22 -0700
Subject: [PATCH] merge: implement Default and Extend on MergeBuilder

`futures::stream::Stream::collect()` requires a collection that
implements `Default` and `Extend`, and I would like to to be able to
collect a stream of trees.
---
 lib/src/merge.rs | 40 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/lib/src/merge.rs b/lib/src/merge.rs
index 102fa4b4b..334427cf4 100644
--- a/lib/src/merge.rs
+++ b/lib/src/merge.rs
@@ -277,6 +277,15 @@ pub struct MergeBuilder<T> {
     adds: Vec<T>,
 }
 
+impl<T> Default for MergeBuilder<T> {
+    fn default() -> Self {
+        Self {
+            removes: Default::default(),
+            adds: Default::default(),
+        }
+    }
+}
+
 impl<T> MergeBuilder<T> {
     /// Requires that exactly one more "adds" than "removes" have been added to
     /// this builder.
@@ -296,15 +305,23 @@ impl<T> IntoIterator for Merge<T> {
 
 impl<T> FromIterator<T> for MergeBuilder<T> {
     fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
-        let mut removes = vec![];
-        let mut adds = vec![];
-        let mut curr = &mut adds;
-        let mut next = &mut removes;
+        let mut builder = MergeBuilder::default();
+        builder.extend(iter);
+        builder
+    }
+}
+
+impl<T> Extend<T> for MergeBuilder<T> {
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
+        let (mut curr, mut next) = if self.adds.len() != self.removes.len() {
+            (&mut self.removes, &mut self.adds)
+        } else {
+            (&mut self.adds, &mut self.removes)
+        };
         for item in iter {
             curr.push(item);
             std::mem::swap(&mut curr, &mut next);
         }
-        MergeBuilder { removes, adds }
     }
 }
 
@@ -789,6 +806,19 @@ mod tests {
         MergeBuilder::from_iter([1, 2]).build();
     }
 
+    #[test]
+    fn test_extend() {
+        // 1-way merge
+        let mut builder: MergeBuilder<i32> = Default::default();
+        builder.extend([1]);
+        assert_eq!(builder.build(), c(&[], &[1]));
+        // 5-way merge
+        let mut builder: MergeBuilder<i32> = Default::default();
+        builder.extend([1, 2]);
+        builder.extend([3, 4, 5]);
+        assert_eq!(builder.build(), c(&[2, 4], &[1, 3, 5]));
+    }
+
     #[test]
     fn test_map() {
         fn increment(i: &i32) -> i32 {