Handle mask=0 case in match and remove, fixing hroi#13

jiegec · Jun 29, 2019 · 6e4e2f1 · 6e4e2f1
1 parent 9540ae5
commit 6e4e2f1
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 4 deletions.
diff --git a/src/tree_bitmap/mod.rs b/src/tree_bitmap/mod.rs
@@ -109,7 +109,16 @@ impl<T: Sized> TreeBitmap<T> {
         let mut bits_searched = 0;
         let mut best_match: Option<(AllocatorHandle, u32)> = None; // result handle + index
 
-        for nibble in nibbles {
+        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);
             let match_mask = node::MATCH_MASKS[*nibble as usize];
 
@@ -256,7 +265,16 @@ impl<T: Sized> TreeBitmap<T> {
         let mut cur_index = 0;
         let mut bits_left = masklen;
 
-        for nibble in nibbles {
+        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);
             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);
@@ -279,7 +297,6 @@ impl<T: Sized> TreeBitmap<T> {
                 _ => return None,
             }
         }
-        None
     }
 
     /// Remove prefix. Returns existing value if the prefix previously existed.
@@ -294,7 +311,7 @@ impl<T: Sized> TreeBitmap<T> {
 
     // remove child and result from node
     fn remove_child(&mut self, node: &mut Node, nibbles: &[u8], masklen: u32) -> Option<T> {
-        let nibble = nibbles[0];
+        let nibble = if masklen > 0 { nibbles[0] } else { 0 };
         let bitmap = node::gen_bitmap(nibble, cmp::min(masklen, 4)) & node::END_BIT_MASK;
         let reached_final_node = masklen < 4 || (node.is_endnode() && masklen == 4);
 

diff --git a/tests/tests.rs b/tests/tests.rs
@@ -186,3 +186,26 @@ fn issue_7() {
 
     println!("len: {}", table.len());
 }
+
+// https://github.com/hroi/treebitmap/issues/13
+#[test]
+fn issue_13() {
+    const ADDR: Ipv4Addr = Ipv4Addr::new(49, 255, 11, 17);
+    let mut table = IpLookupTable::new();
+
+    println!("insert 28");
+    table.insert(Ipv4Addr::new(49, 255, 11, 16), 28, ());
+    assert_eq!(
+        table.exact_match(Ipv4Addr::new(49, 255, 11, 16), 28),
+        Some(&())
+    );
+    println!("insert 32");
+    table.insert(ADDR, 32, ());
+
+    println!("match 32");
+    assert_eq!(table.exact_match(ADDR, 32), Some(&()));
+    assert!(table.longest_match(ADDR).is_some());
+
+    let v = table.remove(ADDR, 32);
+    println!("removed: {:?}", v);
+}