From 7764d72169d76644504dd0115182a9aed6a8e9aa Mon Sep 17 00:00:00 2001
From: Neil Mitchell <ndmitchell@meta.com>
Date: Mon, 2 Sep 2024 14:38:50 -0700
Subject: [PATCH] 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 | 30 +++++++++++++++++++++
 starlark-rust/starlark_map/src/small_set.rs |  8 ++++++
 2 files changed, 38 insertions(+)

diff --git a/starlark-rust/starlark_map/src/small_map.rs b/starlark-rust/starlark_map/src/small_map.rs
index 52c10033a4dd..fcbeb97d80dd 100644
--- a/starlark-rust/starlark_map/src/small_map.rs
+++ b/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));
+    }
 }
diff --git a/starlark-rust/starlark_map/src/small_set.rs b/starlark-rust/starlark_map/src/small_set.rs
index 332ea6cf4a5d..3b98cf213df2 100644
--- a/starlark-rust/starlark_map/src/small_set.rs
+++ b/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> {