diff --git a/src/lib.rs b/src/lib.rs index 0b7197d..aa7ceae 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -144,6 +144,28 @@ where self.inner.exact_match(&ip.nibbles().as_ref(), masklen) } + /// Perform exact match lookup of `ip`/`masklen` and return the + /// value in mutable form. + /// + /// # Examples + /// + /// ``` + /// use treebitmap::IpLookupTable; + /// use std::net::Ipv6Addr; + /// + /// let mut table = IpLookupTable::new(); + /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0); + /// let masklen = 32; + /// table.insert(prefix, masklen, "foo"); + /// + /// assert_eq!(table.exact_match_mut(prefix, masklen), Some(&mut "foo")); + /// // differing mask + /// assert_eq!(table.exact_match_mut(prefix, 48), None); + /// ``` + pub fn exact_match_mut(&mut self, ip: A, masklen: u32) -> Option<&mut T> { + self.inner.exact_match_mut(&ip.nibbles().as_ref(), masklen) + } + /// Perform longest match lookup of `ip` and return the best matching /// prefix, designated by ip, masklen, along with its value. /// @@ -176,6 +198,39 @@ where } } + /// Perform match lookup of `ip` and return the all matching + /// prefixes, designated by ip, masklen, along with its value. + /// + /// # Example + /// + /// ``` + /// use treebitmap::IpLookupTable; + /// use std::net::Ipv6Addr; + /// + /// let mut table = IpLookupTable::new(); + /// let less_specific = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0); + /// let more_specific = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0, 0, 0, 0, 0); + /// table.insert(less_specific, 32, "foo"); + /// table.insert(more_specific, 48, "bar"); + /// + /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, + /// 0xcafe, 0xbabe, 0, 1); + /// let matches = table.matches(lookupip); + /// assert_eq!(matches.len(), 2); + /// + /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xcafe, 0xf00, + /// 0xf00, 0xf00, 0, 1); + /// let matches = table.matches(lookupip); + /// assert_eq!(matches.len(), 1); + /// ``` + pub fn matches(&self, ip: A) -> Vec<(A, u32, &T)> { + self.inner + .matches(&ip.nibbles().as_ref()) + .iter() + .map(|(bits_matched, value)| (ip.mask(*bits_matched), *bits_matched, *value)) + .collect() + } + /// Returns iterator over prefixes and values. /// /// # Examples diff --git a/src/tree_bitmap/allocator.rs b/src/tree_bitmap/allocator.rs index 08e4350..f8c1e17 100644 --- a/src/tree_bitmap/allocator.rs +++ b/src/tree_bitmap/allocator.rs @@ -455,7 +455,7 @@ mod tests { bucket.set_slot_entry(slot, i, 1000 + i); } for i in 0..spacing { - let mut x = bucket.get_slot_entry_mut(slot, i); + let x = bucket.get_slot_entry_mut(slot, i); *x += 1; } for i in 0..spacing { @@ -542,7 +542,7 @@ mod tests { } for i in 0..32 { - let mut x = alloc.get_mut(&hdl, i); + let x = alloc.get_mut(&hdl, i); *x += 1; } diff --git a/src/tree_bitmap/mod.rs b/src/tree_bitmap/mod.rs index 49bbdf6..fc8f2b9 100644 --- a/src/tree_bitmap/mod.rs +++ b/src/tree_bitmap/mod.rs @@ -146,6 +146,45 @@ impl TreeBitmap { } } + /// All matches lookup of ```nibbles```. Returns of Vec of tuples, each containing bits matched as u32 and a reference to T. + pub fn matches(&self, nibbles: &[u8]) -> Vec<(u32, &T)> { + let mut cur_hdl = self.root_handle(); + let mut cur_index = 0; + let mut bits_searched = 0; + let mut matches: Vec<(u32, &T)> = Vec::new(); + + for nibble in nibbles { + let cur_node = *self.trienodes.get(&cur_hdl, cur_index); + let match_mask = node::MATCH_MASKS[*nibble as usize]; + + if let MatchResult::Match(result_hdl, result_index, matching_bit_index) = + cur_node.match_internal(match_mask) + { + let mut bits_matched = bits_searched; + bits_matched += node::BIT_MATCH[matching_bit_index as usize]; + matches.push((bits_matched, self.results.get(&result_hdl, result_index))); + } + + if cur_node.is_endnode() { + break; + } + match cur_node.match_external(match_mask) { + MatchResult::Chase(child_hdl, child_index) => { + bits_searched += 4; + cur_hdl = child_hdl; + cur_index = child_index; + continue; + } + MatchResult::None => { + break; + } + _ => unreachable!(), + } + } + + matches + } + pub fn insert(&mut self, nibbles: &[u8], masklen: u32, value: T) -> Option { let mut cur_hdl = self.root_handle(); let mut cur_index = 0; @@ -282,6 +321,37 @@ impl TreeBitmap { None } + pub fn exact_match_mut(&mut self, nibbles: &[u8], masklen: u32) -> Option<&mut T> { + let mut cur_hdl = self.root_handle(); + let mut cur_index = 0; + let mut bits_left = masklen; + + for nibble in nibbles { + let cur_node = self.trienodes.get(&cur_hdl, cur_index); + let bitmap = node::gen_bitmap(*nibble, cmp::min(bits_left, 4)) & node::END_BIT_MASK; + let reached_final_node = bits_left < 4 || (cur_node.is_endnode() && bits_left == 4); + + if reached_final_node { + match cur_node.match_internal(bitmap) { + MatchResult::Match(result_hdl, result_index, _) => { + return Some(self.results.get_mut(&result_hdl, result_index)); + } + _ => return None, + } + } + + match cur_node.match_external(bitmap) { + MatchResult::Chase(child_hdl, child_index) => { + cur_hdl = child_hdl; + cur_index = child_index; + bits_left -= 4; + } + _ => return None, + } + } + None + } + /// Remove prefix. Returns existing value if the prefix previously existed. pub fn remove(&mut self, nibbles: &[u8], masklen: u32) -> Option { debug_assert!(nibbles.len() >= (masklen / 4) as usize); @@ -564,6 +634,13 @@ mod tests { tbm.insert(nibbles_a, mask_a, "foo"); tbm.insert(nibbles_b, mask_b, "bar"); { + { + let matches = tbm.matches(nibbles_b); + assert_eq!( + true, + matches.contains(&(mask_a, &"foo")) && matches.contains(&(mask_b, &"bar")) + ); + } let value = tbm.remove(nibbles_b, mask_b); assert_eq!(value, Some("bar")); let lookup_result = tbm.longest_match(nibbles_b); diff --git a/tests/tests.rs b/tests/tests.rs index 3ca3986..c943f60 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -68,6 +68,20 @@ fn longest_match6() { assert_eq!(ret, None); } +#[test] +fn matches6() { + let mut tbm = IpLookupTable::new(); + let ip = Ipv6Addr::from_str("2a00::0").unwrap(); + tbm.insert(ip, 32, 1); + tbm.insert(ip, 24, 1); + tbm.insert(ip, 16, 2); + assert_eq!( + 2, + tbm.matches(Ipv6Addr::from_str("2a00:0099::0").unwrap()) + .len() + ); +} + #[test] fn longest_match() { let mut tbm = IpLookupTable::new(); @@ -92,6 +106,14 @@ fn longest_match() { assert_eq!(result, None); } +#[test] +fn matches() { + let mut tbm = IpLookupTable::new(); + tbm.insert(Ipv4Addr::new(10, 0, 0, 0), 8, 1); + tbm.insert(Ipv4Addr::new(10, 1, 0, 0), 16, 2); + assert_eq!(2, tbm.matches(Ipv4Addr::new(10, 1, 0, 30)).len()); +} + #[test] fn iter() { let mut tbl = IpLookupTable::new();