diff --git a/src/iter.rs b/src/iter.rs index 3e0c1ff..e49cd20 100644 --- a/src/iter.rs +++ b/src/iter.rs @@ -1,12 +1,27 @@ use crate::{ - Key, Node, Poptrie, STRIDE, + Entry, Key, Node, Poptrie, STRIDE, bitmap::{PrefixId, StrideId}, value_index::ValueIndex, }; -use alloc::vec; +use alloc::{collections::btree_map, vec}; use alloc::{collections::btree_map::BTreeMap, vec::Vec}; impl FromIterator<((K, u8), V)> for Poptrie { + /// Creates a [`Poptrie`] from an iterator of `((key, prefix_length), value)` tuples. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let trie: Poptrie = [ + /// ((u32::from_be_bytes([10, 0, 0, 0]), 8), 8), + /// ((u32::from_be_bytes([10, 1, 0, 0]), 16), 16), + /// ].into_iter().collect(); + /// + /// assert_eq!(trie.lookup(u32::from_be_bytes([10, 1, 1, 1])), Some(&16)); + /// assert_eq!(trie.lookup(u32::from_be_bytes([10, 2, 1, 1])), Some(&8)); + /// ``` fn from_iter>(iter: I) -> Self { let mut poptrie = Self::new(); @@ -55,7 +70,7 @@ impl FromIterator<((K, u8), V)> for Poptrie { let prefix_id = PrefixId::from_key(key, key_offset, remaining_length); poptrie.entries[parent_node_index] - .insert(prefix_id, current_value_index); + .insert(prefix_id, ((key, len), current_value_index)); } // Last step allows us to calculate the leaves @@ -105,3 +120,358 @@ impl FromIterator<((K, u8), V)> for Poptrie { poptrie } } + +/// An owning iterator over the entries of a [`Poptrie`], in lexicographic +/// order of `(prefix_length, key)`. +/// +/// This `struct` is created by the [`into_iter`] method on [`Poptrie`] +/// (provided by the [`IntoIterator`] trait). See its documentation for more. +/// +/// [`into_iter`]: IntoIterator::into_iter +pub struct IntoIter { + entries: vec::IntoIter>>, + current: btree_map::IntoIter>, + values: alloc::vec::Vec>, +} + +impl Iterator for IntoIter { + type Item = ((K, u8), V); + + fn next(&mut self) -> Option { + loop { + for (_, ((key, key_len), value_index)) in &mut self.current { + if let Some(value) = + value_index.get().and_then(|idx| self.values[idx].take()) + { + return Some(((key, key_len), value)); + } + } + self.current = self.entries.next()?.into_iter(); + } + } +} + +impl IntoIterator for Poptrie { + type Item = ((K, u8), V); + type IntoIter = IntoIter; + + /// Consumes the trie and iterates over all `((key, prefix_length), value)` + /// tuples, in lexicographic order of `(prefix_length, key)`. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, "10/8"); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, "10.1/16"); + /// + /// let entries: Vec<_> = trie.into_iter().collect(); + /// assert_eq!(entries, [ + /// ((u32::from_be_bytes([10, 0, 0, 0]), 8), "10/8"), + /// ((u32::from_be_bytes([10, 1, 0, 0]), 16), "10.1/16"), + /// ]); + /// ``` + fn into_iter(self) -> Self::IntoIter { + let values = self.values.into_iter().map(Some).collect(); + let mut entries = self.entries.into_iter(); + let current = entries.next().unwrap_or_default().into_iter(); + + IntoIter { entries, current, values } + } +} + +/// A borrowing iterator over the entries of a [`Poptrie`], in lexicographic +/// order of `(prefix_length, key)`. +/// +/// This `struct` is created by the [`iter`] method on [`Poptrie`]. +/// See its documentation for more. +/// +/// [`iter`]: Poptrie::iter +pub struct Iter<'a, K: Key, V> { + entries: core::slice::Iter<'a, BTreeMap>>, + current: btree_map::Iter<'a, PrefixId, Entry>, + values: &'a [V], +} + +impl<'a, K: Key, V> Iterator for Iter<'a, K, V> { + type Item = ((&'a K, u8), &'a V); + + fn next(&mut self) -> Option { + loop { + for (_, ((key, key_len), value_index)) in &mut self.current { + if let Some(idx) = value_index.get() { + return Some(((key, *key_len), &self.values[idx])); + } + } + self.current = self.entries.next()?.iter(); + } + } +} + +impl Poptrie { + /// Iterates over all `((&key, prefix_length), &value)` tuples, in + /// lexicographic order of `(prefix_length, key)`. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, "10/8"); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, "10.1/16"); + /// + /// for ((key, len), val) in trie.iter() { + /// assert!(trie.contains_key(*key, len)); + /// } + /// ``` + pub fn iter(&self) -> Iter<'_, K, V> { + let mut entries = self.entries.iter(); + let current = entries.next().map(|m| m.iter()).unwrap_or_default(); + + Iter { entries, current, values: &self.values } + } +} + +/// A mutable borrowing iterator over the entries of a [`Poptrie`], in +/// lexicographic order of `(prefix_length, key)`. +/// +/// This `struct` is created by the [`iter_mut`] method on [`Poptrie`]. +/// See its documentation for more. +/// +/// [`iter_mut`]: Poptrie::iter_mut +pub struct IterMut<'a, K: Key, V> { + entries: core::slice::Iter<'a, BTreeMap>>, + current: btree_map::Iter<'a, PrefixId, Entry>, + values: &'a mut [V], +} + +impl<'a, K: Key, V> Iterator for IterMut<'a, K, V> { + type Item = ((&'a K, u8), &'a mut V); + + fn next(&mut self) -> Option { + loop { + for (_, ((key, key_len), value_index)) in &mut self.current { + if let Some(idx) = value_index.get() { + // SAFETY: Each ValueIndex is unique across all entries so + // no two yielded references alias. + return Some(((key, *key_len), unsafe { + &mut *self.values.as_mut_ptr().add(idx) + })); + } + } + self.current = self.entries.next()?.iter(); + } + } +} + +impl Poptrie { + /// Iterates mutably over all `((&key, prefix_length), &mut value)` tuples, + /// in lexicographic order of `(prefix_length, key)`. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 1u32); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 2u32); + /// + /// for ((key, len), val) in trie.iter_mut() { + /// *val *= 10; + /// } + /// + /// assert_eq!(trie.lookup(u32::from_be_bytes([10, 0, 1, 1])), Some(&10)); + /// assert_eq!(trie.lookup(u32::from_be_bytes([10, 1, 1, 1])), Some(&20)); + /// ``` + pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { + let Poptrie { entries, values, .. } = self; + let mut entries_iter = entries.iter(); + let current = entries_iter.next().map(|m| m.iter()).unwrap_or_default(); + IterMut { + entries: entries_iter, + current, + values: values.as_mut_slice(), + } + } +} + +/// An iterator over the keys of a [`Poptrie`], in lexicographic order of +/// `(prefix_length, key)`. +/// +/// This `struct` is created by the [`keys`] method on [`Poptrie`]. +/// See its documentation for more. +/// +/// [`keys`]: Poptrie::keys +pub struct Keys<'a, K: Key, V>(pub(crate) Iter<'a, K, V>); + +impl<'a, K: Key, V> Iterator for Keys<'a, K, V> { + type Item = (&'a K, u8); + + fn next(&mut self) -> Option { + self.0.next().map(|((k, l), _)| (k, l)) + } +} + +/// An iterator over the values of a [`Poptrie`], in lexicographic order of +/// `(prefix_length, key)`. +/// +/// This `struct` is created by the [`values`] method on [`Poptrie`]. +/// See its documentation for more. +/// +/// [`values`]: Poptrie::values +pub struct Values<'a, K: Key, V>(pub(crate) Iter<'a, K, V>); + +impl<'a, K: Key, V> Iterator for Values<'a, K, V> { + type Item = &'a V; + + fn next(&mut self) -> Option { + self.0.next().map(|(_, v)| v) + } +} + +/// A mutable iterator over the values of a [`Poptrie`], in lexicographic +/// order of `(prefix_length, key)`. +/// +/// This `struct` is created by the [`values_mut`] method on [`Poptrie`]. +/// See its documentation for more. +/// +/// [`values_mut`]: Poptrie::values_mut +pub struct ValuesMut<'a, K: Key, V>(pub(crate) IterMut<'a, K, V>); + +impl<'a, K: Key, V> Iterator for ValuesMut<'a, K, V> { + type Item = &'a mut V; + + fn next(&mut self) -> Option { + self.0.next().map(|(_, v)| v) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn iter_yields_all_entries() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + + let entries: Vec<_> = + trie.iter().map(|((k, l), v)| (*k, l, *v)).collect(); + assert_eq!( + entries, + [ + (u32::from_be_bytes([10, 0, 0, 0]), 8, 8), + (u32::from_be_bytes([10, 1, 0, 0]), 16, 16), + ] + ); + } + + #[test] + fn iter_empty_trie() { + assert_eq!(Poptrie::::new().iter().count(), 0); + } + + #[test] + fn into_iter_consumes_all_entries() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + + let entries: Vec<_> = trie.into_iter().collect(); + assert_eq!( + entries, + [ + ((u32::from_be_bytes([10, 0, 0, 0]), 8), 8), + ((u32::from_be_bytes([10, 1, 0, 0]), 16), 16), + ] + ); + } + + #[test] + fn into_iter_empty_trie() { + assert_eq!(Poptrie::::new().into_iter().count(), 0); + } + + #[test] + fn iter_mut_modifies_values() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 1u32); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 2u32); + + for (_, v) in trie.iter_mut() { + *v *= 10; + } + + assert_eq!(trie.lookup(u32::from_be_bytes([10, 0, 1, 1])), Some(&10)); + assert_eq!(trie.lookup(u32::from_be_bytes([10, 1, 1, 1])), Some(&20)); + } + + #[test] + fn iter_after_remove() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + trie.remove(u32::from_be_bytes([10, 1, 0, 0]), 16); + + let entries: Vec<_> = trie.iter().collect(); + assert_eq!(entries.len(), 1); + assert_eq!(*entries[0].0.0, u32::from_be_bytes([10, 0, 0, 0])); + assert_eq!(entries[0].0.1, 8); + } + + #[test] + fn from_iter_round_trips_with_into_iter() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + trie.insert(u32::from_be_bytes([10, 1, 2, 0]), 24, 24u32); + + let rebuilt: Poptrie = trie.into_iter().collect(); + + let entries: Vec<_> = rebuilt.into_iter().collect(); + assert_eq!( + entries, + [ + ((u32::from_be_bytes([10, 0, 0, 0]), 8), 8), + ((u32::from_be_bytes([10, 1, 0, 0]), 16), 16), + ((u32::from_be_bytes([10, 1, 2, 0]), 24), 24), + ] + ); + } + + #[test] + fn iter_contains_key_consistent() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + + for ((key, len), _) in trie.iter() { + assert!(trie.contains_key(*key, len)); + } + } + + #[test] + fn iter_is_sorted() { + let mut trie = Poptrie::new(); + trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + trie.insert(u32::from_be_bytes([192, 168, 0, 0]), 16, 160u32); + trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + trie.insert(u32::from_be_bytes([10, 1, 2, 0]), 24, 24u32); + + let entries: Vec<_> = trie.iter().map(|((k, l), _)| (*k, l)).collect(); + assert_eq!( + entries, + [ + (u32::from_be_bytes([10, 0, 0, 0]), 8), + (u32::from_be_bytes([10, 1, 0, 0]), 16), + (u32::from_be_bytes([192, 168, 0, 0]), 16), + (u32::from_be_bytes([10, 1, 2, 0]), 24), + ] + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 6464267..189b908 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -24,6 +24,7 @@ mod iter; mod key; mod value_index; +pub use iter::{IntoIter, Iter, IterMut, Keys, Values, ValuesMut}; pub use key::Key; use alloc::collections::btree_map::BTreeMap; @@ -31,7 +32,6 @@ use alloc::vec; use alloc::vec::Vec; use bitmap::*; use core::cmp::min; -use core::marker::PhantomData; use value_index::ValueIndex; /// The maximum number of bits we can consume from the prefix at a time. @@ -40,6 +40,10 @@ use value_index::ValueIndex; /// instruction exists. const STRIDE: u8 = 6; +/// A tuple representing a prefix entry in the trie, consisting of a key, prefix length, +/// and value index. +type Entry = ((K, u8), ValueIndex); + /// A compressed prefix tree optimized for fast longest prefix match (LPM) lookups. /// /// # Type Parameters @@ -81,10 +85,7 @@ where values: Vec, /// The entries associated with each node. - entries: Vec>, - - /// The marker for `K`. - _marker: PhantomData, + entries: Vec>>, } impl Poptrie @@ -114,7 +115,6 @@ where nodes: vec![root_node], // Start with a root node entries: vec![BTreeMap::new()], // Root's entries leaves: vec![ValueIndex::NONE], // Global default value index - _marker: PhantomData, } } @@ -214,8 +214,9 @@ where let remaining_length = key_length - key_offset; let prefix_id = PrefixId::from_key(key, key_offset, remaining_length); - // Store the value index for that prefix chunk - self.entries[parent_node_index].insert(prefix_id, current_value_index); + // Store the key and value index for that prefix chunk + self.entries[parent_node_index] + .insert(prefix_id, ((key, key_length), current_value_index)); // Update the defaults for children self.calculate_leaf_ranges(parent_node_index, default_value_index); @@ -302,13 +303,196 @@ where /// assert!(!trie.contains_key(u32::from_be_bytes([192, 168, 0, 0]), 16)); /// ``` pub fn contains_key(&self, key: K, key_length: u8) -> bool { - // Find the final parent node and prefix id let (parent_node, prefix_id, _) = self.find_parent_node(key, key_length); self.entries[parent_node].contains_key(&prefix_id) } + /// Returns the number of entries in the trie. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// assert_eq!(trie.len(), 0); + /// + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// assert_eq!(trie.len(), 1); + /// ``` + pub fn len(&self) -> usize { + self.values.len() + } + + /// Returns `true` if the trie contains no prefixes. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::::new(); + /// assert!(trie.is_empty()); + /// + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// assert!(!trie.is_empty()); + /// ``` + pub fn is_empty(&self) -> bool { + self.values.is_empty() + } + + /// Returns a reference to the value associated with the exact prefix + /// `(key, key_length)`, or `None` if it was not present. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// + /// assert_eq!(trie.get(u32::from_be_bytes([10, 0, 0, 0]), 8), Some(&8)); + /// assert_eq!(trie.get(u32::from_be_bytes([10, 0, 0, 0]), 16), None); + /// ``` + pub fn get(&self, key: K, key_length: u8) -> Option<&V> { + let (parent_node, prefix_id, _) = + self.find_parent_node(key, key_length); + + self.entries[parent_node] + .get(&prefix_id) + .and_then(|(_, vi)| vi.get().map(|i| &self.values[i])) + } + + /// Returns a mutable reference to the value associated with the exact + /// prefix `(key, key_length)`, or `None` if it was not present. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 1u32); + /// + /// if let Some(v) = trie.get_mut(u32::from_be_bytes([10, 0, 0, 0]), 8) { + /// *v *= 10; + /// } + /// + /// assert_eq!(trie.get(u32::from_be_bytes([10, 0, 0, 0]), 8), Some(&10)); + /// ``` + pub fn get_mut(&mut self, key: K, key_length: u8) -> Option<&mut V> { + let (parent_node, prefix_id, _) = + self.find_parent_node(key, key_length); + + self.entries[parent_node] + .get(&prefix_id) + .and_then(|(_, vi)| vi.get()) + .map(|i| &mut self.values[i]) + } + + /// Returns an iterator over the prefixes `(&key, prefix_length)` of the + /// trie, in lexicographic order of `(prefix_length, key)`. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + /// + /// let keys: Vec<_> = trie.keys().collect(); + /// assert_eq!(keys, [ + /// (&u32::from_be_bytes([10, 0, 0, 0]), 8), + /// (&u32::from_be_bytes([10, 1, 0, 0]), 16), + /// ]); + /// ``` + pub fn keys(&self) -> Keys<'_, K, V> { + Keys(self.iter()) + } + + /// Returns an iterator over the values of the trie, in lexicographic + /// order of `(prefix_length, key)`. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + /// + /// let values: Vec<_> = trie.values().collect(); + /// assert_eq!(values, [&8, &16]); + /// ``` + pub fn values(&self) -> Values<'_, K, V> { + Values(self.iter()) + } + + /// Returns a mutable iterator over the values of the trie, in lexicographic + /// order of `(prefix_length, key)`. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 1u32); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 2u32); + /// + /// for v in trie.values_mut() { + /// *v *= 10; + /// } + /// + /// assert_eq!(trie.get(u32::from_be_bytes([10, 0, 0, 0]), 8), Some(&10)); + /// assert_eq!(trie.get(u32::from_be_bytes([10, 1, 0, 0]), 16), Some(&20)); + /// ``` + pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> { + ValuesMut(self.iter_mut()) + } + + /// Retains only the entries for which the predicate returns `true`, + /// removing all others. + /// + /// # Examples + /// + /// ``` + /// use poptrie::Poptrie; + /// + /// let mut trie = Poptrie::new(); + /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// trie.insert(u32::from_be_bytes([10, 1, 0, 0]), 16, 16u32); + /// trie.insert(u32::from_be_bytes([10, 1, 2, 0]), 24, 24u32); + /// + /// trie.retain(|_, v| *v <= 16); + /// + /// assert_eq!(trie.len(), 2); + /// assert!(trie.contains_key(u32::from_be_bytes([10, 0, 0, 0]), 8)); + /// assert!(trie.contains_key(u32::from_be_bytes([10, 1, 0, 0]), 16)); + /// assert!(!trie.contains_key(u32::from_be_bytes([10, 1, 2, 0]), 24)); + /// ``` + pub fn retain(&mut self, mut f: F) + where + F: FnMut(&K, &mut V) -> bool, + { + let to_remove: Vec<_> = self + .iter_mut() + .filter_map( + |((k, l), v)| if !f(k, v) { Some((*k, l)) } else { None }, + ) + .collect(); + + for (key, len) in to_remove { + self.remove(key, len); + } + } + /// Removes and returns the value associated with the exact prefix /// `(key, key_length)`, or `None` if it was not present. /// @@ -331,11 +515,10 @@ where /// assert!(!trie.contains_key(u32::from_be_bytes([10, 1, 0, 0]), 16)); /// ``` pub fn remove(&mut self, key: K, key_length: u8) -> Option { - // Find the final parent node and prefix id let (parent_node, prefix_id, default_value_index) = self.find_parent_node(key, key_length); - self.entries[parent_node].remove(&prefix_id).map(|v| { + self.entries[parent_node].remove(&prefix_id).map(|(_, v)| { // Update the leaf ranges self.calculate_leaf_ranges(parent_node, default_value_index); @@ -343,7 +526,11 @@ where for higher_v in self .leaves .iter_mut() - .chain(self.entries.iter_mut().flat_map(|s| s.values_mut())) + .chain( + self.entries + .iter_mut() + .flat_map(|s| s.values_mut().map(|(_, vi)| vi)), + ) .filter(|higher_v| **higher_v > v) { higher_v.decrement(); @@ -361,7 +548,6 @@ where key: K, key_length: u8, ) -> (usize, PrefixId, ValueIndex) { - // Traverse the trie similarly to a insert but check the entry let mut key_offset = 0; let mut parent_node_index = 0; let mut parent_node = &self.nodes[parent_node_index]; @@ -371,17 +557,13 @@ where let local_id = StrideId::from_key(key, key_offset, STRIDE); default_value_index = self.get_default(parent_node_index, local_id); - // Check if there's already a node with `local_id` if !parent_node.node_bitmap.contains(local_id) { - // If there's no node with `local_id`, break out break; } - // Traverse to the child node parent_node_index = parent_node.get_child_index(local_id); parent_node = &self.nodes[parent_node_index]; - // Advance the key offset key_offset += STRIDE; } @@ -392,9 +574,6 @@ where } /// Find the next base node and leaf node index for a given parent node index. - /// - /// Used for inserting a new node while keeping the order of nodes. - /// It will get the latest descendant node that's before the new node. fn find_next_base(&self, next_node_index: usize) -> (u32, u32) { // SAFETY: We start with a root node at 0 let last_node = &self.nodes[next_node_index - 1]; @@ -430,8 +609,10 @@ where .map(|p| self.get_default(node_index, StrideId(*p))) .collect(); - let (new_bitmap, new_leaves) = - build_leaf_ranges(&self.entries[node_index], default_value_index); + let (new_bitmap, new_leaves) = build_leaf_ranges( + self.entries[node_index].iter().map(|(&pid, &(_, vi))| (pid, vi)), + default_value_index, + ); let old_end = if node_index < self.nodes.len() - 1 { self.nodes[node_index + 1].leaf_base as usize @@ -490,13 +671,13 @@ where let default = entries .peek() .take_if(|(p, _)| p.prefix_length() == 0) - .map(|(_, v)| **v) + .map(|(_, (_, v))| *v) .unwrap_or(default_value_index); self.leaves.insert(leaf_base, default); leaf_bitmap.set(StrideId(0)); - for (prefix_id, value) in entries { + for (prefix_id, (_, value)) in entries { let (prefix, len) = prefix_id.components(); let leaf_id = prefix_id.stride_id(); let leafvec_index = leaf_bitmap.leafvec_index(leaf_id); @@ -590,17 +771,16 @@ impl Node { // - For `Poptrie::insert`, multiple leaves may have to added, where each insert pushes the following leaves // - We could not always shrink, trading a little cache locality for insertion speed. fn build_leaf_ranges( - entries: &BTreeMap, + entries: impl Iterator, default_value_index: ValueIndex, ) -> (Bitmap, Vec) { let mut leaf_bitmap = Bitmap::new(); - // Prepare default - let mut entries = entries.iter().peekable(); + let mut entries = entries.peekable(); let default = entries .peek() .take_if(|(p, _)| p.prefix_length() == 0) - .map(|(_, v)| **v) + .map(|(_, v)| *v) .unwrap_or(default_value_index); let mut leaves = vec![default]; @@ -612,13 +792,14 @@ fn build_leaf_ranges( let leafvec_index = leaf_bitmap.leafvec_index(leaf_id) as usize; let initial_value = leaves[leafvec_index]; - // Insert new leaf or rewrite a possible terminator let leaf_bitmap_index = leaf_bitmap.bitmap_index(leaf_id) as usize; + + // Insert new leaf or rewrite a possible terminator if !leaf_bitmap.contains(leaf_id) { - leaves.insert(leaf_bitmap_index, *value); + leaves.insert(leaf_bitmap_index, value); leaf_bitmap.set(leaf_id); } else { - leaves[leaf_bitmap_index] = *value; + leaves[leaf_bitmap_index] = value; } let next_id = StrideId((prefix + 1) << (STRIDE - len));