diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs index 53c2c29a9b659..0a2849f6e3961 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, K, V, Type> { /// Returns a raw ptr to avoid asserting exclusive access to the entire node. fn as_leaf_mut(&mut self) -> *mut LeafNode { - // 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, 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::() > mem::align_of::>() && 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::>() > mem::align_of::>() + || mem::size_of::>() != mem::size_of::>()) + && self.is_shared_root() + { &[] } else { - // Thanks to the alignment check above, we know that `keys` will be + // If we are a `LeafNode`, we can always transmute to + // `NodeHeader` 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` 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` to `NodeHeader` - // because we did the alignment check above, and hence `NodeHeader` - // is not bigger than `NodeHeader`! Then we can use `NodeHeader` + // Thus we can use `NodeHeader` // to compute the pointer where the keys start. // This entire hack will become unnecessary once // lands, then we can just take a raw // pointer to the `keys` field of `*const InternalNode`. - - // This is a non-debug-assert because it can be completely compile-time evaluated. - assert!(mem::size_of::>() == mem::size_of::>()); let header = self.as_header() as *const _ as *const NodeHeader; 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, 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, 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::() > mem::align_of::>() && self.is_shared_root() { + // Same as for `into_key_slice` above, we try to avoid a run-time check. + if (mem::align_of::>() > mem::align_of::>() + || mem::size_of::>() != mem::size_of::>()) + && self.is_shared_root() + { &mut [] } else { unsafe { diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs index 3177f19927eda..35ce1354f52e0 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![&1]); + assert_eq!(map.values().collect::>(), 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![&1]); + assert_eq!(map.values().collect::>(), 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![&1, &2]); + assert_eq!(map.values().collect::>(), 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![&2]); + assert_eq!(map.values().collect::>(), 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(size: usize) +where + T: Copy + Debug + Ord + TryFrom, + >::Error: std::fmt::Debug, +{ + let zero = T::try_from(0).unwrap(); + let mut map: BTreeMap = (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 for Align32 { + type Error = (); + + fn try_from(s: usize) -> Result { + 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::(0); + do_test_iter_mut_mutation::(1); + do_test_iter_mut_mutation::(12); + do_test_iter_mut_mutation::(127); // not enough unique values to test 144 + do_test_iter_mut_mutation::(1); + do_test_iter_mut_mutation::(12); + do_test_iter_mut_mutation::(144); + do_test_iter_mut_mutation::(1); + do_test_iter_mut_mutation::(12); + do_test_iter_mut_mutation::(144); + do_test_iter_mut_mutation::(1); + do_test_iter_mut_mutation::(12); + do_test_iter_mut_mutation::(144); + do_test_iter_mut_mutation::(1); + do_test_iter_mut_mutation::(12); + do_test_iter_mut_mutation::(144); + do_test_iter_mut_mutation::(1); + do_test_iter_mut_mutation::(12); + do_test_iter_mut_mutation::(144); +} + +#[test] +fn test_into_key_slice_with_shared_root_past_bounds() { + let mut map: BTreeMap = 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 ed29ed62b1b3e..265ef758cc5bc 100644 --- a/src/liballoc/tests/btree/set.rs +++ b/src/liballoc/tests/btree/set.rs @@ -302,6 +302,15 @@ fn test_is_subset() { assert_eq!(is_subset(&[99, 100], &large), false); } +#[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();