diff --git a/README.md b/README.md index e03a77a..ccd76aa 100644 --- a/README.md +++ b/README.md @@ -58,10 +58,10 @@ assert_eq!(trie.lookup(u32::from_be_bytes([192, 168, 1, 5])), Some(&"192.168.0.0 | Function | Description | |---|---| | `new()` | Construct a new, empty poptrie | -| `insert(key, key_length, value)` | Insert a value associated with the given prefix | -| `lookup(key)` | Longest-prefix match lookup, returns `Option<&V>` | -| `contains_key(key, key_length)` | Returns `true` if the exact prefix is present | -| `remove(key, key_length)` | Remove a prefix and return its value | +| `insert(prefix, value)` | Insert a value associated with the given prefix | +| `lookup(address)` | Longest-prefix match lookup, returns `Option<&V>` | +| `contains_key(prefix)` | Returns `true` if the exact prefix is present | +| `remove(prefix)` | Remove a prefix and return its value | ### Lookup performance diff --git a/benches/lpm_benches.rs b/benches/lpm_benches.rs index 0a422e6..34ae1e2 100644 --- a/benches/lpm_benches.rs +++ b/benches/lpm_benches.rs @@ -97,8 +97,7 @@ fn bench_insert(c: &mut Criterion) { let mut poptrie = Poptrie::new(); for &((prefix, length), val) in prefixes { poptrie.insert( - black_box(prefix), - black_box(length), + black_box((prefix, length)), black_box(val), ); } diff --git a/src/key.rs b/src/address.rs similarity index 77% rename from src/key.rs rename to src/address.rs index bf1ce01..c356e96 100644 --- a/src/key.rs +++ b/src/address.rs @@ -1,19 +1,24 @@ #![allow(clippy::unusual_byte_groupings)] // Grouped 6 by 6 because that's the current STRIDE use core::ops::{Shl, Shr}; -/// A trait for types that can be used as keys in the [poptrie](crate::Poptrie). +/// A trait for types that can be used as addresses in the `Prefix`. /// -/// This is currently not sealed to allow for custom `Key` types. +/// This is currently not sealed to allow for custom types. /// Future versions may change this, which would be a breaking change. /// /// - Needs to inform its bit width and a `to_u8` method that returns its 8 least significant bits. /// - Needs to implement `rotate_right` method that rotates the key by `n` bits to the right. /// +/// # Implementation Details +/// +/// The current bounds are required to guarantee performance, else there would be a more flexible +/// trait with their own bit extraction functions. This may change in the future. +/// /// # Examples /// -/// Implementing `Key` for a custom newtype: +/// Implementing `Address` for a custom newtype: /// ``` -/// use poptrie::Key; +/// use poptrie::Address; /// use core::ops::{Shl, Shr}; /// /// #[derive(Clone, Copy)] @@ -29,7 +34,7 @@ use core::ops::{Shl, Shr}; /// fn shr(self, n: u8) -> Self { MyAddr(self.0 >> n) } /// } /// -/// impl Key for MyAddr { +/// impl Address for MyAddr { /// const BITS: u8 = 32; /// fn to_u8(self) -> u8 { self.0 as u8 } /// fn rotate_right(self, n: u32) -> Self { MyAddr(self.0.rotate_right(n)) } @@ -38,18 +43,18 @@ use core::ops::{Shl, Shr}; /// assert_eq!(MyAddr::BITS, 32); /// assert_eq!(MyAddr(0xdeadbeef).to_u8(), 0xef); /// ``` -pub trait Key: +pub trait Address: Copy + Shl + Shr + Sized { - /// The number of bits in the key. + /// The number of bits in the address. const BITS: u8; - /// Converts the least significant bits of the key to a u8. + /// Converts the least significant bits of the address to a u8. fn to_u8(self) -> u8; - /// Rotates the key to the right by `n` bits. + /// Rotates the address to the right by `n` bits. fn rotate_right(self, n: u32) -> Self; } -impl Key for u32 { +impl Address for u32 { const BITS: u8 = 32; #[inline(always)] fn to_u8(self) -> u8 { @@ -62,7 +67,7 @@ impl Key for u32 { } } -impl Key for u128 { +impl Address for u128 { const BITS: u8 = 128; #[inline(always)] fn to_u8(self) -> u8 { @@ -84,7 +89,7 @@ impl Key for u128 { /// ^^^^^^^^^^^^^^ /// extracted bits ///``` -/// If `len` + `offset` > K::BITS, the extraction will be zero-padded from the right: +/// If `len` + `offset` > A::BITS, the extraction will be zero-padded from the right: ///``` text /// MSB |------------------------key-------------------------| LSB /// |---------------- offset -----------------|----- len ----| @@ -92,11 +97,11 @@ impl Key for u128 { /// extracted bits ///``` #[inline(always)] -pub(crate) fn extract_bits(key: K, offset: u8, len: u8) -> u8 +pub(crate) fn extract_bits(address: A, offset: u8, len: u8) -> u8 where - K: Key, + A: Address, { - (key.rotate_right((K::BITS - offset).wrapping_sub(len) as u32)).to_u8() + (address.rotate_right((A::BITS - offset).wrapping_sub(len) as u32)).to_u8() & ((1u16 << len) - 1) as u8 } @@ -110,7 +115,7 @@ where /// ^^^^^^^^^^^^^^ /// extracted bits ///``` -/// If `len` + `offset` > `K::BITS`, the extraction will be saturated: +/// If `len` + `offset` > `A::BITS`, the extraction will be saturated: /// /// ``` text /// MSB |------------------------key-------------------------| LSB @@ -119,16 +124,16 @@ where /// extracted bits ///``` #[inline(always)] -pub(crate) fn extract_bits_saturated(key: K, offset: u8, len: u8) -> u8 +pub(crate) fn extract_bits_saturated(address: A, offset: u8, len: u8) -> u8 where - K: Key, + A: Address, { // TODO: Check if mask version is faster - let remaining = u8::saturating_sub(K::BITS - offset, len); - if remaining + offset == K::BITS { + let remaining = u8::saturating_sub(A::BITS - offset, len); + if remaining + offset == A::BITS { return 0; } - (key << offset >> (remaining + offset)).to_u8() + (address << offset >> (remaining + offset)).to_u8() } #[cfg(test)] diff --git a/src/bitmap.rs b/src/bitmap.rs index 6b65128..e7fabfa 100644 --- a/src/bitmap.rs +++ b/src/bitmap.rs @@ -2,7 +2,7 @@ use alloc::{collections::btree_map::BTreeMap, vec::Vec}; use crate::{ STRIDE, - key::{Key, extract_bits, extract_bits_saturated}, + address::{Address, extract_bits, extract_bits_saturated}, }; /// A generic bitmap for storing u8 encoded ids of 0..63 @@ -79,8 +79,12 @@ impl PrefixId { PrefixId((1u8 << len) - 1 + prefix) } - pub(crate) fn from_key(key: K, key_offset: u8, stride: u8) -> Self { - let prefix = extract_bits_saturated(key, key_offset, stride); + pub(crate) fn from_address( + address: A, + offset: u8, + stride: u8, + ) -> Self { + let prefix = extract_bits_saturated(address, offset, stride); PrefixId::new(prefix, stride) } @@ -137,12 +141,12 @@ where #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub(crate) struct StrideId(pub(crate) u8); impl StrideId { - pub(crate) fn from_key( - key: K, - key_offset: u8, + pub(crate) fn from_address( + address: A, + offset: u8, length: u8, ) -> StrideId { - StrideId(extract_bits(key, key_offset, length)) + StrideId(extract_bits(address, offset, length)) } } @@ -209,11 +213,11 @@ mod tests { #[test] fn test_prefix_to_stride() { - for key in 0u32..256 { + for address in 0u32..256 { for stride in 0u8..STRIDE { // PrefixIds are never full strides - let prefix = PrefixId::from_key(key, 0, stride); - let stride = StrideId::from_key(key, 0, stride); + let prefix = PrefixId::from_address(address, 0, stride); + let stride = StrideId::from_address(address, 0, stride); assert_eq!(prefix.stride_id(), stride); } } diff --git a/src/iter.rs b/src/iter.rs index e49cd20..403897c 100644 --- a/src/iter.rs +++ b/src/iter.rs @@ -1,20 +1,20 @@ use crate::{ - Entry, Key, Node, Poptrie, STRIDE, + Entry, Node, Poptrie, Prefix, STRIDE, bitmap::{PrefixId, StrideId}, value_index::ValueIndex, }; 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. +impl FromIterator<(P, V)> for Poptrie { + /// Creates a [`Poptrie`] from an iterator of `(prefix, value)` tuples. /// /// # Examples /// /// ``` /// use poptrie::Poptrie; /// - /// let trie: Poptrie = [ + /// let trie: Poptrie<(u32, u8), u32> = [ /// ((u32::from_be_bytes([10, 0, 0, 0]), 8), 8), /// ((u32::from_be_bytes([10, 1, 0, 0]), 16), 16), /// ].into_iter().collect(); @@ -22,17 +22,21 @@ impl FromIterator<((K, u8), V)> for Poptrie { /// 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 { + fn from_iter>(iter: I) -> Self { let mut poptrie = Self::new(); let mut items: Vec<_> = iter .into_iter() - .map(|((key, len), value)| { + .map(|(prefix, value)| { + let address = prefix.address(); + let len = prefix.prefix_length(); let path: Vec<_> = (0..(len / STRIDE)) - .map(|i| StrideId::from_key(key, i * STRIDE, STRIDE)) + .map(|i| { + StrideId::from_address(address, i * STRIDE, STRIDE) + }) .collect(); // Let's add the path and the last parent - (path, 0, len, key, value) + (path, 0usize, prefix, address, len, value) }) .collect(); @@ -52,7 +56,7 @@ impl FromIterator<((K, u8), V)> for Poptrie { while !items.is_empty() { // Remove all leaves for this level and add them to entries - for (path, mut parent_node_index, len, key, value) in + for (path, mut parent_node_index, prefix, address, len, value) in items.extract_if(.., |(path, ..)| path.len() <= level) { poptrie.values.push(value); @@ -65,12 +69,12 @@ impl FromIterator<((K, u8), V)> for Poptrie { .get_child_index(local_id); } - let key_offset = path.len() as u8 * STRIDE; - let remaining_length = len - key_offset; + let offset = path.len() as u8 * STRIDE; + let remaining_length = len - offset; let prefix_id = - PrefixId::from_key(key, key_offset, remaining_length); + PrefixId::from_address(address, offset, remaining_length); poptrie.entries[parent_node_index] - .insert(prefix_id, ((key, len), current_value_index)); + .insert(prefix_id, (prefix, current_value_index)); } // Last step allows us to calculate the leaves @@ -128,22 +132,22 @@ impl FromIterator<((K, u8), V)> for 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>, +pub struct IntoIter { + entries: vec::IntoIter>>, + current: btree_map::IntoIter>, values: alloc::vec::Vec>, } -impl Iterator for IntoIter { - type Item = ((K, u8), V); +impl Iterator for IntoIter { + type Item = (P, V); fn next(&mut self) -> Option { loop { - for (_, ((key, key_len), value_index)) in &mut self.current { + for (_, (prefix, 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)); + return Some((prefix, value)); } } self.current = self.entries.next()?.into_iter(); @@ -151,12 +155,12 @@ impl Iterator for IntoIter { } } -impl IntoIterator for Poptrie { - type Item = ((K, u8), V); - type IntoIter = IntoIter; +impl IntoIterator for Poptrie { + type Item = (P, V); + type IntoIter = IntoIter; - /// Consumes the trie and iterates over all `((key, prefix_length), value)` - /// tuples, in lexicographic order of `(prefix_length, key)`. + /// Consumes the trie and iterates over all `(prefix, value)` tuples, in + /// lexicographic order of `(prefix_length, key)`. /// /// # Examples /// @@ -164,8 +168,8 @@ impl IntoIterator for Poptrie { /// 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"); + /// 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, [ @@ -189,20 +193,20 @@ impl IntoIterator for 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>, +pub struct Iter<'a, P: Prefix, 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); +impl<'a, P: Prefix, V> Iterator for Iter<'a, P, V> { + type Item = (&'a P, &'a V); fn next(&mut self) -> Option { loop { - for (_, ((key, key_len), value_index)) in &mut self.current { + for (_, (prefix, value_index)) in &mut self.current { if let Some(idx) = value_index.get() { - return Some(((key, *key_len), &self.values[idx])); + return Some((prefix, &self.values[idx])); } } self.current = self.entries.next()?.iter(); @@ -210,9 +214,9 @@ impl<'a, K: Key, V> Iterator for Iter<'a, K, V> { } } -impl Poptrie { - /// Iterates over all `((&key, prefix_length), &value)` tuples, in - /// lexicographic order of `(prefix_length, key)`. +impl Poptrie { + /// Iterates over all `(&prefix, &value)` pairs, in lexicographic order of + /// `(prefix_length, key)`. /// /// # Examples /// @@ -220,17 +224,16 @@ impl Poptrie { /// 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"); + /// 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)); + /// for (prefix, val) in trie.iter() { + /// assert!(trie.contains_key(*prefix)); /// } /// ``` - pub fn iter(&self) -> Iter<'_, K, V> { + pub fn iter(&self) -> Iter<'_, P, V> { let mut entries = self.entries.iter(); let current = entries.next().map(|m| m.iter()).unwrap_or_default(); - Iter { entries, current, values: &self.values } } } @@ -242,22 +245,22 @@ impl 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>, +pub struct IterMut<'a, P: Prefix, 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); +impl<'a, P: Prefix, V> Iterator for IterMut<'a, P, V> { + type Item = (&'a P, &'a mut V); fn next(&mut self) -> Option { loop { - for (_, ((key, key_len), value_index)) in &mut self.current { + for (_, (prefix, 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 { + return Some((prefix, unsafe { &mut *self.values.as_mut_ptr().add(idx) })); } @@ -267,9 +270,9 @@ impl<'a, K: Key, V> Iterator for IterMut<'a, K, V> { } } -impl Poptrie { - /// Iterates mutably over all `((&key, prefix_length), &mut value)` tuples, - /// in lexicographic order of `(prefix_length, key)`. +impl Poptrie { + /// Iterates mutably over all `(&prefix, &mut value)` pairs, in lexicographic + /// order of `(prefix_length, key)`. /// /// # Examples /// @@ -277,17 +280,17 @@ impl Poptrie { /// 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); + /// 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() { + /// for (_, 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> { + pub fn iter_mut(&mut self) -> IterMut<'_, P, V> { let Poptrie { entries, values, .. } = self; let mut entries_iter = entries.iter(); let current = entries_iter.next().map(|m| m.iter()).unwrap_or_default(); @@ -306,13 +309,13 @@ impl Poptrie { /// See its documentation for more. /// /// [`keys`]: Poptrie::keys -pub struct Keys<'a, K: Key, V>(pub(crate) Iter<'a, K, V>); +pub struct Keys<'a, P: Prefix, V>(pub(crate) Iter<'a, P, V>); -impl<'a, K: Key, V> Iterator for Keys<'a, K, V> { - type Item = (&'a K, u8); +impl<'a, P: Prefix, V> Iterator for Keys<'a, P, V> { + type Item = &'a P; fn next(&mut self) -> Option { - self.0.next().map(|((k, l), _)| (k, l)) + self.0.next().map(|(p, _)| p) } } @@ -323,9 +326,9 @@ impl<'a, K: Key, V> Iterator for Keys<'a, K, V> { /// See its documentation for more. /// /// [`values`]: Poptrie::values -pub struct Values<'a, K: Key, V>(pub(crate) Iter<'a, K, V>); +pub struct Values<'a, P: Prefix, V>(pub(crate) Iter<'a, P, V>); -impl<'a, K: Key, V> Iterator for Values<'a, K, V> { +impl<'a, P: Prefix, V> Iterator for Values<'a, P, V> { type Item = &'a V; fn next(&mut self) -> Option { @@ -340,9 +343,9 @@ impl<'a, K: Key, V> Iterator for Values<'a, K, V> { /// See its documentation for more. /// /// [`values_mut`]: Poptrie::values_mut -pub struct ValuesMut<'a, K: Key, V>(pub(crate) IterMut<'a, K, V>); +pub struct ValuesMut<'a, P: Prefix, V>(pub(crate) IterMut<'a, P, V>); -impl<'a, K: Key, V> Iterator for ValuesMut<'a, K, V> { +impl<'a, P: Prefix, V> Iterator for ValuesMut<'a, P, V> { type Item = &'a mut V; fn next(&mut self) -> Option { @@ -357,30 +360,29 @@ mod tests { #[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); + 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(); + let entries: Vec<_> = trie.iter().map(|(p, v)| (*p, *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), + ((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); + assert_eq!(Poptrie::<(u32, u8), u32>::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); + 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!( @@ -394,14 +396,14 @@ mod tests { #[test] fn into_iter_empty_trie() { - assert_eq!(Poptrie::::new().into_iter().count(), 0); + assert_eq!(Poptrie::<(u32, u8), u32>::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); + 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; @@ -414,24 +416,23 @@ mod tests { #[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); + 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); + assert_eq!(*entries[0].0, (u32::from_be_bytes([10, 0, 0, 0]), 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); + 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 rebuilt: Poptrie<(u32, u8), u32> = trie.into_iter().collect(); let entries: Vec<_> = rebuilt.into_iter().collect(); assert_eq!( @@ -447,23 +448,23 @@ mod tests { #[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); + 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)); + for (prefix, _) in trie.iter() { + assert!(trie.contains_key(*prefix)); } } #[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); + 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(); + let entries: Vec<_> = trie.iter().map(|(p, _)| *p).collect(); assert_eq!( entries, [ diff --git a/src/lib.rs b/src/lib.rs index 1f4c1e2..2b00fea 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,13 +19,15 @@ #![no_std] extern crate alloc; +mod address; mod bitmap; mod iter; -mod key; +mod prefix; mod value_index; +pub use address::Address; pub use iter::{IntoIter, Iter, IterMut, Keys, Values, ValuesMut}; -pub use key::Key; +pub use prefix::Prefix; use alloc::collections::btree_map::BTreeMap; use alloc::vec; @@ -40,16 +42,15 @@ 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 tuple representing a prefix entry in the trie, consisting of a prefix and value index. +type Entry

= (P, ValueIndex); /// A compressed prefix tree optimized for fast longest prefix match (LPM) lookups. /// /// # Type Parameters /// -/// * `K`: [`Key`] - The key type (e.g., `u32` for IPv4, `u128` for IPv6) -/// * `V` - The value type associated with each prefix +/// * `P`: [`Prefix`] - The prefix type (e.g. `(u32, u8)` for IPv4 or `(u128, u8)` for IPv6), +/// * `V` - The value type associated with each prefix. /// /// # Examples /// @@ -57,12 +58,11 @@ type Entry = ((K, u8), ValueIndex); /// use poptrie::Poptrie; /// /// // Create a routing table for IPv4 addresses -/// let mut trie = Poptrie::::new(); +/// let mut trie = Poptrie::<(u32, u8), &str>::new(); /// -/// // Insert prefixes with their associated values -/// trie.insert(u32::from_be_bytes([192, 168, 0, 0]), 16, "192.168.0.0/16"); -/// trie.insert(u32::from_be_bytes([192, 168, 1, 0]), 24, "192.168.1.0/24"); -/// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, "10.0.0.0/8"); +/// trie.insert((u32::from_be_bytes([192, 168, 0, 0]), 16), "192.168.0.0/16"); +/// trie.insert((u32::from_be_bytes([192, 168, 1, 0]), 24), "192.168.1.0/24"); +/// trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), "10.0.0.0/8"); /// /// // Perform longest prefix match lookups /// assert_eq!(trie.lookup(u32::from_be_bytes([192, 168, 1, 5])), Some(&"192.168.1.0/24")); @@ -71,9 +71,9 @@ type Entry = ((K, u8), ValueIndex); /// assert_eq!(trie.lookup(u32::from_be_bytes([8, 8, 8, 8])), None); /// ``` #[derive(Debug, Clone, Default)] -pub struct Poptrie +pub struct Poptrie where - K: Key, + P: Prefix, { /// The internal nodes of the trie. nodes: Vec, @@ -85,12 +85,12 @@ where values: Vec, /// The entries associated with each node. - entries: Vec>>, + entries: Vec>>, } -impl Poptrie +impl Poptrie where - K: Key, + P: Prefix, { /// Construct a new, empty poptrie. /// @@ -99,7 +99,7 @@ where /// ``` /// use poptrie::Poptrie; /// - /// let trie = Poptrie::::new(); + /// let trie = Poptrie::<(u32, u8), ()>::new(); /// ``` pub fn new() -> Self { let mut root_node = Node::new( @@ -110,7 +110,7 @@ where ); // Register the default value bit root_node.leaf_bitmap.set(StrideId(0)); - Poptrie:: { + Poptrie:: { values: Vec::new(), nodes: vec![root_node], // Start with a root node entries: vec![BTreeMap::new()], // Root's entries @@ -123,7 +123,7 @@ where /// /// # Panics /// - /// Panics if `key_length > K::BITS`. + /// Panics if `prefix.prefix_length() > P::ADDRESS::BITS`. /// /// # Examples /// @@ -133,27 +133,25 @@ where /// let mut trie = Poptrie::new(); /// /// // Insert a /8 prefix - /// assert_eq!(trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, "10.0.0.0/8"), None); + /// assert_eq!(trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), "10.0.0.0/8"), None); /// /// // Insert a more specific /24 prefix - /// assert_eq!(trie.insert(u32::from_be_bytes([10, 1, 2, 0]), 24, "10.1.2.0/24"), None); + /// assert_eq!(trie.insert((u32::from_be_bytes([10, 1, 2, 0]), 24), "10.1.2.0/24"), None); /// /// // Replacing an existing prefix returns the old value - /// assert_eq!(trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, "new"), Some("10.0.0.0/8")); + /// assert_eq!(trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), "new"), Some("10.0.0.0/8")); /// /// // Insert a default route (0-length prefix matches everything) - /// assert_eq!(trie.insert(0u32, 0, "default"), None); + /// assert_eq!(trie.insert((0u32, 0), "default"), None); /// ``` - pub fn insert( - &mut self, - key: K, - key_length: u8, - mut value: V, - ) -> Option { - assert!(key_length <= K::BITS); + pub fn insert(&mut self, prefix: P, mut value: V) -> Option { + let key = prefix.address(); + let prefix_length = prefix.prefix_length(); + + assert!(prefix_length <= P::ADDRESS::BITS); let mut default_value_index = ValueIndex::NONE; - let mut key_offset = 0; + let mut offset = 0; // First node is root let mut parent_node_index = 0; @@ -161,8 +159,8 @@ where // Check if it's in the correct depth // We MUST use '>=' here to ensure that full strides always direct towards inner nodes. - while key_length >= key_offset + STRIDE { - let local_id = StrideId::from_key(key, key_offset, STRIDE); + while prefix_length >= offset + STRIDE { + let local_id = StrideId::from_address(key, offset, STRIDE); let full_node_index = parent_node.get_child_index(local_id); // Find the default from the parent @@ -212,12 +210,12 @@ where parent_node_index = full_node_index; parent_node = &self.nodes[parent_node_index]; - key_offset += STRIDE; + offset += STRIDE; } // Can't consume a whole STRIDE, so we handle the remainder. - let remaining_length = key_length - key_offset; - let prefix_id = PrefixId::from_key(key, key_offset, remaining_length); + let remaining_length = prefix_length - offset; + let prefix_id = PrefixId::from_address(key, offset, remaining_length); // If an entry already exists, reuse it and return the old value let old_value = if let Some(idx) = self.entries[parent_node_index] @@ -231,7 +229,7 @@ where let current_value_index = ValueIndex::new((self.values.len() - 1) as u32); self.entries[parent_node_index] - .insert(prefix_id, ((key, key_length), current_value_index)); + .insert(prefix_id, (prefix, current_value_index)); None }; @@ -241,7 +239,7 @@ where old_value } - /// Lookup a key in the trie, performing longest-prefix match. + /// Lookup an address in the trie, performing longest-prefix match. /// /// Returns `None` if no prefix matches the key. /// @@ -255,9 +253,9 @@ where /// // No match without a default route /// assert_eq!(trie.lookup(u32::from_be_bytes([8, 8, 8, 8])), None); /// - /// trie.insert(0u32, 0, "default"); - /// 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"); + /// trie.insert((0u32, 0), "default"); + /// 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"); /// /// // Longest prefix match: 10.1.2.3 matches 10.1/16 /// assert_eq!(trie.lookup(u32::from_be_bytes([10, 1, 2, 3])), Some(&"10.1/16")); @@ -268,13 +266,14 @@ where /// // Falls back to default /// assert_eq!(trie.lookup(u32::from_be_bytes([8, 8, 8, 8])), Some(&"default")); /// ``` - pub fn lookup(&self, key: K) -> Option<&V> { - let mut key_offset = 0; + pub fn lookup(&self, address: P::ADDRESS) -> Option<&V> { + // lookup_inner(&self.nodes, &self.leaves, key).map(|i| &self.values[i]) + let mut offset = 0; // First node is root let mut parent_node_index = 0; let mut parent_node = &self.nodes[parent_node_index]; - let mut local_id = StrideId::from_key(key, key_offset, STRIDE); + let mut local_id = StrideId::from_address(address, offset, STRIDE); // Should always try internal nodes first. while parent_node.node_bitmap.contains(local_id) { @@ -283,8 +282,8 @@ where parent_node = &self.nodes[parent_node_index]; // Update key offset and local ID - key_offset += STRIDE; - local_id = StrideId::from_key(key, key_offset, STRIDE); + offset += STRIDE; + local_id = StrideId::from_address(address, offset, STRIDE); } // There will always be at least a 0th leaf (e.g. with the default) @@ -296,12 +295,7 @@ where value_index.get().map(|i| &self.values[i]) } - /// Returns `true` if the trie contains an entry for the exact prefix - /// `(key, key_length)`. - /// - /// Note that this checks for an exact prefix match, not a longest-prefix - /// match. A key that would resolve via [`lookup`](Self::lookup) may still - /// return `false` here if it was never explicitly inserted with that length. + /// Returns `true` if the trie contains an entry for the exact prefix. /// /// # Examples /// @@ -309,22 +303,14 @@ where /// use poptrie::Poptrie; /// /// let mut trie = Poptrie::new(); - /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); - /// - /// // Exact prefix is present - /// assert!(trie.contains_key(u32::from_be_bytes([10, 0, 0, 0]), 8)); + /// trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), 8u32); /// - /// // Different length — not present, even though lookup would succeed - /// assert!(!trie.contains_key(u32::from_be_bytes([10, 0, 0, 0]), 7)); - /// assert!(!trie.contains_key(u32::from_be_bytes([10, 0, 0, 0]), 9)); - /// - /// // Completely absent prefix - /// assert!(!trie.contains_key(u32::from_be_bytes([192, 168, 0, 0]), 16)); + /// assert!(trie.contains_key((u32::from_be_bytes([10, 0, 0, 0]), 8))); + /// assert!(!trie.contains_key((u32::from_be_bytes([10, 0, 0, 0]), 7))); + /// assert!(!trie.contains_key((u32::from_be_bytes([192, 168, 0, 0]), 16))); /// ``` - pub fn contains_key(&self, key: K, key_length: u8) -> bool { - let (parent_node, prefix_id, _) = - self.find_parent_node(key, key_length); - + pub fn contains_key(&self, prefix: P) -> bool { + let (parent_node, prefix_id, _) = self.find_parent_node(prefix); self.entries[parent_node].contains_key(&prefix_id) } @@ -335,10 +321,10 @@ where /// ``` /// use poptrie::Poptrie; /// - /// let mut trie = Poptrie::new(); + /// let mut trie = Poptrie::<(u32, u8), u32>::new(); /// assert_eq!(trie.len(), 0); /// - /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), 8u32); /// assert_eq!(trie.len(), 1); /// ``` pub fn len(&self) -> usize { @@ -352,18 +338,18 @@ where /// ``` /// use poptrie::Poptrie; /// - /// let mut trie = Poptrie::::new(); + /// let mut trie = Poptrie::<(u32, u8), u32>::new(); /// assert!(trie.is_empty()); /// - /// trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 8u32); + /// 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. + /// Returns a reference to the value associated with the exact prefix, or + /// `None` if it was not present. /// /// # Examples /// @@ -371,22 +357,20 @@ where /// 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, 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); + /// 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); - + pub fn get(&self, prefix: P) -> Option<&V> { + let (parent_node, prefix_id, _) = self.find_parent_node(prefix); 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. + /// prefix, or `None` if it was not present. /// /// # Examples /// @@ -394,26 +378,23 @@ where /// 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, 0, 0, 0]), 8), 1u32); /// - /// if let Some(v) = trie.get_mut(u32::from_be_bytes([10, 0, 0, 0]), 8) { + /// 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)); + /// 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); - + pub fn get_mut(&mut self, prefix: P) -> Option<&mut V> { + let (parent_node, prefix_id, _) = self.find_parent_node(prefix); 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)`. + /// Returns an iterator over the prefixes of the trie, in lexicographic order of `(prefix_length, address)`. /// /// # Examples /// @@ -421,21 +402,21 @@ where /// 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, 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), + /// &(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> { + pub fn keys(&self) -> Keys<'_, P, V> { Keys(self.iter()) } /// Returns an iterator over the values of the trie, in lexicographic - /// order of `(prefix_length, key)`. + /// order of `(prefix_length, address)`. /// /// # Examples /// @@ -443,18 +424,17 @@ where /// 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, 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> { + pub fn values(&self) -> Values<'_, P, V> { Values(self.iter()) } - /// Returns a mutable iterator over the values of the trie, in lexicographic - /// order of `(prefix_length, key)`. + /// Returns a mutable iterator over the values of the trie. /// /// # Examples /// @@ -462,22 +442,19 @@ where /// 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); + /// 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; - /// } + /// 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)); + /// 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> { + pub fn values_mut(&mut self) -> ValuesMut<'_, P, V> { ValuesMut(self.iter_mut()) } - /// Retains only the entries for which the predicate returns `true`, - /// removing all others. + /// Retains only the entries for which the predicate returns `true`. /// /// # Examples /// @@ -485,35 +462,33 @@ where /// 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.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)); + /// 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, + F: FnMut(&P, &mut V) -> bool, { let to_remove: Vec<_> = self .iter_mut() - .filter_map( - |((k, l), v)| if !f(k, v) { Some((*k, l)) } else { None }, - ) + .filter_map(|(p, v)| if !f(p, v) { Some(*p) } else { None }) .collect(); - for (key, len) in to_remove { - self.remove(key, len); + for prefix in to_remove { + self.remove(prefix); } } - /// Removes and returns the value associated with the exact prefix - /// `(key, key_length)`, or `None` if it was not present. + /// Removes and returns the value associated with the exact prefix, or + /// `None` if it was not present. /// /// # Examples /// @@ -521,21 +496,21 @@ where /// 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, 0, 0, 0]), 8), 8u32); + /// trie.insert((u32::from_be_bytes([10, 1, 0, 0]), 16), 16u32); /// /// // Remove the /16 prefix - /// assert_eq!(trie.remove(u32::from_be_bytes([10, 1, 0, 0]), 16), Some(16)); + /// assert_eq!(trie.remove((u32::from_be_bytes([10, 1, 0, 0]), 16)), Some(16)); /// /// // Addresses previously matched by /16 now fall back to /8 /// assert_eq!(trie.lookup(u32::from_be_bytes([10, 1, 2, 3])), Some(&8)); /// - /// // Removing a prefix that doesn't exist returns None - /// assert!(!trie.contains_key(u32::from_be_bytes([10, 1, 0, 0]), 16)); + /// // Removing a prefix that doesn't exist returns `None` + /// assert!(!trie.contains_key((u32::from_be_bytes([10, 1, 0, 0]), 16))); /// ``` - pub fn remove(&mut self, key: K, key_length: u8) -> Option { + pub fn remove(&mut self, prefix: P) -> Option { let (parent_node, prefix_id, default_value_index) = - self.find_parent_node(key, key_length); + self.find_parent_node(prefix); self.entries[parent_node].remove(&prefix_id).map(|(_, v)| { // Update the leaf ranges @@ -562,18 +537,16 @@ where } /// Find the final parent node and the `PrefixId` of the given key. - fn find_parent_node( - &self, - key: K, - key_length: u8, - ) -> (usize, PrefixId, ValueIndex) { - let mut key_offset = 0; + fn find_parent_node(&self, prefix: P) -> (usize, PrefixId, ValueIndex) { + let address = prefix.address(); + let prefix_length = prefix.prefix_length(); + let mut offset = 0; let mut parent_node_index = 0; let mut parent_node = &self.nodes[parent_node_index]; let mut default_value_index = ValueIndex::NONE; - while key_length >= key_offset + STRIDE { - let local_id = StrideId::from_key(key, key_offset, STRIDE); + while prefix_length >= offset + STRIDE { + let local_id = StrideId::from_address(address, offset, STRIDE); default_value_index = self.get_default(parent_node_index, local_id); if !parent_node.node_bitmap.contains(local_id) { @@ -583,11 +556,12 @@ where parent_node_index = parent_node.get_child_index(local_id); parent_node = &self.nodes[parent_node_index]; - key_offset += STRIDE; + offset += STRIDE; } - let remaining_length = min(key_length - key_offset, STRIDE - 1); - let prefix_id = PrefixId::from_key(key, key_offset, remaining_length); + let remaining_length = min(prefix_length - offset, STRIDE - 1); + let prefix_id = + PrefixId::from_address(address, offset, remaining_length); (parent_node_index, prefix_id, default_value_index) } @@ -761,10 +735,7 @@ impl Node { /// Returns the index of the child node pointed by `local_id`. #[inline(always)] fn get_child_index(&self, local_id: StrideId) -> usize { - let node_base = self.node_base; - let node_offset = self.node_bitmap.bitmap_index(local_id); - - (node_base + node_offset) as usize + (self.node_base + self.node_bitmap.bitmap_index(local_id)) as usize } } diff --git a/src/prefix.rs b/src/prefix.rs new file mode 100644 index 0000000..b503556 --- /dev/null +++ b/src/prefix.rs @@ -0,0 +1,53 @@ +use crate::address::Address; + +/// A prefix is a pattern to match the beginning of a sequence, in this case called an +/// `ADDRESS`. The prefix is represented by an `ADDRESS` and a `prefix_length` that +/// defines the number of bits in the prefix (counted from the most significant bits). +/// +/// # Examples +/// +/// ``` +/// use poptrie::{Poptrie, Prefix}; +/// +/// let mut trie = Poptrie::<(u32, u8), &str>::new(); +/// trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), "10/8"); +/// assert_eq!(trie.lookup(u32::from_be_bytes([10, 1, 2, 3])), Some(&"10/8")); +/// ``` +/// +/// Implementing `Prefix` for a custom type: +/// ``` +/// use poptrie::{Prefix}; +/// +/// #[derive(Clone, Copy, PartialEq, Eq)] +/// struct Ipv4Net { addr: u32, len: u8 } +/// +/// impl Prefix for Ipv4Net { +/// type ADDRESS = u32; +/// fn address(&self) -> u32 { self.addr } +/// fn prefix_length(&self) -> u8 { self.len } +/// } +/// ``` +pub trait Prefix: Copy { + /// The underlying address type. + type ADDRESS: Address; + + /// Returns the address part of the prefix. + fn address(&self) -> Self::ADDRESS; + + /// Returns the prefix length (number of significant bits). + fn prefix_length(&self) -> u8; +} + +impl Prefix for (A, u8) { + type ADDRESS = A; + + #[inline(always)] + fn address(&self) -> Self::ADDRESS { + self.0 + } + + #[inline(always)] + fn prefix_length(&self) -> u8 { + self.1 + } +} diff --git a/tests/api.rs b/tests/api.rs index d2a9ae3..df632bd 100644 --- a/tests/api.rs +++ b/tests/api.rs @@ -13,7 +13,7 @@ use poptrie::Poptrie; #[test] fn empty_trie_returns_none() { - let trie = Poptrie::::new(); + let trie = Poptrie::<(u32, u8), ()>::new(); assert_eq!(trie.lookup(0u32), None); assert_eq!(trie.lookup(u32::MAX), None); } @@ -25,7 +25,7 @@ fn empty_trie_returns_none() { #[test] fn default_route_matches_everything() { let mut trie = Poptrie::new(); - trie.insert(0u32, 0, 0); + trie.insert((0u32, 0), 0); assert_eq!(trie.lookup(0u32), Some(&0)); assert_eq!(trie.lookup(u32::MAX), Some(&0)); assert_eq!( @@ -37,8 +37,8 @@ fn default_route_matches_everything() { #[test] fn specific_prefix_overrides_default_route() { let mut trie = Poptrie::new(); - trie.insert(0u32, 0, 0); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 1); + trie.insert((0u32, 0), 0); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 1); // Inside the /6 prefix assert_eq!( trie.lookup(0b000001_111111_000000_000000_000000_00u32), @@ -54,10 +54,10 @@ fn specific_prefix_overrides_default_route() { #[test] fn default_route_overwrite_propagates_to_children() { let mut trie = Poptrie::new(); - trie.insert(0u32, 0, 0); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 15, 2); + trie.insert((0u32, 0), 0); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 15), 2); // Overwrite the default - trie.insert(0u32, 0, 1); + trie.insert((0u32, 0), 1); // The /15 child is still there assert_eq!( trie.lookup(0b000001_000001_000000_000000_000000_00u32), @@ -78,7 +78,7 @@ fn default_route_overwrite_propagates_to_children() { fn prefix_at_first_stride_boundary() { // /6: the entire first stride let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); assert_eq!( trie.lookup(0b000001_111111_000000_000000_000000_00u32), Some(&6) @@ -90,7 +90,7 @@ fn prefix_at_first_stride_boundary() { fn prefix_at_second_stride_boundary() { // /12: exactly two strides consumed let mut trie = Poptrie::new(); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); assert_eq!( trie.lookup(0b000001_000001_111111_000000_000000_00u32), Some(&12) @@ -102,7 +102,7 @@ fn prefix_at_second_stride_boundary() { fn prefix_spanning_two_stride_levels() { // /7: one full stride (6 bits) plus 1 bit into the next let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 7, 7); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 7), 7); // Same first stride, second bit of second stride = 0 → match assert_eq!( trie.lookup(0b000001_011111_000000_000000_000000_00u32), @@ -119,9 +119,9 @@ fn prefix_spanning_two_stride_levels() { #[test] fn lpm_selects_deepest_matching_prefix() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); - trie.insert(0b000001_000001_000001_000000_000000_00u32, 18, 18); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); + trie.insert((0b000001_000001_000001_000000_000000_00u32, 18), 18); // Matches /18 assert_eq!( @@ -145,8 +145,8 @@ fn lpm_selects_deepest_matching_prefix() { #[test] fn lpm_sibling_strides_do_not_bleed() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 12, 1); - trie.insert(0b000001_000010_000000_000000_000000_00u32, 12, 2); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 12), 1); + trie.insert((0b000001_000010_000000_000000_000000_00u32, 12), 2); assert_eq!( trie.lookup(0b000001_000000_111111_000000_000000_00u32), @@ -164,9 +164,9 @@ fn lpm_sibling_strides_do_not_bleed() { fn lpm_child_prefix_inserted_before_parent() { // Insert deeper prefix first, shallower prefix second let mut trie = Poptrie::new(); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 13, 13); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 12, 12); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 7, 7); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 13), 13); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 12), 12); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 7), 7); // Address outside all three prefixes assert_eq!(trie.lookup(0b000001_100001_100000_000000_000000_00u32), None); } @@ -178,8 +178,8 @@ fn lpm_child_prefix_inserted_before_parent() { #[test] fn two_prefixes_one_stride_apart_with_shared_parent() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000001_000001_000001_000000_00u32, 25, 25); - trie.insert(0b000001_000001_000001_000001_000011_01u32, 32, 32); + trie.insert((0b000001_000001_000001_000001_000000_00u32, 25), 25); + trie.insert((0b000001_000001_000001_000001_000011_01u32, 32), 32); assert_eq!( trie.lookup(0b000001_000001_000001_000001_000001_00u32), Some(&25) @@ -197,8 +197,8 @@ fn two_prefixes_one_stride_apart_with_shared_parent() { #[test] fn default_at_first_level_propagates_into_second_level_node() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); - trie.insert(0b000000_000000_000000_000000_000000_00u32, 12, 12); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); + trie.insert((0b000000_000000_000000_000000_000000_00u32, 12), 12); // The /6 is shallower and its stride doesn't overlap with /12's stride assert_eq!( trie.lookup(0b000001_000000_000000_000000_000000_00u32), @@ -209,8 +209,8 @@ fn default_at_first_level_propagates_into_second_level_node() { #[test] fn default_at_second_level_propagates_into_third_level_node() { let mut trie = Poptrie::new(); - trie.insert(0b000000_110000_000000_000000_000000_00u32, 12, 12); - trie.insert(0b000000_000000_000000_000000_000000_00u32, 18, 18); + trie.insert((0b000000_110000_000000_000000_000000_00u32, 12), 12); + trie.insert((0b000000_000000_000000_000000_000000_00u32, 18), 18); assert_eq!( trie.lookup(0b000000_110000_000000_000000_000000_00u32), Some(&12) @@ -220,8 +220,8 @@ fn default_at_second_level_propagates_into_third_level_node() { #[test] fn non_full_length_default_does_not_cover_sibling_subtree() { let mut trie = Poptrie::new(); - trie.insert(0b001100_000000_000000_000000_000000_00u32, 10, 10); - trie.insert(0b000000_000000_000000_000000_000000_00u32, 1, 1); + trie.insert((0b001100_000000_000000_000000_000000_00u32, 10), 10); + trie.insert((0b000000_000000_000000_000000_000000_00u32, 1), 1); // This address shares the first stride's leading bit with /1 but not /10 assert_eq!( trie.lookup(0b001100_000100_000000_000000_000000_00u32), @@ -236,22 +236,28 @@ fn non_full_length_default_does_not_cover_sibling_subtree() { #[test] fn contains_key_true_for_inserted_prefix() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); - assert!(trie.contains_key(0b000001_000001_000000_000000_000000_00u32, 12)); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); + assert!( + trie.contains_key((0b000001_000001_000000_000000_000000_00u32, 12)) + ); } #[test] fn contains_key_false_for_different_prefix_length() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); - assert!(!trie.contains_key(0b000001_000001_000000_000000_000000_00u32, 11)); - assert!(!trie.contains_key(0b000001_000001_000000_000000_000000_00u32, 13)); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); + assert!( + !trie.contains_key((0b000001_000001_000000_000000_000000_00u32, 11)) + ); + assert!( + !trie.contains_key((0b000001_000001_000000_000000_000000_00u32, 13)) + ); } #[test] fn contains_key_false_on_empty_trie() { - let trie = Poptrie::::new(); - assert!(!trie.contains_key(0u32, 0)); + let trie = Poptrie::<(u32, u8), ()>::new(); + assert!(!trie.contains_key((0u32, 0))); } // ------------------------------------------------------------------------- @@ -261,9 +267,9 @@ fn contains_key_false_on_empty_trie() { #[test] fn insert_returns_old_value() { let mut trie = Poptrie::new(); - assert_eq!(trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 1u32), None); + assert_eq!(trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), 1u32), None); assert_eq!( - trie.insert(u32::from_be_bytes([10, 0, 0, 0]), 8, 2u32), + trie.insert((u32::from_be_bytes([10, 0, 0, 0]), 8), 2u32), Some(1) ); assert_eq!(trie.lookup(u32::from_be_bytes([10, 0, 1, 1])), Some(&2)); @@ -276,9 +282,9 @@ fn insert_returns_old_value() { #[test] fn remove_returns_value_and_clears_lookup() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); assert_eq!( - trie.remove(0b000001_000000_000000_000000_000000_00u32, 6), + trie.remove((0b000001_000000_000000_000000_000000_00u32, 6)), Some(6) ); assert_eq!(trie.lookup(0b000001_111111_000000_000000_000000_00u32), None); @@ -287,9 +293,9 @@ fn remove_returns_value_and_clears_lookup() { #[test] fn remove_specific_falls_back_to_less_specific() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); - trie.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); - trie.remove(0b000001_000001_000000_000000_000000_00u32, 12); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); + trie.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); + trie.remove((0b000001_000001_000000_000000_000000_00u32, 12)); // After /12 is gone, its range falls back to /6 assert_eq!( trie.lookup(0b000001_000001_111111_000000_000000_00u32), @@ -300,9 +306,9 @@ fn remove_specific_falls_back_to_less_specific() { #[test] fn remove_and_reinsert_gives_new_value() { let mut trie = Poptrie::new(); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); - trie.remove(0b000001_000000_000000_000000_000000_00u32, 6); - trie.insert(0b000001_000000_000000_000000_000000_00u32, 6, 99); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); + trie.remove((0b000001_000000_000000_000000_000000_00u32, 6)); + trie.insert((0b000001_000000_000000_000000_000000_00u32, 6), 99); assert_eq!( trie.lookup(0b000001_111111_000000_000000_000000_00u32), Some(&99) @@ -316,14 +322,14 @@ fn remove_and_reinsert_gives_new_value() { #[test] fn insert_order_does_not_affect_lpm() { let mut forward = Poptrie::new(); - forward.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); - forward.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); - forward.insert(0b000001_000001_000001_000000_000000_00u32, 18, 18); + forward.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); + forward.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); + forward.insert((0b000001_000001_000001_000000_000000_00u32, 18), 18); let mut reverse = Poptrie::new(); - reverse.insert(0b000001_000001_000001_000000_000000_00u32, 18, 18); - reverse.insert(0b000001_000001_000000_000000_000000_00u32, 12, 12); - reverse.insert(0b000001_000000_000000_000000_000000_00u32, 6, 6); + reverse.insert((0b000001_000001_000001_000000_000000_00u32, 18), 18); + reverse.insert((0b000001_000001_000000_000000_000000_00u32, 12), 12); + reverse.insert((0b000001_000000_000000_000000_000000_00u32, 6), 6); let lookups = [ 0b000001_000001_000001_111111_000000_00u32, @@ -347,7 +353,7 @@ fn ip6(s: &str) -> u128 { #[test] fn ipv6_basic_lookup() { let mut trie = Poptrie::new(); - trie.insert(ip6("2001:db8::"), 32, 32u32); + trie.insert((ip6("2001:db8::"), 32), 32u32); assert_eq!(trie.lookup(ip6("2001:db8::1")), Some(&32)); assert_eq!(trie.lookup(ip6("2001:db9::1")), None); } @@ -355,8 +361,8 @@ fn ipv6_basic_lookup() { #[test] fn ipv6_lpm_selects_most_specific() { let mut trie = Poptrie::new(); - trie.insert(ip6("2001:db8::"), 32, 32u32); - trie.insert(ip6("2001:db8:dead::"), 48, 48u32); + trie.insert((ip6("2001:db8::"), 32), 32u32); + trie.insert((ip6("2001:db8:dead::"), 48), 48u32); assert_eq!(trie.lookup(ip6("2001:db8:dead::1")), Some(&48)); assert_eq!(trie.lookup(ip6("2001:db8:beef::1")), Some(&32)); @@ -366,8 +372,8 @@ fn ipv6_lpm_selects_most_specific() { #[test] fn ipv6_default_route() { let mut trie = Poptrie::new(); - trie.insert(ip6("::"), 0, 0u32); - trie.insert(ip6("2001:db8::"), 32, 32u32); + trie.insert((ip6("::"), 0), 0u32); + trie.insert((ip6("2001:db8::"), 32), 32u32); assert_eq!(trie.lookup(ip6("2001:db8::1")), Some(&32)); assert_eq!(trie.lookup(ip6("fe80::1")), Some(&0)); @@ -377,9 +383,9 @@ fn ipv6_default_route() { #[test] fn ipv6_remove_falls_back_to_less_specific() { let mut trie = Poptrie::new(); - trie.insert(ip6("2001:db8::"), 32, 32u32); - trie.insert(ip6("2001:db8:dead::"), 48, 48u32); - trie.remove(ip6("2001:db8:dead::"), 48); + trie.insert((ip6("2001:db8::"), 32), 32u32); + trie.insert((ip6("2001:db8:dead::"), 48), 48u32); + trie.remove((ip6("2001:db8:dead::"), 48)); assert_eq!(trie.lookup(ip6("2001:db8:dead::1")), Some(&32)); } @@ -387,7 +393,7 @@ fn ipv6_remove_falls_back_to_less_specific() { #[test] fn ipv6_host_route() { let mut trie = Poptrie::new(); - trie.insert(ip6("::1"), 128, 128u32); + trie.insert((ip6("::1"), 128), 128u32); assert_eq!(trie.lookup(ip6("::1")), Some(&128)); assert_eq!(trie.lookup(ip6("::2")), None); } diff --git a/tests/bulk_proptests.rs b/tests/bulk_proptests.rs index 9f4a83b..242d99e 100644 --- a/tests/bulk_proptests.rs +++ b/tests/bulk_proptests.rs @@ -15,7 +15,7 @@ proptest! { let mut manually_inserted = Poptrie::new(); for ((key, len), value) in entries.iter().cloned() { - manually_inserted.insert(key, len, value); + manually_inserted.insert((key, len), value); } let bulk_inserted: Poptrie<_,_> = entries.iter().cloned().collect(); diff --git a/tests/lpm_proptests.rs b/tests/lpm_proptests.rs index 18f23b3..36ae05f 100644 --- a/tests/lpm_proptests.rs +++ b/tests/lpm_proptests.rs @@ -49,7 +49,7 @@ proptest! { // Insert all prefixes into both implementations for (prefix, value) in &prefixes { - poptrie.insert(prefix.addr.to_bits(),prefix.prefix_len, *value); + poptrie.insert((prefix.addr.to_bits(), prefix.prefix_len), *value); reference.insert(prefix.addr, prefix.prefix_len, *value); } @@ -57,14 +57,14 @@ proptest! { let delete_prefixes = prefixes.iter().take(prefixes.len() / 10); for (prefix, _) in delete_prefixes { // Assert that the poptrie contains the prefix before deletion - prop_assert!(poptrie.contains_key(prefix.addr.to_bits(), prefix.prefix_len)); + prop_assert!(poptrie.contains_key((prefix.addr.to_bits(), prefix.prefix_len))); // Delete - poptrie.remove(prefix.addr.to_bits(), prefix.prefix_len); + poptrie.remove((prefix.addr.to_bits(), prefix.prefix_len)); reference.remove(prefix.addr, prefix.prefix_len); // Assert that the poptrie no longer contains the prefix after deletion - prop_assert!(!poptrie.contains_key(prefix.addr.to_bits(), prefix.prefix_len)); + prop_assert!(!poptrie.contains_key((prefix.addr.to_bits(), prefix.prefix_len))); } // Compare lookups