Implement core::iter::Extend for storage collections (#341)

* Implement extend for Vec

* Implement extend for BinaryHeap

* Remove redundant core::iter imports

* Implement extend for BitVec

* Implement extend for BTreeMap

* Implement extend for HashMap

* Fmt

* Move trait bounds to where clause

* Cloned -> Copied

* More bounds moved to where clause

Co-authored-by: Hero Bird <robbepop@web.de>

core/src/storage/collections/binary_heap/impls.rs

@@ -440,3 +440,29 @@ where
         Iter::new(self)
     }
 }
+
+impl<T> Extend<T> for BinaryHeap<T>
+where
+    T: Codec + Ord,
+{
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = T>,
+    {
+        for i in iter {
+            self.push(i)
+        }
+    }
+}
+
+impl<'a, T> Extend<&'a T> for BinaryHeap<T>
+where
+    T: Codec + Ord + Copy + 'a,
+{
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = &'a T>,
+    {
+        self.extend(iter.into_iter().copied())
+    }
+}

core/src/storage/collections/binary_heap/tests.rs

@@ -28,6 +28,7 @@ use core::{
     cmp::Ord,
     fmt::Debug,
 };
+use ink_prelude::collections::HashSet;
 use ink_primitives::Key;
 use scale::{
     Codec,
@@ -346,3 +347,26 @@ fn min_heap_with_multiple_levels() -> Result<()> {
         Ok(())
     })
 }
+
+#[test]
+fn extend() -> Result<()> {
+    env::test::run_test::<env::DefaultEnvTypes, _>(|_| {
+        // given
+        let mut heap = filled_heap();
+        let arr = [1, 2, 3];
+
+        // when
+        heap.extend(&arr);
+
+        // then
+        let mut expected = HashSet::new();
+        expected.extend(heap.values());
+        expected.extend(arr.iter());
+
+        let actual = heap.values().collect::<HashSet<_>>();
+
+        assert_eq!(actual, expected);
+
+        Ok(())
+    })
+}

core/src/storage/collections/bitvec/tests.rs

@@ -115,6 +115,20 @@ fn push_filled() {
     assert_eq!(filled.len(), len + 1);
 }
 
+#[test]
+fn extend() {
+    let mut filled = new_filled_bitvec();
+    let arr = [true, false, true];
+
+    let mut expected = Vec::new();
+    expected.extend(filled.iter());
+    expected.extend(&arr);
+
+    filled.extend(&arr);
+
+    assert!(filled.iter().eq(expected.iter().cloned()));
+}
+
 #[test]
 fn pop_empty() {
     assert_eq!(new_empty_bitvec().pop(), None);

core/src/storage/collections/bitvec/vec.rs

@@ -117,6 +117,20 @@ impl Drop for BitVec {
     }
 }
 
+impl Extend<bool> for BitVec {
+    fn extend<T: IntoIterator<Item = bool>>(&mut self, iter: T) {
+        for b in iter {
+            self.push(b)
+        }
+    }
+}
+
+impl<'a> Extend<&'a bool> for BitVec {
+    fn extend<T: IntoIterator<Item = &'a bool>>(&mut self, iter: T) {
+        self.extend(iter.into_iter().copied())
+    }
+}
+
 impl BitVec {
     /// Returns the number of bits.
     pub fn len(&self) -> u32 {

core/src/storage/collections/btree_map/impls.rs

@@ -262,6 +262,38 @@ impl<K, V> AllocateUsing for BTreeMap<K, V> {
     }
 }
 
+impl<K, V> Extend<(K, V)> for BTreeMap<K, V>
+where
+    K: Codec + Ord,
+    V: Codec,
+{
+    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
+        for (k, v) in iter {
+            self.insert(k, v);
+        }
+    }
+}
+
+impl<'a, K, V> Extend<(&'a K, &'a V)> for BTreeMap<K, V>
+where
+    K: Codec + Ord + Copy,
+    V: Codec + Copy,
+{
+    fn extend<I: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: I) {
+        self.extend(iter.into_iter().map(|(&key, &value)| (key, value)));
+    }
+}
+
+impl<'a, K, V> Extend<&'a (K, V)> for BTreeMap<K, V>
+where
+    K: Codec + Ord + Copy + 'a,
+    V: Codec + Copy + 'a,
+{
+    fn extend<I: IntoIterator<Item = &'a (K, V)>>(&mut self, iter: I) {
+        self.extend(iter.into_iter().copied());
+    }
+}
+
 impl<K, V> BTreeMap<K, V> {
     /// Returns the index of the root node.
     pub(super) fn root(&self) -> &Option<NodeHandle> {

core/src/storage/collections/btree_map/tests.rs

@@ -34,6 +34,7 @@ use crate::{
         BTreeMap,
     },
 };
+use ink_prelude::collections::HashMap;
 use ink_primitives::Key;
 use itertools::Itertools;
 
@@ -442,3 +443,29 @@ fn complex_trees_work() -> Result<()> {
         Ok(())
     })
 }
+
+#[test]
+fn extend_works() -> Result<()> {
+    env::test::run_test::<env::DefaultEnvTypes, _>(|_| {
+        // given
+        let arr1 = [(1i32, 2i32), (3, 4), (5, 6)];
+        let arr2 = [(7i32, 8i32), (9, 10)];
+        let mut map = empty_map();
+
+        let mut expected = HashMap::new();
+
+        expected.extend(arr1.iter().cloned());
+        expected.extend(arr2.iter().cloned());
+
+        // when
+        map.extend(arr1.iter());
+        map.extend(arr2.iter());
+
+        // then
+        assert_eq!(map.len() as usize, expected.len());
+        for (k, v) in &expected {
+            assert_eq!(Some(v), map.get(k))
+        }
+        Ok(())
+    })
+}

core/src/storage/collections/hash_map/impls.rs

@@ -194,6 +194,38 @@ impl<K, V> Initialize for HashMap<K, V> {
     }
 }
 
+impl<K, V> Extend<(K, V)> for HashMap<K, V>
+where
+    K: scale::Codec + Hash + Eq,
+    V: scale::Codec,
+{
+    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
+        for (k, v) in iter {
+            self.insert(k, v);
+        }
+    }
+}
+
+impl<'a, K, V> Extend<(&'a K, &'a V)> for HashMap<K, V>
+where
+    K: scale::Codec + Hash + Eq + Copy,
+    V: scale::Codec + Copy,
+{
+    fn extend<I: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: I) {
+        self.extend(iter.into_iter().map(|(&key, &value)| (key, value)));
+    }
+}
+
+impl<'a, K, V> Extend<&'a (K, V)> for HashMap<K, V>
+where
+    K: scale::Codec + Hash + Eq + Copy + 'a,
+    V: scale::Codec + Copy + 'a,
+{
+    fn extend<I: IntoIterator<Item = &'a (K, V)>>(&mut self, iter: I) {
+        self.extend(iter.into_iter().copied());
+    }
+}
+
 impl<K, V> HashMap<K, V> {
     /// Returns the number of key-value pairs in the map.
     pub fn len(&self) -> u32 {

core/src/storage/collections/hash_map/tests.rs

@@ -188,3 +188,29 @@ fn mutate_with() -> Result<()> {
         Ok(())
     })
 }
+
+#[test]
+fn extend_works() -> Result<()> {
+    env::test::run_test::<env::DefaultEnvTypes, _>(|_| {
+        // given
+        let arr1 = [(1i32, 2i32), (3, 4), (5, 6)];
+        let arr2 = [(7i32, 8i32), (9, 10)];
+        let mut map = new_empty::<i32, i32>();
+
+        let mut expected = ink_prelude::collections::HashMap::new();
+
+        expected.extend(arr1.iter().cloned());
+        expected.extend(arr2.iter().cloned());
+
+        // when
+        map.extend(arr1.iter());
+        map.extend(arr2.iter());
+
+        // then
+        assert_eq!(map.len() as usize, expected.len());
+        for (k, v) in &expected {
+            assert_eq!(Some(v), map.get(k))
+        }
+        Ok(())
+    })
+}

core/src/storage/collections/vec/impls.rs

@@ -12,11 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use core::iter::{
-    DoubleEndedIterator,
-    ExactSizeIterator,
-};
-
 #[cfg(feature = "ink-generate-abi")]
 use ink_abi::{
     HasLayout,
@@ -149,6 +144,32 @@ where
     }
 }
 
+impl<T> Extend<T> for Vec<T>
+where
+    T: scale::Codec,
+{
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = T>,
+    {
+        for i in iter {
+            self.push(i)
+        }
+    }
+}
+
+impl<'a, T> Extend<&'a T> for Vec<T>
+where
+    T: scale::Codec + Copy + 'a,
+{
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = &'a T>,
+    {
+        self.extend(iter.into_iter().copied())
+    }
+}
+
 impl<T> scale::Encode for Vec<T> {
     fn encode_to<W: scale::Output>(&self, dest: &mut W) {
         self.len.encode_to(dest);

core/src/storage/collections/vec/tests.rs

@@ -241,6 +241,20 @@ fn iter_size_hint() {
     assert_eq!(iter.size_hint(), (3, Some(3)));
 }
 
+#[test]
+fn extend() {
+    let mut vec1 = new_filled_vec();
+    let arr = [1, 2, 3];
+
+    let mut expected = ink_prelude::vec::Vec::new();
+    expected.extend(vec1.iter());
+    expected.extend(&arr);
+
+    vec1.extend(&arr);
+
+    assert!(vec1.iter().eq(expected.iter()));
+}
+
 #[test]
 fn regression_issue_193() {
     let mut vec = new_empty_vec();