prune ill-conceived BTreeMap iter_mut assertion and test more

3 files changed, 152 insertions, 25 deletions

diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs
index 53c2c29a9b6..0a2849f6e39 100644
--- a/src/liballoc/collections/btree/node.rs
+++ b/src/liballoc/collections/btree/node.rs
@@ -493,7 +493,7 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
 
     /// Returns a raw ptr to avoid asserting exclusive access to the entire node.
     fn as_leaf_mut(&mut self) -> *mut LeafNode<K, V> {
-        // We are mutable, so we cannot be the root, so accessing this as a leaf is okay.
+        // We are mutable, so we cannot be the shared root, so accessing this as a leaf is okay.
         self.node.as_ptr()
     }
 
@@ -514,33 +514,37 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
         // but we want to avoid that run-time check.
         // Instead, we create a slice pointing into the node whenever possible.
         // We can sometimes do this even for the shared root, as the slice will be
-        // empty.  We cannot *always* do this because if the type is too highly
-        // aligned, the offset of `keys` in a "full node" might be outside the bounds
-        // of the header!  So we do an alignment check first, that will be
-        // evaluated at compile-time, and only do any run-time check in the rare case
-        // that the alignment is very big.
-        if mem::align_of::<K>() > mem::align_of::<LeafNode<(), ()>>() && self.is_shared_root() {
+        // empty and `NodeHeader` contains an empty `keys_start` array.
+        // We cannot *always* do this because:
+        // - `keys_start` is not correctly typed because we want `NodeHeader`'s size to
+        //   not depend on the alignment of `K` (needed because `as_header` should be safe).
+        //   For this reason, `NodeHeader` has this `K2` parameter (that's usually `()`
+        //   and hence just adds a size-0-align-1 field, not affecting layout).
+        //   If the correctly typed header is more highly aligned than the allocated header,
+        //   we cannot transmute safely.
+        // - Even if we can transmute, the offset of a correctly typed `keys_start` might
+        //   be different and outside the bounds of the allocated header!
+        // So we do an alignment check and a size check first, that will be evaluated
+        // at compile-time, and only do any run-time check in the rare case that
+        // the compile-time checks signal danger.
+        if (mem::align_of::<NodeHeader<K, V, K>>() > mem::align_of::<NodeHeader<K, V>>()
+            || mem::size_of::<NodeHeader<K, V, K>>() != mem::size_of::<NodeHeader<K, V>>())
+            && self.is_shared_root()
+        {
             &[]
         } else {
-            // Thanks to the alignment check above, we know that `keys` will be
+            // If we are a `LeafNode<K, V>`, we can always transmute to
+            // `NodeHeader<K, V, K>` and `keys_start` always has the same offset
+            // as the actual `keys`.
+            // Thanks to the checks above, we know that we can transmute to
+            // `NodeHeader<K, V, K>` and that `keys_start` will be
             // in-bounds of some allocation even if this is the shared root!
             // (We might be one-past-the-end, but that is allowed by LLVM.)
-            // Getting the pointer is tricky though.  `NodeHeader` does not have a `keys`
-            // field because we want its size to not depend on the alignment of `K`
-            // (needed because `as_header` should be safe).  We cannot call `as_leaf`
-            // because we might be the shared root.
-            // For this reason, `NodeHeader` has this `K2` parameter (that's usually `()`
-            // and hence just adds a size-0-align-1 field, not affecting layout).
-            // We know that we can transmute `NodeHeader<K, V, ()>` to `NodeHeader<K, V, K>`
-            // because we did the alignment check above, and hence `NodeHeader<K, V, K>`
-            // is not bigger than `NodeHeader<K, V, ()>`!  Then we can use `NodeHeader<K, V, K>`
+            // Thus we can use `NodeHeader<K, V, K>`
             // to compute the pointer where the keys start.
             // This entire hack will become unnecessary once
             // <https://github.com/rust-lang/rfcs/pull/2582> lands, then we can just take a raw
             // pointer to the `keys` field of `*const InternalNode<K, V>`.
-
-            // This is a non-debug-assert because it can be completely compile-time evaluated.
-            assert!(mem::size_of::<NodeHeader<K, V>>() == mem::size_of::<NodeHeader<K, V, K>>());
             let header = self.as_header() as *const _ as *const NodeHeader<K, V, K>;
             let keys = unsafe { &(*header).keys_start as *const _ as *const K };
             unsafe { slice::from_raw_parts(keys, self.len()) }
@@ -549,7 +553,7 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
 
     fn into_val_slice(self) -> &'a [V] {
         debug_assert!(!self.is_shared_root());
-        // We cannot be the root, so `as_leaf` is okay
+        // We cannot be the shared root, so `as_leaf` is okay
         unsafe { slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().vals), self.len()) }
     }
 
@@ -567,9 +571,11 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
     }
 
     fn into_key_slice_mut(mut self) -> &'a mut [K] {
-        // Same as for `into_key_slice` above, we try to avoid a run-time check
-        // (the alignment comparison will usually be performed at compile-time).
-        if mem::align_of::<K>() > mem::align_of::<LeafNode<(), ()>>() && self.is_shared_root() {
+        // Same as for `into_key_slice` above, we try to avoid a run-time check.
+        if (mem::align_of::<NodeHeader<K, V, K>>() > mem::align_of::<NodeHeader<K, V>>()
+            || mem::size_of::<NodeHeader<K, V, K>>() != mem::size_of::<NodeHeader<K, V>>())
+            && self.is_shared_root()
+        {
             &mut []
         } else {
             unsafe {
diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs
index 3177f19927e..35ce1354f52 100644
--- a/src/liballoc/tests/btree/map.rs
+++ b/src/liballoc/tests/btree/map.rs
@@ -1,5 +1,7 @@
 use std::collections::btree_map::Entry::{Occupied, Vacant};
 use std::collections::BTreeMap;
+use std::convert::TryFrom;
+use std::fmt::Debug;
 use std::iter::FromIterator;
 use std::ops::Bound::{self, Excluded, Included, Unbounded};
 use std::rc::Rc;
@@ -57,36 +59,65 @@ fn test_basic_large() {
 #[test]
 fn test_basic_small() {
     let mut map = BTreeMap::new();
+    // Empty, shared root:
     assert_eq!(map.remove(&1), None);
     assert_eq!(map.len(), 0);
+    assert_eq!(map.get(&1), None);
+    assert_eq!(map.get_mut(&1), None);
     assert_eq!(map.first_key_value(), None);
     assert_eq!(map.last_key_value(), None);
-    assert_eq!(map.get(&1), None);
+    assert_eq!(map.keys().count(), 0);
+    assert_eq!(map.values().count(), 0);
     assert_eq!(map.insert(1, 1), None);
+
+    // 1 key-value pair:
     assert_eq!(map.len(), 1);
     assert_eq!(map.get(&1), Some(&1));
+    assert_eq!(map.get_mut(&1), Some(&mut 1));
     assert_eq!(map.first_key_value(), Some((&1, &1)));
     assert_eq!(map.last_key_value(), Some((&1, &1)));
+    assert_eq!(map.keys().collect::<Vec<_>>(), vec![&1]);
+    assert_eq!(map.values().collect::<Vec<_>>(), vec![&1]);
     assert_eq!(map.insert(1, 2), Some(1));
     assert_eq!(map.len(), 1);
     assert_eq!(map.get(&1), Some(&2));
+    assert_eq!(map.get_mut(&1), Some(&mut 2));
     assert_eq!(map.first_key_value(), Some((&1, &2)));
     assert_eq!(map.last_key_value(), Some((&1, &2)));
+    assert_eq!(map.keys().collect::<Vec<_>>(), vec![&1]);
+    assert_eq!(map.values().collect::<Vec<_>>(), vec![&2]);
     assert_eq!(map.insert(2, 4), None);
+
+    // 2 key-value pairs:
     assert_eq!(map.len(), 2);
     assert_eq!(map.get(&2), Some(&4));
+    assert_eq!(map.get_mut(&2), Some(&mut 4));
     assert_eq!(map.first_key_value(), Some((&1, &2)));
     assert_eq!(map.last_key_value(), Some((&2, &4)));
+    assert_eq!(map.keys().collect::<Vec<_>>(), vec![&1, &2]);
+    assert_eq!(map.values().collect::<Vec<_>>(), vec![&2, &4]);
     assert_eq!(map.remove(&1), Some(2));
+
+    // 1 key-value pair:
     assert_eq!(map.len(), 1);
     assert_eq!(map.get(&1), None);
+    assert_eq!(map.get_mut(&1), None);
     assert_eq!(map.get(&2), Some(&4));
+    assert_eq!(map.get_mut(&2), Some(&mut 4));
     assert_eq!(map.first_key_value(), Some((&2, &4)));
     assert_eq!(map.last_key_value(), Some((&2, &4)));
+    assert_eq!(map.keys().collect::<Vec<_>>(), vec![&2]);
+    assert_eq!(map.values().collect::<Vec<_>>(), vec![&4]);
     assert_eq!(map.remove(&2), Some(4));
+
+    // Empty but private root:
     assert_eq!(map.len(), 0);
+    assert_eq!(map.get(&1), None);
+    assert_eq!(map.get_mut(&1), None);
     assert_eq!(map.first_key_value(), None);
     assert_eq!(map.last_key_value(), None);
+    assert_eq!(map.keys().count(), 0);
+    assert_eq!(map.values().count(), 0);
     assert_eq!(map.remove(&1), None);
 }
 
@@ -142,6 +173,87 @@ fn test_iter_rev() {
     test(size, map.into_iter().rev());
 }
 
+/// Specifically tests iter_mut's ability to mutate the value of pairs in-line
+fn do_test_iter_mut_mutation<T>(size: usize)
+where
+    T: Copy + Debug + Ord + TryFrom<usize>,
+    <T as std::convert::TryFrom<usize>>::Error: std::fmt::Debug,
+{
+    let zero = T::try_from(0).unwrap();
+    let mut map: BTreeMap<T, T> = (0..size).map(|i| (T::try_from(i).unwrap(), zero)).collect();
+
+    // Forward and backward iteration sees enough pairs (also tested elsewhere)
+    assert_eq!(map.iter_mut().count(), size);
+    assert_eq!(map.iter_mut().rev().count(), size);
+
+    // Iterate forwards, trying to mutate to unique values
+    for (i, (k, v)) in map.iter_mut().enumerate() {
+        assert_eq!(*k, T::try_from(i).unwrap());
+        assert_eq!(*v, zero);
+        *v = T::try_from(i + 1).unwrap();
+    }
+
+    // Iterate backwards, checking that mutations succeeded and trying to mutate again
+    for (i, (k, v)) in map.iter_mut().rev().enumerate() {
+        assert_eq!(*k, T::try_from(size - i - 1).unwrap());
+        assert_eq!(*v, T::try_from(size - i).unwrap());
+        *v = T::try_from(2 * size - i).unwrap();
+    }
+
+    // Check that backward mutations succeeded
+    for (i, (k, v)) in map.iter_mut().enumerate() {
+        assert_eq!(*k, T::try_from(i).unwrap());
+        assert_eq!(*v, T::try_from(size + i + 1).unwrap());
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord)]
+#[repr(align(32))]
+struct Align32(usize);
+
+impl TryFrom<usize> for Align32 {
+    type Error = ();
+
+    fn try_from(s: usize) -> Result<Align32, ()> {
+        Ok(Align32(s))
+    }
+}
+
+#[test]
+fn test_iter_mut_mutation() {
+    // Check many alignments because various fields precede array in NodeHeader.
+    // Check with size 0 which should not iterate at all.
+    // Check with size 1 for a tree with one kind of node (root = leaf).
+    // Check with size 12 for a tree with two kinds of nodes (root and leaves).
+    // Check with size 144 for a tree with all kinds of nodes (root, internals and leaves).
+    do_test_iter_mut_mutation::<u8>(0);
+    do_test_iter_mut_mutation::<u8>(1);
+    do_test_iter_mut_mutation::<u8>(12);
+    do_test_iter_mut_mutation::<u8>(127); // not enough unique values to test 144
+    do_test_iter_mut_mutation::<u16>(1);
+    do_test_iter_mut_mutation::<u16>(12);
+    do_test_iter_mut_mutation::<u16>(144);
+    do_test_iter_mut_mutation::<u32>(1);
+    do_test_iter_mut_mutation::<u32>(12);
+    do_test_iter_mut_mutation::<u32>(144);
+    do_test_iter_mut_mutation::<u64>(1);
+    do_test_iter_mut_mutation::<u64>(12);
+    do_test_iter_mut_mutation::<u64>(144);
+    do_test_iter_mut_mutation::<u128>(1);
+    do_test_iter_mut_mutation::<u128>(12);
+    do_test_iter_mut_mutation::<u128>(144);
+    do_test_iter_mut_mutation::<Align32>(1);
+    do_test_iter_mut_mutation::<Align32>(12);
+    do_test_iter_mut_mutation::<Align32>(144);
+}
+
+#[test]
+fn test_into_key_slice_with_shared_root_past_bounds() {
+    let mut map: BTreeMap<Align32, ()> = BTreeMap::new();
+    assert_eq!(map.get(&Align32(1)), None);
+    assert_eq!(map.get_mut(&Align32(1)), None);
+}
+
 #[test]
 fn test_values_mut() {
     let mut a = BTreeMap::new();
diff --git a/src/liballoc/tests/btree/set.rs b/src/liballoc/tests/btree/set.rs
index ed29ed62b1b..265ef758cc5 100644
--- a/src/liballoc/tests/btree/set.rs
+++ b/src/liballoc/tests/btree/set.rs
@@ -303,6 +303,15 @@ fn test_is_subset() {
 }
 
 #[test]
+fn test_clear() {
+    let mut x = BTreeSet::new();
+    x.insert(1);
+
+    x.clear();
+    assert!(x.is_empty());
+}
+
+#[test]
 fn test_zip() {
     let mut x = BTreeSet::new();
     x.insert(5);