Implement TryExtend<V> for OOM-handling PrimaryMap<K, V> (#12610)

Author: fitzgen

crates/environ/src/collections.rs

@@ -14,7 +14,7 @@ pub use index_map::IndexMap;
 pub use primary_map::PrimaryMap;
 pub use secondary_map::SecondaryMap;
 pub use wasmtime_core::{
-    alloc::{String, TryClone, TryNew, Vec, try_new},
+    alloc::{String, TryClone, TryCollect, TryExtend, TryFromIterator, TryNew, Vec, try_new},
     vec,
 };
 

crates/environ/src/collections/primary_map.rs

@@ -1,10 +1,10 @@
+use crate::{collections::TryExtend, error::OutOfMemory};
 use core::{
     fmt,
     ops::{Index, IndexMut},
 };
 use cranelift_entity::EntityRef;
 use serde::{Serialize, ser::SerializeSeq};
-use wasmtime_core::error::OutOfMemory;
 
 /// Like [`cranelift_entity::PrimaryMap`] but enforces fallible allocation for
 /// all methods that allocate.
@@ -217,6 +217,25 @@ where
     }
 }
 
+impl<K, V> TryExtend<V> for PrimaryMap<K, V>
+where
+    K: EntityRef,
+{
+    fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory>
+    where
+        I: IntoIterator<Item = V>,
+    {
+        let iter = iter.into_iter();
+        let (min, max) = iter.size_hint();
+        let cap = max.unwrap_or(min);
+        self.reserve(cap)?;
+        for v in iter {
+            self.push(v)?;
+        }
+        Ok(())
+    }
+}
+
 impl<K, V> Index<K> for PrimaryMap<K, V>
 where
     K: EntityRef,