From a960730add4ad4d3d5248187da6a14a78b621097 Mon Sep 17 00:00:00 2001
From: Jakob Stoklund Olesen <jolesen@mozilla.com>
Date: Wed, 22 Nov 2017 13:43:41 -0800
Subject: [PATCH] Add retain() methods to bforest::{Set,Map}.

These work just like their counterparts in HashMap and HashSet.
---
 lib/cretonne/src/bforest/map.rs | 37 +++++++++++++++++++++++++++++++++
 lib/cretonne/src/bforest/set.rs | 22 ++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/lib/cretonne/src/bforest/map.rs b/lib/cretonne/src/bforest/map.rs
index 601e9febe6..bf288582eb 100644
--- a/lib/cretonne/src/bforest/map.rs
+++ b/lib/cretonne/src/bforest/map.rs
@@ -130,6 +130,32 @@ where
         }
     }
 
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// Remove all key-value pairs where the predicate returns false.
+    ///
+    /// The predicate is allowed to update the values stored in the map.
+    pub fn retain<F>(&mut self, forest: &mut MapForest<K, V, C>, mut predicate: F)
+    where
+        F: FnMut(K, &mut V) -> bool,
+    {
+        let mut path = Path::default();
+        if let Some(root) = self.root.expand() {
+            path.first(root, &forest.nodes);
+        }
+        while let Some((node, entry)) = path.leaf_pos() {
+            let keep = {
+                let (ks, vs) = forest.nodes[node].unwrap_leaf_mut();
+                predicate(ks[entry], &mut vs[entry])
+            };
+            if keep {
+                path.next(&forest.nodes);
+            } else {
+                self.root = path.remove(&mut forest.nodes).into();
+            }
+        }
+    }
+
     /// Create a cursor for navigating this map. The cursor is initially positioned off the end of
     /// the map.
     pub fn cursor<'a>(
@@ -394,6 +420,7 @@ mod test {
 
         assert_eq!(m.get(7, &f, &()), None);
         assert_eq!(m.iter(&f).next(), None);
+        m.retain(&mut f, |_, _| unreachable!());
 
         let mut c = m.cursor(&mut f, &());
         assert!(c.is_empty());
@@ -525,6 +552,16 @@ mod test {
         m.verify(f, &());
         assert_eq!(m.get(5, f, &()), Some(4.2));
 
+        // Retain even entries, with altered values.
+        m.retain(f, |k, v| {
+            *v = (k / 10) as f32;
+            (k % 20) == 0
+        });
+        assert_eq!(
+            m.iter(f).collect::<Vec<_>>(),
+            [(20, 2.0), (40, 4.0), (60, 6.0)]
+        );
+
         // Insert at back of leaf.
         let mut m = full_leaf(f);
         m.insert(80, 4.2, f, &());
diff --git a/lib/cretonne/src/bforest/set.rs b/lib/cretonne/src/bforest/set.rs
index 025efc0225..6992265773 100644
--- a/lib/cretonne/src/bforest/set.rs
+++ b/lib/cretonne/src/bforest/set.rs
@@ -127,6 +127,26 @@ where
         }
     }
 
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// Remove all elements where the predicate returns false.
+    pub fn retain<F>(&mut self, forest: &mut SetForest<K, C>, mut predicate: F)
+    where
+        F: FnMut(K) -> bool,
+    {
+        let mut path = Path::default();
+        if let Some(root) = self.root.expand() {
+            path.first(root, &forest.nodes);
+        }
+        while let Some((node, entry)) = path.leaf_pos() {
+            if predicate(forest.nodes[node].unwrap_leaf().0[entry]) {
+                path.next(&forest.nodes);
+            } else {
+                self.root = path.remove(&mut forest.nodes).into();
+            }
+        }
+    }
+
     /// Create a cursor for navigating this set. The cursor is initially positioned off the end of
     /// the set.
     pub fn cursor<'a>(
@@ -354,6 +374,8 @@ mod test {
         // Iterator for an empty set.
         assert_eq!(s.iter(&f).next(), None);
 
+        s.retain(&mut f, |_| unreachable!());
+
         let mut c = SetCursor::new(&mut s, &mut f, &());
         c.verify();
         assert_eq!(c.elem(), None);