Auto merge of #12956 - oxalica:feat/la-arena-apis, r=lnicola

More methods and traits for `la_arena::ArenaMap`

Continue of #12931. Seems that I forgot some methods in the previous PR :(

I also changed `ArenaMap::insert` to return the old value, to match the map-like collection API of std. **So this is a breaking change.**

r? `@lnicola`

4 files changed, 51 insertions, 5 deletions

diff --git a/crates/hir-def/src/generics.rs b/crates/hir-def/src/generics.rs
index 2397cf50158..469b28c2d9e 100644
--- a/crates/hir-def/src/generics.rs
+++ b/crates/hir-def/src/generics.rs
@@ -451,7 +451,7 @@ impl HasChildSource<LocalTypeOrConstParamId> for GenericDefId {
         if let GenericDefId::TraitId(id) = *self {
             let trait_ref = id.lookup(db).source(db).value;
             let idx = idx_iter.next().unwrap();
-            params.insert(idx, Either::Right(trait_ref))
+            params.insert(idx, Either::Right(trait_ref));
         }
 
         if let Some(generic_params_list) = generic_params_list {
diff --git a/crates/hir-def/src/visibility.rs b/crates/hir-def/src/visibility.rs
index 6e22a877a9f..087268a9ece 100644
--- a/crates/hir-def/src/visibility.rs
+++ b/crates/hir-def/src/visibility.rs
@@ -224,7 +224,7 @@ pub(crate) fn field_visibilities_query(
     let resolver = variant_id.module(db).resolver(db);
     let mut res = ArenaMap::default();
     for (field_id, field_data) in var_data.fields().iter() {
-        res.insert(field_id, field_data.visibility.resolve(db, &resolver))
+        res.insert(field_id, field_data.visibility.resolve(db, &resolver));
     }
     Arc::new(res)
 }
diff --git a/crates/hir-ty/src/lower.rs b/crates/hir-ty/src/lower.rs
index 3ed9c941f47..239f66bcb7e 100644
--- a/crates/hir-ty/src/lower.rs
+++ b/crates/hir-ty/src/lower.rs
@@ -1126,7 +1126,7 @@ pub(crate) fn field_types_query(
     let ctx =
         TyLoweringContext::new(db, &resolver).with_type_param_mode(ParamLoweringMode::Variable);
     for (field_id, field_data) in var_data.fields().iter() {
-        res.insert(field_id, make_binders(db, &generics, ctx.lower_ty(&field_data.type_ref)))
+        res.insert(field_id, make_binders(db, &generics, ctx.lower_ty(&field_data.type_ref)));
     }
     Arc::new(res)
 }
diff --git a/lib/la-arena/src/map.rs b/lib/la-arena/src/map.rs
index f74ec5d97cb..5f347e27450 100644
--- a/lib/la-arena/src/map.rs
+++ b/lib/la-arena/src/map.rs
@@ -21,12 +21,42 @@ impl<T, V> ArenaMap<Idx<T>, V> {
         Self { v: Vec::with_capacity(capacity), _ty: PhantomData }
     }
 
+    /// Reserves capacity for at least additional more elements to be inserted in the map.
+    pub fn reserve(&mut self, additional: usize) {
+        self.v.reserve(additional);
+    }
+
+    /// Clears the map, removing all elements.
+    pub fn clear(&mut self) {
+        self.v.clear();
+    }
+
+    /// Shrinks the capacity of the map as much as possible.
+    pub fn shrink_to_fit(&mut self) {
+        let min_len = self.v.iter().rposition(|slot| slot.is_some()).map_or(0, |i| i + 1);
+        self.v.truncate(min_len);
+        self.v.shrink_to_fit();
+    }
+
+    /// Returns whether the map contains a value for the specified index.
+    pub fn contains_idx(&self, idx: Idx<T>) -> bool {
+        matches!(self.v.get(Self::to_idx(idx)), Some(Some(_)))
+    }
+
+    /// Removes an index from the map, returning the value at the index if the index was previously in the map.
+    pub fn remove(&mut self, idx: Idx<T>) -> Option<V> {
+        self.v.get_mut(Self::to_idx(idx))?.take()
+    }
+
     /// Inserts a value associated with a given arena index into the map.
-    pub fn insert(&mut self, idx: Idx<T>, t: V) {
+    ///
+    /// If the map did not have this index present, None is returned.
+    /// Otherwise, the value is updated, and the old value is returned.
+    pub fn insert(&mut self, idx: Idx<T>, t: V) -> Option<V> {
         let idx = Self::to_idx(idx);
 
         self.v.resize_with((idx + 1).max(self.v.len()), || None);
-        self.v[idx] = Some(t);
+        self.v[idx].replace(t)
     }
 
     /// Returns a reference to the value associated with the provided index
@@ -94,6 +124,22 @@ impl<T, V> Default for ArenaMap<Idx<V>, T> {
     }
 }
 
+impl<T, V> Extend<(Idx<V>, T)> for ArenaMap<Idx<V>, T> {
+    fn extend<I: IntoIterator<Item = (Idx<V>, T)>>(&mut self, iter: I) {
+        iter.into_iter().for_each(move |(k, v)| {
+            self.insert(k, v);
+        });
+    }
+}
+
+impl<T, V> FromIterator<(Idx<V>, T)> for ArenaMap<Idx<V>, T> {
+    fn from_iter<I: IntoIterator<Item = (Idx<V>, T)>>(iter: I) -> Self {
+        let mut this = Self::new();
+        this.extend(iter);
+        this
+    }
+}
+
 /// A view into a single entry in a map, which may either be vacant or occupied.
 ///
 /// This `enum` is constructed from the [`entry`] method on [`ArenaMap`].