BTreeMap: test full nodes a little more

Author: ssomers

library/alloc/src/collections/btree/append.rs

@@ -30,7 +30,7 @@ impl<K, V> Root<K, V> {
     /// Pushes all key-value pairs to the end of the tree, incrementing a
     /// `length` variable along the way. The latter makes it easier for the
     /// caller to avoid a leak when the iterator panicks.
-    fn bulk_push<I>(&mut self, iter: I, length: &mut usize)
+    pub fn bulk_push<I>(&mut self, iter: I, length: &mut usize)
     where
         I: Iterator<Item = (K, V)>,
     {

library/alloc/src/collections/btree/map/tests.rs

@@ -111,6 +111,18 @@ impl<K, V> BTreeMap<K, V> {
             }
         }
     }
+
+    // Transform the tree to minimize wasted space, obtaining fewer nodes that
+    // are mostly filled up to their capacity. The same compact tree could have
+    // been obtained by inserting keys in a shrewd order.
+    fn compact(&mut self)
+    where
+        K: Ord,
+    {
+        let iter = mem::take(self).into_iter();
+        let root = BTreeMap::ensure_is_owned(&mut self.root);
+        root.bulk_push(iter, &mut self.length);
+    }
 }
 
 impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
@@ -1679,17 +1691,29 @@ fn test_first_last_entry() {
 }
 
 #[test]
-fn test_insert_into_full_left() {
-    let mut map: BTreeMap<_, _> = (0..NODE_CAPACITY).map(|i| (i * 2, ())).collect();
-    assert!(map.insert(NODE_CAPACITY, ()).is_none());
-    map.check();
+fn test_insert_into_full_height_0() {
+    let size = NODE_CAPACITY;
+    for pos in 0..=size {
+        let mut map: BTreeMap<_, _> = (0..size).map(|i| (i * 2 + 1, ())).collect();
+        assert!(map.insert(pos * 2, ()).is_none());
+        map.check();
+    }
 }
 
 #[test]
-fn test_insert_into_full_right() {
-    let mut map: BTreeMap<_, _> = (0..NODE_CAPACITY).map(|i| (i * 2, ())).collect();
-    assert!(map.insert(NODE_CAPACITY + 2, ()).is_none());
-    map.check();
+fn test_insert_into_full_height_1() {
+    let size = NODE_CAPACITY + 1 + NODE_CAPACITY;
+    for pos in 0..=size {
+        let mut map: BTreeMap<_, _> = (0..size).map(|i| (i * 2 + 1, ())).collect();
+        map.compact();
+        let root_node = map.root.as_ref().unwrap().reborrow();
+        assert_eq!(root_node.len(), 1);
+        assert_eq!(root_node.first_leaf_edge().into_node().len(), NODE_CAPACITY);
+        assert_eq!(root_node.last_leaf_edge().into_node().len(), NODE_CAPACITY);
+
+        assert!(map.insert(pos * 2, ()).is_none());
+        map.check();
+    }
 }
 
 macro_rules! create_append_test {