New method IpLookupTable::matches

src/lib.rs

@@ -208,6 +208,45 @@ 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.count(), 2);
+    ///
+    /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xcafe, 0xf00,
+    ///                              0xf00, 0xf00, 0, 1);
+    /// let matches = table.matches(lookupip);
+    /// assert_eq!(matches.count(), 1);
+    /// ```
+    pub fn matches(&self, ip: A) -> impl Iterator<Item = (A, u32, &T)> {
+        self.inner
+            .matches(ip.nibbles().as_ref())
+            .map(move |(bits_matched, value)| (ip.mask(bits_matched), bits_matched, value))
+    }
+
+    /// Perform match lookup of `ip` and return the all matching
+    /// prefixes, designated by ip, masklen, along with its mutable value.
+    pub fn matches_mut(&mut self, ip: A) -> impl Iterator<Item = (A, u32, &mut T)> {
+        self.inner
+            .matches_mut(ip.nibbles().as_ref())
+            .map(move |(bits_matched, value)| (ip.mask(bits_matched), bits_matched, value))
+    }
+
     /// Returns iterator over prefixes and values.
     ///
     /// # Examples

src/tree_bitmap/mod.rs

@@ -121,7 +121,7 @@ impl<T: Sized> TreeBitmap<T> {
             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)
+                cur_node.match_internal(match_mask)
             {
                 let bits_matched = bits_searched + node::BIT_MATCH[matching_bit_index as usize];
                 best_match = Some((result_hdl, result_index, bits_matched));
@@ -159,13 +159,79 @@ impl<T: Sized> TreeBitmap<T> {
     /// longest match lookup of ```nibbles```. Returns bits matched as u32, and mutable reference to T.
     pub fn longest_match_mut(&mut self, nibbles: &[u8]) -> Option<(u32, &mut T)> {
         match self.longest_match_internal(&nibbles) {
-            Some((result_hdl, result_index, bits_matched)) => {
-                Some((bits_matched, self.results.get_mut(&result_hdl, result_index)))
-            }
+            Some((result_hdl, result_index, bits_matched)) => Some((
+                bits_matched,
+                self.results.get_mut(&result_hdl, result_index),
+            )),
             None => None,
         }
     }
 
+    /// All matches lookup of ```nibbles```. Returns of Vec of tuples, each containing bits matched as u32 and a reference to T.
+    fn matches_internal(&self, nibbles: &[u8]) -> Vec<(u32, AllocatorHandle, u32)> {
+        let mut cur_hdl = self.root_handle();
+        let mut cur_index = 0;
+        let mut bits_searched = 0;
+        let mut matches = Vec::new();
+
+        let mut loop_count = 0;
+        loop {
+            let nibble = if loop_count < nibbles.len() {
+                nibbles[loop_count]
+            } else {
+                0
+            };
+            loop_count += 1;
+            let nibble = &nibble;
+            let cur_node = *self.trienodes.get(&cur_hdl, cur_index);
+
+            for i in 0..5 {
+                let prefix = *nibble & (!0 << (4 - i));
+                let bitmap = node::gen_bitmap(prefix, i as u32) & node::END_BIT_MASK;
+                if let MatchResult::Match(result_hdl, result_index, _) =
+                    cur_node.match_internal(bitmap)
+                {
+                    let bits_matched = bits_searched + (i as u32);
+                    matches.push((bits_matched, result_hdl, result_index));
+                }
+            }
+
+            if cur_node.is_endnode() {
+                break;
+            }
+
+            let match_mask = node::MATCH_MASKS[*nibble as usize];
+            match cur_node.match_external(match_mask) {
+                MatchResult::Chase(child_hdl, child_index) => {
+                    bits_searched += 4;
+                    cur_hdl = child_hdl;
+                    cur_index = child_index;
+                }
+                MatchResult::None => {
+                    break;
+                }
+                _ => unreachable!(),
+            }
+        }
+
+        matches
+    }
+
+    /// 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]) -> impl Iterator<Item = (u32, &T)> {
+        self.matches_internal(nibbles).into_iter().map(
+            move |(bits_matched, result_hdl, result_index)| {
+                (bits_matched, self.results.get(&result_hdl, result_index))
+            },
+        )
+    }
+
+    /// All matches lookup of ```nibbles```. Returns of Vec of tuples, each containing bits matched as u32 and a reference to T.
+    pub fn matches_mut(&mut self, nibbles: &[u8]) -> MatchesMut<T> {
+        let path = self.matches_internal(nibbles).into_iter();
+        MatchesMut { inner: self, path }
+    }
+
     pub fn insert(&mut self, nibbles: &[u8], masklen: u32, value: T) -> Option<T> {
         let mut cur_hdl = self.root_handle();
         let mut cur_index = 0;
@@ -543,6 +609,26 @@ impl<T> Drop for IntoIter<T> {
     }
 }
 
+pub struct MatchesMut<'a, T: 'a> {
+    inner: &'a mut TreeBitmap<T>,
+    path: std::vec::IntoIter<(u32, AllocatorHandle, u32)>,
+}
+
+impl<'a, T: 'a> Iterator for MatchesMut<'a, T> {
+    type Item = (u32, &'a mut T);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.path.next() {
+            Some((bits_matched, hdl, index)) => unsafe {
+                let ptr: *mut T = self.inner.results.get_mut(&hdl, index);
+                let val_ref = &mut *ptr;
+                Some((bits_matched, val_ref))
+            },
+            None => None,
+        }
+    }
+}
+
 impl<T> Drop for TreeBitmap<T> {
     fn drop(&mut self) {
         if self.should_drop {
@@ -591,6 +677,15 @@ mod tests {
         let (nibbles_b, mask_b) = (&[0, 10, 0, 10, 0, 10, 0, 0], 24);
         tbm.insert(nibbles_a, mask_a, "foo");
         tbm.insert(nibbles_b, mask_b, "bar");
+
+        {
+            let mut matches = tbm.matches(nibbles_b);
+            assert_eq!(
+                true,
+                matches.any(|p| p == (mask_a, &"foo")) && matches.any(|p| p == (mask_b, &"bar"))
+            );
+        }
+
         {
             let value = tbm.remove(nibbles_b, mask_b);
             assert_eq!(value, Some("bar"));

tests/tests.rs

@@ -70,6 +70,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())
+            .count()
+    );
+}
+
 #[test]
 fn longest_match() {
     let mut tbm = IpLookupTable::new();
@@ -88,19 +102,28 @@ fn longest_match() {
     assert_eq!(result, Some((Ipv4Addr::new(100, 64, 0, 0), 10, &100004)));
 
     let result = tbm.longest_match_mut(Ipv4Addr::new(100, 100, 100, 100));
-    assert_eq!(result, Some((Ipv4Addr::new(100, 64, 0, 0), 10, &mut 100004)));
+    assert_eq!(
+        result,
+        Some((Ipv4Addr::new(100, 64, 0, 0), 10, &mut 100004))
+    );
 
     let result = tbm.longest_match(Ipv4Addr::new(100, 64, 0, 100));
     assert_eq!(result, Some((Ipv4Addr::new(100, 64, 0, 0), 24, &10064024)));
 
     let result = tbm.longest_match_mut(Ipv4Addr::new(100, 64, 0, 100));
-    assert_eq!(result, Some((Ipv4Addr::new(100, 64, 0, 0), 24, &mut 10064024)));
+    assert_eq!(
+        result,
+        Some((Ipv4Addr::new(100, 64, 0, 0), 24, &mut 10064024))
+    );
 
     let result = tbm.longest_match(Ipv4Addr::new(100, 64, 1, 100));
     assert_eq!(result, Some((Ipv4Addr::new(100, 64, 1, 0), 24, &10064124)));
 
     let result = tbm.longest_match_mut(Ipv4Addr::new(100, 64, 1, 100));
-    assert_eq!(result, Some((Ipv4Addr::new(100, 64, 1, 0), 24, &mut 10064124)));
+    assert_eq!(
+        result,
+        Some((Ipv4Addr::new(100, 64, 1, 0), 24, &mut 10064124))
+    );
 
     let result = tbm.longest_match(Ipv4Addr::new(200, 200, 200, 200));
     assert_eq!(result, None);
@@ -169,10 +192,70 @@ fn into_iter() {
 
 #[test]
 fn send() {
-    fn check_if_send<T: Send>() { }
+    fn check_if_send<T: Send>() {}
     check_if_send::<IpLookupTable<Ipv4Addr, ()>>();
 }
 
+#[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)).count());
+}
+
+#[test]
+fn matches_all_zeros() {
+    let mut table = IpLookupTable::new();
+    for i in 0..=32 {
+        table.insert(Ipv4Addr::new(0, 0, 0, 0), i, i);
+    }
+
+    for (ip, matched, value) in table.matches(Ipv4Addr::new(0, 0, 0, 0)) {
+        assert_eq!(ip, Ipv4Addr::new(0, 0, 0, 0));
+        assert_eq!(matched, *value);
+    }
+
+    assert_eq!(table.matches(Ipv4Addr::new(0, 0, 0, 0)).count(), 33);
+    assert_eq!(table.matches(Ipv4Addr::new(1, 0, 0, 0)).count(), 8);
+    assert_eq!(table.matches(Ipv4Addr::new(0, 0, 255, 255)).count(), 17);
+    assert_eq!(table.matches(Ipv4Addr::new(255, 255, 255, 255)).count(), 1);
+    assert_eq!(table.matches(Ipv4Addr::new(64, 0, 0, 0)).count(), 2);
+}
+
+#[test]
+fn matches_10_range() {
+    let mut table = IpLookupTable::new();
+    table.insert(Ipv4Addr::new(10, 0, 0, 0), 8, 0);
+    table.insert(Ipv4Addr::new(10, 6, 0, 0), 16, 0);
+    table.insert(Ipv4Addr::new(10, 6, 252, 0), 24, 0);
+
+    assert_eq!(table.matches(Ipv4Addr::new(10, 6, 252, 3)).count(), 3);
+    assert_eq!(table.matches(Ipv4Addr::new(10, 6, 255, 3)).count(), 2);
+    assert_eq!(table.matches(Ipv4Addr::new(11, 6, 255, 3)).count(), 0);
+}
+
+#[test]
+fn matches_empty() {
+    let table: IpLookupTable<Ipv4Addr, ()> = IpLookupTable::new();
+    assert_eq!(table.matches(Ipv4Addr::new(0, 0, 0, 0)).count(), 0);
+}
+
+#[test]
+fn matches_ipv6() {
+    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");
+    assert_eq!(table.matches(less_specific).count(), 1);
+    assert_eq!(table.matches(more_specific).count(), 2);
+    assert_eq!(
+        table.matches(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)).count(),
+        0
+    );
+}
+
 // https://github.com/hroi/treebitmap/issues/7
 #[test]
 fn issue_7() {
@@ -198,19 +281,21 @@ fn issue_13() {
     let mut table = IpLookupTable::new();
 
     println!("insert 28");
-    table.insert(Ipv4Addr::new(49, 255, 11, 16), 28, ());
+    table.insert(Ipv4Addr::new(49, 255, 11, 16), 28, 28);
     assert_eq!(
         table.exact_match(Ipv4Addr::new(49, 255, 11, 16), 28),
-        Some(&())
+        Some(&28)
     );
     println!("insert 32");
-    table.insert(ADDR, 32, ());
+    table.insert(ADDR, 32, 32);
 
     println!("match 32");
-    assert_eq!(table.exact_match(ADDR, 32), Some(&()));
+    assert_eq!(table.exact_match(ADDR, 32), Some(&32));
     assert!(table.longest_match(ADDR).is_some());
     assert!(table.longest_match_mut(ADDR).is_some());
 
+    assert_eq!(table.matches(ADDR).count(), 2);
+
     let v = table.remove(ADDR, 32);
     println!("removed: {:?}", v);
 }