From f327a352b89a571dde76e176471d12fbd8fce0dd Mon Sep 17 00:00:00 2001
From: Stein Somers <git@steinsomers.be>
Date: Wed, 18 Nov 2020 18:19:38 +0100
Subject: [PATCH] BTreeMap: test full nodes a little more

---
 library/alloc/src/collections/btree/append.rs |  2 +-
 .../alloc/src/collections/btree/map/tests.rs  | 40 +++++++++++++++----
 2 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs
index bd99c4ed2f1..9c53c84f5b2 100644
--- a/library/alloc/src/collections/btree/append.rs
+++ b/library/alloc/src/collections/btree/append.rs
@@ -30,7 +30,7 @@ pub fn append_from_sorted_iters<I>(&mut self, left: I, right: I, length: &mut us
     /// 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)>,
     {
diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs
index 97df8ea07d2..565dc175393 100644
--- a/library/alloc/src/collections/btree/map/tests.rs
+++ b/library/alloc/src/collections/btree/map/tests.rs
@@ -111,6 +111,18 @@ fn assert_strictly_ascending(&self)
             }
         }
     }
+
+    // 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 {