remove optimization and prune

Author: ekinimo

Diff for: src/proptest_triemap.rs

@@ -2,6 +2,7 @@ use super::*;
 use proptest::prelude::*;
 use std::collections::BTreeMap;
 use std::collections::HashMap;
+use std::collections::HashSet;
 
 fn key_value_pairs(
     min_pairs: usize,
@@ -163,7 +164,178 @@ fn test_empty_triemap_behavior() {
 }
 
 proptest! {
-  #[test]
+
+              #[test]
+        fn pruning_preserves_values(pairs in key_value_pairs(1, 100), to_remove in key_value_pairs(1, 50)) {
+            let mut trie = TrieMap::new();
+            let mut reference_map = BTreeMap::new();
+
+            for (key, value) in &pairs {
+                trie.insert(key, *value);
+                reference_map.insert(key.clone(), *value);
+            }
+
+            for (key, _) in &to_remove {
+                trie.remove(key);
+                reference_map.remove(key);
+            }
+
+            trie.prune();
+
+            for (key, expected_value) in &reference_map {
+                prop_assert_eq!(trie.get(key), Some(expected_value));
+            }
+
+            for (key, _) in &to_remove {
+                if !reference_map.contains_key(key) {
+                    prop_assert_eq!(trie.get(key), None);
+                }
+            }
+
+            prop_assert_eq!(trie.len(), reference_map.len());
+        }
+        // Test that multiple prune operations are idempotent
+        #[test]
+        fn multiple_prunes_are_idempotent(
+            pairs in key_value_pairs(1, 100),
+            to_remove in key_value_pairs(1, 50)
+        ) {
+            let mut trie = TrieMap::new();
+
+            // Insert all pairs
+            for (key, value) in &pairs {
+                trie.insert(key, *value);
+            }
+
+            // Remove some keys
+            for (key, _) in &to_remove {
+                trie.remove(key);
+            }
+
+            // First prune
+            let first_pruned = trie.prune();
+            let size_after_first = trie.len();
+
+            // Second prune
+            let second_pruned = trie.prune();
+            let size_after_second = trie.len();
+
+            // The second prune should not remove any nodes
+            prop_assert_eq!(second_pruned, 0);
+
+            // Sizes should be the same
+            prop_assert_eq!(size_after_first, size_after_second);
+        }
+
+            #[test]
+        fn removal_works_correctly(pairs in key_value_pairs(5, 100)) {
+            let mut trie = TrieMap::new();
+            let mut reference_map = BTreeMap::new();
+
+            // Insert all pairs
+            for (key, value) in &pairs {
+                trie.insert(key, *value);
+                reference_map.insert(key.clone(), *value);
+            }
+
+            // Get the unique keys that were actually inserted (last value wins for duplicates)
+            let unique_keys: Vec<String> = reference_map.keys().cloned().collect();
+
+            // Remove half of the unique keys
+            let mut removed = 0;
+            for (i, key) in unique_keys.iter().enumerate() {
+                if i % 2 == 0 {  // Only remove every other key
+                    let expected_value = reference_map.get(key).copied();
+                    let trie_removed = trie.remove(key);
+                    let ref_removed = reference_map.remove(key);
+
+                    // The removed value should match between trie and reference map
+                    prop_assert_eq!(trie_removed, ref_removed);
+                    // And should match what we expected to remove
+                    prop_assert_eq!(trie_removed, expected_value);
+
+                    removed += 1;
+                }
+            }
+
+            // Check that the size is correct
+            prop_assert_eq!(trie.len(), reference_map.len());
+
+            // Check that all remaining keys are accessible
+            for (key, value) in &reference_map {
+                prop_assert_eq!(trie.get(key), Some(value));
+            }
+
+            // Check that removed keys are not accessible
+            for (i, key) in unique_keys.iter().enumerate() {
+                if i % 2 == 0 {  // These were removed
+                    prop_assert_eq!(trie.get(key), None);
+                }
+            }
+        }
+
+        #[test]
+        fn iteration_after_removal_is_correct(
+            pairs in key_value_pairs(5, 100),
+            removal_indices in proptest::collection::vec(0..100usize, 1..50)
+        ) {
+            let mut trie = TrieMap::new();
+            let mut reference_map = BTreeMap::new();
+
+            // Insert all pairs
+            for (key, value) in &pairs {
+                trie.insert(key, *value);
+                reference_map.insert(key.clone(), *value);
+            }
+
+            // Create list of keys to remove (using indices to prevent duplicates)
+            let keys_to_remove: Vec<String> = removal_indices.iter()
+                .filter(|&idx| *idx < pairs.len())
+                .map(|&idx| pairs[idx].0.clone())
+                .collect();
+
+            // Remove selected keys
+            for key in &keys_to_remove {
+                trie.remove(key);
+                reference_map.remove(key);
+            }
+
+            // Get all key-value pairs from trie iteration
+            let mut trie_pairs: Vec<(String, i32)> = trie.iter()
+                .map(|(k, &v)| (String::from_utf8(k).unwrap(), v))
+                .collect();
+
+            // Get all key-value pairs from reference map
+            let mut ref_pairs: Vec<(String, i32)> = reference_map
+                .iter()
+                .map(|(k, &v)| (k.clone(), v))
+                .collect();
+
+            // Sort both for comparison
+            trie_pairs.sort_by(|(k1, _), (k2, _)| k1.cmp(k2));
+            ref_pairs.sort_by(|(k1, _), (k2, _)| k1.cmp(k2));
+
+            // They should be equal
+            prop_assert_eq!(trie_pairs, ref_pairs);
+
+            // Size should match
+            prop_assert_eq!(trie.len(), reference_map.len());
+
+            // Check iterator count matches expected count
+            prop_assert_eq!(trie.iter().count(), reference_map.len());
+
+            // Also check keys() and values() iterators
+            prop_assert_eq!(trie.keys().count(), reference_map.len());
+            prop_assert_eq!(trie.values().count(), reference_map.len());
+
+            // Verify removed keys are not present
+            for key in &keys_to_remove {
+                prop_assert_eq!(trie.get(key), None);
+                prop_assert!(!trie.keys().any(|k| String::from_utf8(k).unwrap() == *key));
+            }
+        }
+
+    #[test]
     fn triemap_correctly_handles_common_prefixes(
         pairs in prefixed_keys(vec!["app", "ban", "car", "dog"], 5, 50)
     ) {

Diff for: src/trie_map.rs

@@ -530,7 +530,9 @@ impl<T> TrieMap<T> {
     }
 
     /// Removes a key from the map, returning the value at the key if the key was previously in the map.
-    ///
+    /// This does not remove the nodes used to register the key
+    /// Use `remove_and_prune` method for removing intermittent nodes as well.
+    /// Use `prune` method to remove *all* nodes that leads to a thombstone
     /// # Examples
     ///
     /// ```
@@ -548,6 +550,56 @@ impl<T> TrieMap<T> {
     }
 
     fn remove_internal(&mut self, bytes: &[u8]) -> Option<T> {
+        let mut current = &mut self.root;
+        let mut found = true;
+
+        for &byte in bytes {
+            if !test_bit(&current.is_present, byte) {
+                found = false;
+                break;
+            }
+            let idx = popcount(&current.is_present, byte) as usize;
+            if idx >= current.children.len() {
+                found = false;
+                break;
+            }
+            current = &mut current.children[idx];
+        }
+
+        if found && current.data_idx.is_some() {
+            let data_idx = current.data_idx.unwrap();
+
+            if data_idx < self.data.len() && self.data[data_idx].is_some() {
+                let value = self.data[data_idx].take();
+                current.data_idx = None;
+                self.free_indices.push(data_idx);
+                self.size -= 1;
+                return value;
+            }
+        }
+
+        None
+    }
+
+    /// Removes a key from the map, returning the value at the key if the key was previously in the map.
+    /// This method also removes the nodes used to register the key
+    /// # Examples
+    ///
+    /// ```
+    /// # use triemap::TrieMap;
+    /// let mut map = TrieMap::new();
+    /// map.insert("a", 1);
+    ///
+    /// assert_eq!(map.remove("a"), Some(1));
+    /// assert_eq!(map.remove("a"), None);
+    /// ```
+    pub fn remove_and_prune<K: AsBytes>(&mut self, key: K) -> Option<T> {
+        let bytes = key.as_bytes();
+
+        self.remove_and_prune_internal(bytes)
+    }
+
+    fn remove_and_prune_internal(&mut self, bytes: &[u8]) -> Option<T> {
         let mut path = Vec::with_capacity(bytes.len());
         let mut path_indices = Vec::with_capacity(bytes.len());
 
@@ -618,6 +670,82 @@ impl<T> TrieMap<T> {
         }
     }
 
+    /// Prunes unused nodes from the trie to reclaim memory.
+    ///
+    /// This method removes all nodes that don't contain values and don't lead to nodes with values.
+    /// It's useful to call periodically if you've removed many items from the trie.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use triemap::TrieMap;
+    /// let mut map = TrieMap::new();
+    /// map.insert("apple", 1);
+    /// map.insert("application", 2);
+    ///
+    /// map.remove("apple");
+    /// map.remove("application");
+    ///
+    /// // The trie structure still contains nodes for "apple" and "application"
+    /// // even though the values have been removed
+    ///
+    /// map.prune();
+    /// // Now the unused nodes have been removed
+    /// ```
+    pub fn prune(&mut self) -> usize {
+        Self::prune_node(&mut self.root)
+    }
+
+    // Helper method to recursively prune nodes
+    fn prune_node(node: &mut TrieNode) -> usize {
+        let mut pruned_nodes = 0;
+        let mut bytes_to_clear = Vec::new();
+
+        // Check each byte in the is_present array
+        for byte in 0..=255u8 {
+            if test_bit(&node.is_present, byte) {
+                let idx = popcount(&node.is_present, byte) as usize;
+                if idx < node.children.len() {
+                    // Recursively prune the child node
+                    let child_pruned = Self::prune_node(&mut node.children[idx]);
+                    pruned_nodes += child_pruned;
+
+                    // Check if the child node is now empty and can be removed
+                    if node.children[idx].data_idx.is_none()
+                        && node.children[idx].children.is_empty()
+                    {
+                        bytes_to_clear.push(byte);
+                    }
+                }
+            }
+        }
+
+        // Remove empty children that were marked for removal
+        for &byte in &bytes_to_clear {
+            let idx = popcount(&node.is_present, byte) as usize;
+
+            // Create a new children array without the empty node
+            let mut new_children = Vec::with_capacity(node.children.len() - 1);
+
+            // Copy all children except the one being removed
+            for i in 0..node.children.len() {
+                if i != idx {
+                    new_children.push(std::mem::replace(&mut node.children[i], TrieNode::new()));
+                }
+            }
+
+            // Update the node's children
+            node.children = new_children.into_boxed_slice();
+
+            // Update the is_present bits - need to clear the bit for the removed node
+            clear_bit(&mut node.is_present, byte);
+
+            // Update counts
+            pruned_nodes += 1;
+        }
+
+        pruned_nodes
+    }
     /// Returns an iterator over the key-value pairs of the map.
     ///
     /// # Examples
@@ -1638,8 +1766,7 @@ impl<T> TrieMap<T> {
         &'a self,
         other: &'a TrieMap<T>,
     ) -> impl Iterator<Item = (Vec<u8>, &'a T)> + 'a {
-        self.iter()
-            .filter(move |(key, _)| other.contains_key(key))
+        self.iter().filter(move |(key, _)| other.contains_key(key))
     }
 
     /// Returns an iterator over the entries whose keys are in this map but not in the other map.
@@ -1667,8 +1794,7 @@ impl<T> TrieMap<T> {
         &'a self,
         other: &'a TrieMap<T>,
     ) -> impl Iterator<Item = (Vec<u8>, &'a T)> + 'a {
-        self.iter()
-            .filter(move |(key, _)| !other.contains_key(key))
+        self.iter().filter(move |(key, _)| !other.contains_key(key))
     }
 
     /// Returns an iterator over entries whose keys are in exactly one of the maps.