Add SmallMap::retain

Summary: Useful to have.

Reviewed By: stepancheg

Differential Revision: D62106988

fbshipit-source-id: 7bbbe100ccaa8108d6801a30495f42d59333e347

starlark-rust/starlark_map/src/small_map.rs

@@ -724,6 +724,20 @@ impl<K, V> SmallMap<K, V> {
             }
         }
     }
+
+    /// Retains only the elements specified by the predicate.
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        let mut res = SmallMap::new();
+        for (k, mut v) in mem::take(self).into_iter_hashed() {
+            if f(&*k, &mut v) {
+                res.insert_hashed_unique_unchecked(k, v);
+            }
+        }
+        *self = res;
+    }
 }
 
 /// Reference to the actual entry in the map.
@@ -1386,4 +1400,20 @@ mod tests {
                 .collect::<Vec<_>>()
         );
     }
+
+    #[test]
+    fn test_retain() {
+        let mut map = SmallMap::new();
+        for i in 0..100 {
+            map.insert(i.to_string(), i);
+        }
+        map.retain(|_, v| {
+            let res = *v % 2 == 0;
+            *v += 3;
+            res
+        });
+        assert_eq!(map.len(), 50);
+        assert_eq!(map.get("7"), None);
+        assert_eq!(map.get("8"), Some(&11));
+    }
 }

starlark-rust/starlark_map/src/small_set.rs

@@ -405,6 +405,14 @@ impl<T> SmallSet<T> {
     pub fn reverse(&mut self) {
         self.0.reverse();
     }
+
+    /// Retains only the elements specified by the predicate.
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&T) -> bool,
+    {
+        self.0.retain(|k, _| f(k))
+    }
 }
 
 impl<'a, T> IntoIterator for &'a SmallSet<T> {