fix assertion error during remove

src/trie/patricia.rs

@@ -63,6 +63,8 @@ impl<K: IpPrefix, V> RadixTrie<K, V> {
 
         let (deepestbranching, deepestleaf) =
             self.branching.search_deepest_candidate(&addedpfx.bitslot());
+        let deepeststructuralleaf =
+            self.structural_leaf_for_insert(deepestbranching, deepestleaf);
         let mut l = deepestleaf;
         let mut b = deepestbranching;
         if l != self[b].escape && !self[l].covers(&addedpfx) {
@@ -82,9 +84,9 @@ impl<K: IpPrefix, V> RadixTrie<K, V> {
                         &addedpfx.bitslot(),
                         addedpfx.len(),
                         deepestbranching,
-                        deepestleaf,
-                        &self[deepestleaf].bitslot(),
-                        self[deepestleaf].len(),
+                        deepeststructuralleaf,
+                        &self[deepeststructuralleaf].bitslot(),
+                        self[deepeststructuralleaf].len(),
                     );
                     return None;
                 }
@@ -103,6 +105,8 @@ impl<K: IpPrefix, V> RadixTrie<K, V> {
 
         let (deepestbranching, deepestleaf) =
             self.branching.search_deepest_candidate(&addedpfx.bitslot());
+        let deepeststructuralleaf =
+            self.structural_leaf_for_insert(deepestbranching, deepestleaf);
         let mut l = deepestleaf;
         let mut b = deepestbranching;
         if l != self[b].escape && !self[l].covers(&addedpfx) {
@@ -122,9 +126,9 @@ impl<K: IpPrefix, V> RadixTrie<K, V> {
                         &addedpfx.bitslot(),
                         addedpfx.len(),
                         deepestbranching,
-                        deepestleaf,
-                        &self[deepestleaf].bitslot(),
-                        self[deepestleaf].len(),
+                        deepeststructuralleaf,
+                        &self[deepeststructuralleaf].bitslot(),
+                        self[deepeststructuralleaf].len(),
                     );
                     return None;
                 }
@@ -139,6 +143,19 @@ impl<K: IpPrefix, V> RadixTrie<K, V> {
 }
 
 impl<K: IpPrefix, V> RadixTrie<K, V> {
+    #[inline]
+    fn structural_leaf_for_insert(
+        &self,
+        deepestbranching: BranchingIndex,
+        deepestleaf: LeafIndex,
+    ) -> LeafIndex {
+        if deepestleaf == self[deepestbranching].escape {
+            self.branching.search_one_matching_leaf(deepestbranching)
+        } else {
+            deepestleaf
+        }
+    }
+
     pub fn get<Q>(&self, k: &Q) -> Option<(&K, &V)>
     where
         Q: IpPrefix<Addr = K::Addr>,
@@ -187,19 +204,28 @@ impl<K: IpPrefix, V> RadixTrie<K, V> {
 
             // todo: some branching possibly becomes useless and should be removed here
 
-            // reindex the leaf which will be swapped with the removed one
+            // Reindex the leaf which will be swapped with the removed one.
+            //
+            // If the removed leaf is already the last leaf, `swap_remove` will not move
+            // anything into its slot. In that case there is no surviving leaf to
+            // reindex, and rewriting `lastleaf` to `l` would leave stale references to
+            // the leaf that is about to disappear.
             let lastleaf = LeafIndex::from(self.leaves.len() - 1);
-            let (mut bb, _) = self.inner_lookup(&self[lastleaf]);
-            //debug_assert_eq!( dbg!(self[lastleaf]).len(), dbg!(self[_ll]).len() );
-            if self[bb].child[0] == lastleaf {
-                self[bb].child[0] = l.into();
-            }
-            if self[bb].child[1] == lastleaf {
-                self[bb].child[1] = l.into();
-            }
-            while self[bb].escape == lastleaf {
-                self[bb].escape = l;
-                bb = self[bb].parent; // climb up the escape chain
+            if l != lastleaf {
+                let (mut bb, _) = self.inner_lookup(&self[lastleaf]);
+                //debug_assert_eq!( dbg!(self[lastleaf]).len(), dbg!(self[_ll]).len() );
+                if self[bb].child[0] == lastleaf {
+                    self[bb].child[0] = l.into();
+                }
+                if self[bb].child[1] == lastleaf {
+                    self[bb].child[1] = l.into();
+                }
+                if self[bb].escape == lastleaf {
+                    while self[self[bb].parent].escape == lastleaf {
+                        bb = self[bb].parent;
+                    }
+                    self.branching.replace_escape_leaf(bb, lastleaf, l);
+                }
             }
             // effective removal of the leaf
             Some(self.leaves.0.swap_remove(l.index()).1)
@@ -436,7 +462,6 @@ impl BranchingTree {
         }
     }
 
-    #[allow(dead_code)]
     pub fn search_one_matching_leaf(&self, mut b: BranchingIndex) -> LeafIndex {
         loop {
             let bb = &self[b];

src/trie/tests.rs

@@ -37,3 +37,65 @@ fn ipv6_tries() {
             assert!(p2.covers_equally(p3));
         });
 }
+
+#[test]
+fn remove_reindexes_moved_escape_leaf() {
+    let mut trie = Ipv4RTrieMap::new();
+    let leaf_to_keep = "80.0.0.0/4".parse::<Ipv4Prefix>().unwrap();
+    let unrelated_leaf = "128.0.0.0/3".parse::<Ipv4Prefix>().unwrap();
+    let leaf_to_remove = "96.0.0.0/4".parse::<Ipv4Prefix>().unwrap();
+    let moved_escape_leaf = "0.0.0.0/1".parse::<Ipv4Prefix>().unwrap();
+
+    // The insertion order is important: removing `leaf_to_remove` moves the final
+    // leaf into its slot. That final leaf is also inherited as an escape leaf by a
+    // descendant branch, so all propagated escape references must be reindexed.
+    trie.insert(leaf_to_keep, 1);
+    trie.insert(unrelated_leaf, 2);
+    trie.insert(leaf_to_remove, 3);
+    trie.insert(moved_escape_leaf, 4);
+
+    assert_eq!(trie.remove(&leaf_to_remove), Some(3));
+    assert_eq!(trie.get(&leaf_to_remove), None);
+    assert_eq!(trie.get(&leaf_to_keep), Some(&1));
+    assert_eq!(trie.get(&unrelated_leaf), Some(&2));
+    assert_eq!(trie.get(&moved_escape_leaf), Some(&4));
+}
+
+#[test]
+fn insert_after_removed_specific_leaf_keeps_branch_ordering_valid() {
+    let mut trie = Ipv4RTrieMap::new();
+    let neighboring_leaf = "205.66.33.0/24".parse::<Ipv4Prefix>().unwrap();
+    let leaf_to_remove = "205.66.32.0/24".parse::<Ipv4Prefix>().unwrap();
+    let covering_prefix = "205.66.32.0/22".parse::<Ipv4Prefix>().unwrap();
+
+    // This sequence was minimized from the prefix-trie RIS mutation benchmark.
+    // Removing the /24 leaves behind branching state that must still support
+    // inserting a covering prefix and then the same /24 again.
+    trie.insert(neighboring_leaf, 1);
+    trie.insert(leaf_to_remove, 2);
+    assert_eq!(trie.remove(&leaf_to_remove), Some(2));
+    trie.insert(covering_prefix, 3);
+    trie.insert(leaf_to_remove, 4);
+
+    assert_eq!(trie.get(&neighboring_leaf), Some(&1));
+    assert_eq!(trie.get(&covering_prefix), Some(&3));
+    assert_eq!(trie.get(&leaf_to_remove), Some(&4));
+}
+
+#[test]
+fn replace_after_removed_specific_leaf_keeps_branch_ordering_valid() {
+    let mut trie = Ipv4RTrieSet::new();
+    let neighboring_leaf = "205.66.33.0/24".parse::<Ipv4Prefix>().unwrap();
+    let leaf_to_remove = "205.66.32.0/24".parse::<Ipv4Prefix>().unwrap();
+    let covering_prefix = "205.66.32.0/22".parse::<Ipv4Prefix>().unwrap();
+
+    trie.insert(neighboring_leaf);
+    trie.insert(leaf_to_remove);
+    assert!(trie.remove(&leaf_to_remove));
+    assert_eq!(trie.replace(covering_prefix), None);
+    assert_eq!(trie.replace(leaf_to_remove), None);
+
+    assert!(trie.contains(&neighboring_leaf));
+    assert!(trie.contains(&covering_prefix));
+    assert!(trie.contains(&leaf_to_remove));
+}