Skip to content

Commit

Permalink
Fix BTreeSet's range API panic message, document
Browse files Browse the repository at this point in the history
  • Loading branch information
tnballo committed Jun 23, 2022
1 parent 4334739 commit 774e814
Show file tree
Hide file tree
Showing 7 changed files with 117 additions and 13 deletions.
5 changes: 3 additions & 2 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ use super::dedup_sorted_iter::DedupSortedIter;
use super::navigate::{LazyLeafRange, LeafRange};
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
use super::search::SearchResult::*;
use super::set_val::SetValZST;

mod entry;

Expand Down Expand Up @@ -271,7 +272,7 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
}
}

impl<K, Q: ?Sized, A: Allocator + Clone> super::Recover<Q> for BTreeMap<K, (), A>
impl<K, Q: ?Sized, A: Allocator + Clone> super::Recover<Q> for BTreeMap<K, SetValZST, A>
where
K: Borrow<Q> + Ord,
Q: Ord,
Expand Down Expand Up @@ -318,7 +319,7 @@ where
alloc: (*map.alloc).clone(),
_marker: PhantomData,
}
.insert(());
.insert(SetValZST::default());
None
}
}
Expand Down
33 changes: 33 additions & 0 deletions library/alloc/src/collections/btree/map/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -897,6 +897,39 @@ fn test_range_mut() {
map.check();
}

#[should_panic(expected = "range start is greater than range end in BTreeMap")]
#[test]
fn test_range_panic_1() {
let mut map = BTreeMap::new();
map.insert(3, "a");
map.insert(5, "b");
map.insert(8, "c");

let _invalid_range = map.range((Included(&8), Included(&3)));
}

#[should_panic(expected = "range start and end are equal and excluded in BTreeMap")]
#[test]
fn test_range_panic_2() {
let mut map = BTreeMap::new();
map.insert(3, "a");
map.insert(5, "b");
map.insert(8, "c");

let _invalid_range = map.range((Excluded(&5), Excluded(&5)));
}

#[should_panic(expected = "range start and end are equal and excluded in BTreeMap")]
#[test]
fn test_range_panic_3() {
let mut map: BTreeMap<i32, ()> = BTreeMap::new();
map.insert(3, ());
map.insert(5, ());
map.insert(8, ());

let _invalid_range = map.range((Excluded(&5), Excluded(&5)));
}

#[test]
fn test_retain() {
let mut map = BTreeMap::from_iter((0..100).map(|x| (x, x * 10)));
Expand Down
1 change: 1 addition & 0 deletions library/alloc/src/collections/btree/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ mod node;
mod remove;
mod search;
pub mod set;
mod set_val;
mod split;

#[doc(hidden)]
Expand Down
15 changes: 13 additions & 2 deletions library/alloc/src/collections/btree/search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,17 +97,28 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
K: Borrow<Q>,
R: RangeBounds<Q>,
{
// Determine if map or set is being searched
let is_set = <V as super::set_val::IsSetVal>::is_set_val();

// Inlining these variables should be avoided. We assume the bounds reported by `range`
// remain the same, but an adversarial implementation could change between calls (#81138).
let (start, end) = (range.start_bound(), range.end_bound());
match (start, end) {
(Bound::Excluded(s), Bound::Excluded(e)) if s == e => {
panic!("range start and end are equal and excluded in BTreeMap")
if is_set {
panic!("range start and end are equal and excluded in BTreeSet")
} else {
panic!("range start and end are equal and excluded in BTreeMap")
}
}
(Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e))
if s > e =>
{
panic!("range start is greater than range end in BTreeMap")
if is_set {
panic!("range start is greater than range end in BTreeSet")
} else {
panic!("range start is greater than range end in BTreeMap")
}
}
_ => {}
}
Expand Down
24 changes: 15 additions & 9 deletions library/alloc/src/collections/btree/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use core::ops::{BitAnd, BitOr, BitXor, RangeBounds, Sub};

use super::map::{BTreeMap, Keys};
use super::merge_iter::MergeIterInner;
use super::set_val::SetValZST;
use super::Recover;

use crate::alloc::{Allocator, Global};
Expand Down Expand Up @@ -81,7 +82,7 @@ pub struct BTreeSet<
T,
#[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global,
> {
map: BTreeMap<T, (), A>,
map: BTreeMap<T, SetValZST, A>,
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down Expand Up @@ -135,7 +136,7 @@ impl<T: Clone, A: Allocator + Clone> Clone for BTreeSet<T, A> {
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Iter<'a, T: 'a> {
iter: Keys<'a, T, ()>,
iter: Keys<'a, T, SetValZST>,
}

#[stable(feature = "collection_debug", since = "1.17.0")]
Expand All @@ -158,7 +159,7 @@ pub struct IntoIter<
T,
#[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global,
> {
iter: super::map::IntoIter<T, (), A>,
iter: super::map::IntoIter<T, SetValZST, A>,
}

/// An iterator over a sub-range of items in a `BTreeSet`.
Expand All @@ -171,7 +172,7 @@ pub struct IntoIter<
#[derive(Debug)]
#[stable(feature = "btree_range", since = "1.17.0")]
pub struct Range<'a, T: 'a> {
iter: super::map::Range<'a, T, ()>,
iter: super::map::Range<'a, T, SetValZST>,
}

/// A lazy iterator producing elements in the difference of `BTreeSet`s.
Expand Down Expand Up @@ -375,6 +376,11 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
/// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
/// range from 4 to 10.
///
/// # Panics
///
/// Panics if range `start > end`.
/// Panics if range `start == end` and both bounds are `Excluded`.
///
/// # Examples
///
/// ```
Expand Down Expand Up @@ -905,7 +911,7 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
where
T: Ord,
{
self.map.insert(value, ()).is_none()
self.map.insert(value, SetValZST::default()).is_none()
}

/// Adds a value to the set, replacing the existing element, if any, that is
Expand Down Expand Up @@ -1210,7 +1216,7 @@ impl<T: Ord> FromIterator<T> for BTreeSet<T> {

impl<T: Ord, A: Allocator + Clone> BTreeSet<T, A> {
fn from_sorted_iter<I: Iterator<Item = T>>(iter: I, alloc: A) -> BTreeSet<T, A> {
let iter = iter.map(|k| (k, ()));
let iter = iter.map(|k| (k, SetValZST::default()));
let map = BTreeMap::bulk_build_from_sorted_iter(iter, alloc);
BTreeSet { map }
}
Expand All @@ -1234,7 +1240,7 @@ impl<T: Ord, const N: usize> From<[T; N]> for BTreeSet<T> {

// use stable sort to preserve the insertion order.
arr.sort();
let iter = IntoIterator::into_iter(arr).map(|k| (k, ()));
let iter = IntoIterator::into_iter(arr).map(|k| (k, SetValZST::default()));
let map = BTreeMap::bulk_build_from_sorted_iter(iter, Global);
BTreeSet { map }
}
Expand Down Expand Up @@ -1284,7 +1290,7 @@ pub struct DrainFilter<
F: 'a + FnMut(&T) -> bool,
{
pred: F,
inner: super::map::DrainFilterInner<'a, T, ()>,
inner: super::map::DrainFilterInner<'a, T, SetValZST>,
/// The BTreeMap will outlive this IntoIter so we don't care about drop order for `alloc`.
alloc: A,
}
Expand Down Expand Up @@ -1319,7 +1325,7 @@ where

fn next(&mut self) -> Option<T> {
let pred = &mut self.pred;
let mut mapped_pred = |k: &T, _v: &mut ()| pred(k);
let mut mapped_pred = |k: &T, _v: &mut SetValZST| pred(k);
self.inner.next(&mut mapped_pred, self.alloc.clone()).map(|(k, _)| k)
}

Expand Down
23 changes: 23 additions & 0 deletions library/alloc/src/collections/btree/set/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use crate::vec::Vec;
use std::cmp::Ordering;
use std::hash::{Hash, Hasher};
use std::iter::FromIterator;
use std::ops::Bound::{Excluded, Included};
use std::panic::{catch_unwind, AssertUnwindSafe};

#[test]
Expand Down Expand Up @@ -831,3 +832,25 @@ fn from_array() {
let unordered_duplicates = BTreeSet::from([4, 1, 4, 3, 2]);
assert_eq!(set, unordered_duplicates);
}

#[should_panic(expected = "range start is greater than range end in BTreeSet")]
#[test]
fn test_range_panic_1() {
let mut set = BTreeSet::new();
set.insert(3);
set.insert(5);
set.insert(8);

let _invalid_range = set.range((Included(&8), Included(&3)));
}

#[should_panic(expected = "range start and end are equal and excluded in BTreeSet")]
#[test]
fn test_range_panic_2() {
let mut set = BTreeSet::new();
set.insert(3);
set.insert(5);
set.insert(8);

let _invalid_range = set.range((Excluded(&5), Excluded(&5)));
}
29 changes: 29 additions & 0 deletions library/alloc/src/collections/btree/set_val.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
/// Zero-Sized Type (ZST) for internal `BTreeSet` values.
/// Used instead of `()` to differentiate between:
/// * `BTreeMap<T, ()>` (possible user-defined map)
/// * `BTreeMap<T, SetValZST>` (internal set representation)
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Clone, Default)]
pub struct SetValZST;

/// A trait to differentiate between `BTreeMap` and `BTreeSet` values.
/// Returns `true` only for type `SetValZST`, `false` for all other types (blanket implementation).
/// `TypeId` requires a `'static` lifetime, use of this trait avoids that restriction.
///
/// [`TypeId`]: std::any::TypeId
pub trait IsSetVal {
fn is_set_val() -> bool;
}

// Blanket implementation
impl<V> IsSetVal for V {
default fn is_set_val() -> bool {
false
}
}

// Specialization
impl IsSetVal for SetValZST {
fn is_set_val() -> bool {
true
}
}

0 comments on commit 774e814

Please sign in to comment.