use std::default::Default;
use std::cmp::Ordering::{self, Less, Equal, Greater};
use std::fmt::{self, Debug};
use std::iter::{self, IntoIterator};
use std::mem::{replace, swap};
use std::ops;
use std::ops::Deref;
use std::ptr;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use compare::{Compare, Natural, natural};
use super::Bound;
#[derive(Clone)]
pub struct TreeMap<K, V, C: Compare<K> = Natural<K>> {
root: Option<Box<TreeNode<K, V>>>,
length: usize,
cmp: C,
}
impl<K: PartialEq + Ord, V: PartialEq> PartialEq for TreeMap<K, V> {
#[inline]
fn eq(&self, other: &TreeMap<K, V>) -> bool {
self.iter().eq(other)
}
}
impl<K: Eq + Ord, V: Eq> Eq for TreeMap<K, V> {}
impl<K: Ord, V: PartialOrd> PartialOrd for TreeMap<K, V> {
#[inline]
fn partial_cmp(&self, other: &TreeMap<K, V>) -> Option<Ordering> {
self.iter().partial_cmp(other)
}
}
impl<K: Ord, V: Ord> Ord for TreeMap<K, V> {
#[inline]
fn cmp(&self, other: &TreeMap<K, V>) -> Ordering {
self.iter().cmp(other)
}
}
impl<K: Debug, V: Debug, C> Debug for TreeMap<K, V, C>
where C: Compare<K>
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
try!(write!(f, "{{"));
for (i, (k, v)) in self.iter().enumerate() {
if i != 0 {
try!(write!(f, ", "));
}
try!(write!(f, "{:?}: {:?}", *k, *v));
}
write!(f, "}}")
}
}
impl<K, V, C> Default for TreeMap<K, V, C>
where C: Compare<K> + Default
{
#[inline]
fn default() -> TreeMap<K, V, C> {
TreeMap::with_comparator(Default::default())
}
}
impl<'a, K, V, C, Q: ?Sized> ops::Index<&'a Q> for TreeMap<K, V, C>
where C: Compare<K> + Compare<Q, K>
{
type Output = V;
#[inline]
fn index(&self, i: &'a Q) -> &V {
self.get(i).expect("no entry found for key")
}
}
impl<'a, K, V, C, Q: ?Sized> ops::IndexMut<&'a Q> for TreeMap<K, V, C>
where C: Compare<K> + Compare<Q, K>
{
#[inline]
fn index_mut(&mut self, i: &'a Q) -> &mut V {
self.get_mut(i).expect("no entry found for key")
}
}
impl<K: Ord, V> TreeMap<K, V> {
pub fn new() -> TreeMap<K, V> {
TreeMap::with_comparator(natural())
}
}
impl<K, V, C> TreeMap<K, V, C>
where C: Compare<K>
{
pub fn with_comparator(cmp: C) -> TreeMap<K, V, C> {
TreeMap {
root: None,
length: 0,
cmp: cmp,
}
}
pub fn comparator(&self) -> &C {
&self.cmp
}
pub fn keys<'a>(&'a self) -> Keys<'a, K, V> {
fn first<A, B>((a, _): (A, B)) -> A {
a
}
let first: fn((&'a K, &'a V)) -> &'a K = first;
Keys(self.iter().map(first))
}
pub fn values<'a>(&'a self) -> Values<'a, K, V> {
fn second<A, B>((_, b): (A, B)) -> B {
b
}
let second: fn((&'a K, &'a V)) -> &'a V = second;
Values(self.iter().map(second))
}
pub fn values_mut<'a>(&'a mut self) -> ValuesMut<'a, K, V> {
fn second<A, B>((_, b): (A, B)) -> B {
b
}
let second: fn((&'a K, &'a mut V)) -> &'a mut V = second;
ValuesMut(self.iter_mut().map(second))
}
pub fn iter(&self) -> Iter<K, V, Forward> {
Iter { iter_mut: (unsafe { &mut *(self as *const Self as *mut Self) }).iter_mut() }
}
fn iter_mut_dir<D: Direction>(&mut self) -> IterMut<K, V, D> {
IterMut {
stack: vec![],
node: deref_mut(&mut self.root),
direction: PhantomData,
}
}
pub fn iter_mut(&mut self) -> IterMut<K, V, Forward> {
self.iter_mut_dir()
}
pub fn into_iter(self) -> IntoIter<K, V> {
let TreeMap { root, length, .. } = self;
let stk = match root {
None => vec![],
Some(b) => vec![*b],
};
IntoIter {
stack: stk,
remaining: length,
}
}
pub fn len(&self) -> usize {
self.length
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn clear(&mut self) {
self.root = None;
self.length = 0
}
#[inline]
pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
where C: Compare<Q, K>
{
fn f<'r, K, V, C, Q: ?Sized>(node: &'r Option<Box<TreeNode<K, V>>>,
cmp: &C,
key: &Q)
-> Option<&'r V>
where C: Compare<Q, K>
{
tree_find_with(node, |k| cmp.compare(key, k))
}
f(&self.root, &self.cmp, key)
}
#[inline]
pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool
where C: Compare<Q, K>
{
self.get(key).is_some()
}
#[inline]
pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V>
where C: Compare<Q, K>
{
fn f<'r, K, V, C, Q: ?Sized>(node: &'r mut Option<Box<TreeNode<K, V>>>,
cmp: &C,
key: &Q)
-> Option<&'r mut V>
where C: Compare<Q, K>
{
tree_find_with_mut(node, |k| cmp.compare(key, k))
}
f(&mut self.root, &self.cmp, key)
}
pub fn insert(&mut self, key: K, value: V) -> Option<V> {
let mut val = Some(value);
let ret = self.get_or_insert(key, || val.take().unwrap());
match val {
None => None,
Some(val) => Some(replace(ret, val)),
}
}
pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
where C: Compare<Q, K>
{
let ret = remove(&mut self.root, key, &self.cmp);
if ret.is_some() {
self.length -= 1
}
ret
}
pub fn get_or_insert<F>(&mut self, key: K, default: F) -> &mut V
where F: FnOnce() -> V
{
let (inserted, ret) = insert(&mut self.root, key, default, &self.cmp);
self.length += inserted;
unsafe { &mut *ret }
}
#[inline]
pub fn find_with<F>(&self, f: F) -> Option<&V>
where F: FnMut(&K) -> Ordering
{
tree_find_with(&self.root, f)
}
#[inline]
pub fn find_with_mut<F>(&mut self, f: F) -> Option<&mut V>
where F: FnMut(&K) -> Ordering
{
tree_find_with_mut(&mut self.root, f)
}
}
impl<K, V, C> TreeMap<K, V, C>
where C: Compare<K>
{
fn compare_bound<D, Q: ?Sized>(&self, bound: Bound<&Q>, key: &K) -> Ordering
where C: Compare<Q, K>,
D: Direction
{
if D::forward() {
match bound {
Bound::Unbounded => Less,
Bound::Included(k) => self.cmp.compare(k, key),
Bound::Excluded(k) => {
match self.cmp.compare(k, key) {
Less => Less,
Greater | Equal => Greater,
}
}
}
} else {
match bound {
Bound::Unbounded => Greater,
Bound::Included(k) => self.cmp.compare(k, key),
Bound::Excluded(k) => {
match self.cmp.compare(k, key) {
Less | Equal => Less,
Greater => Greater,
}
}
}
}
}
fn bound_setup<'a, D, Q: ?Sized>(&self,
mut iter: IterMut<'a, K, V, D>,
bound: Bound<&Q>)
-> IterMut<'a, K, V, D>
where C: Compare<Q, K>,
D: Direction
{
loop {
if !iter.node.is_null() {
let node_k = unsafe { &(*iter.node).key };
match self.compare_bound::<D, Q>(bound, node_k) {
Less => iter.traverse_before(),
Greater => iter.traverse_after(),
Equal => {
iter.traverse_complete();
return iter;
}
}
} else {
iter.traverse_complete();
return iter;
}
}
}
pub fn range_mut<'a, Min: ?Sized, Max: ?Sized>(&'a mut self,
min: Bound<&Min>,
max: Bound<&Max>)
-> RangeMut<'a, K, V>
where C: Compare<Min, K> + Compare<Max, K>
{
let none = None;
let mut node = &self.root;
loop {
match node {
&None => {
return RangeMut {
start: (unsafe { &mut *(self as *const Self as *mut Self) }).iter_mut_dir(),
end: (unsafe { &mut *(self as *const Self as *mut Self) }).iter_mut_dir(),
empty: true,
}
}
&Some(ref n) => {
match (self.compare_bound::<Forward, Min>(min, &n.key),
self.compare_bound::<Backward, Max>(max, &n.key)) {
(Less, Less) => node = &n.left,
(Greater, Greater) => node = &n.right,
(Equal, Less) | (Greater, Less) | (Greater, Equal) => node = &none,
(Less, Equal) | (Less, Greater) | (Equal, Equal) | (Equal, Greater) => {
let start: IterMut<K, V, Forward> = IterMut {
stack: vec![],
node: n.deref() as *const TreeNode<K, V> as *mut TreeNode<K, V>,
direction: PhantomData,
};
let end: IterMut<K, V, Backward> = IterMut {
stack: vec![],
node: n.deref() as *const TreeNode<K, V> as *mut TreeNode<K, V>,
direction: PhantomData,
};
return RangeMut {
start: self.bound_setup(start, min),
end: self.bound_setup(end, max),
empty: false,
};
}
};
}
}
}
}
pub fn range<'a, Min: ?Sized, Max: ?Sized>(&'a self,
min: Bound<&Min>,
max: Bound<&Max>)
-> Range<'a, K, V>
where C: Compare<Min, K> + Compare<Max, K>
{
Range { range: (unsafe { &mut *(self as *const Self as *mut Self) }).range_mut(min, max) }
}
}
pub struct RangeMut<'a, K: 'a, V: 'a> {
start: IterMut<'a, K, V, Forward>,
end: IterMut<'a, K, V, Backward>,
empty: bool,
}
pub struct Range<'a, K: 'a, V: 'a> {
range: RangeMut<'a, K, V>,
}
pub trait Direction {
fn forward() -> bool;
}
pub enum Forward {}
impl Direction for Forward {
fn forward() -> bool {
true
}
}
pub enum Backward {}
impl Direction for Backward {
fn forward() -> bool {
false
}
}
pub struct Iter<'a, K: 'a, V: 'a, D: Direction> {
iter_mut: IterMut<'a, K, V, D>,
}
pub struct IterMut<'a, K: 'a, V: 'a, D: Direction> {
stack: Vec<&'a mut TreeNode<K, V>>,
node: *mut TreeNode<K, V>,
direction: PhantomData<D>,
}
pub struct Keys<'a, K: 'a, V: 'a>(iter::Map<Iter<'a, K, V, Forward>, fn((&'a K, &'a V)) -> &'a K>);
pub struct Values<'a, K: 'a, V: 'a>(iter::Map<Iter<'a, K, V, Forward>,
fn((&'a K, &'a V)) -> &'a V>);
pub struct ValuesMut<'a, K: 'a, V: 'a>(iter::Map<IterMut<'a, K, V, Forward>,
fn((&'a K, &'a mut V)) -> &'a mut V>);
impl<'a, K, V, D: Direction> IterMut<'a, K, V, D> {
#[inline(always)]
fn next_(&mut self) -> Option<(&'a K, &'a mut V)> {
self.normalize();
self.next_node()
}
fn normalize(&mut self) {
while !self.node.is_null() {
let node = unsafe { &mut *self.node };
{
let next_node = if D::forward() {
&mut node.left
} else {
&mut node.right
};
self.node = deref_mut(next_node);
}
self.stack.push(node);
}
}
fn next_node(&mut self) -> Option<(&'a K, &'a mut V)> {
return self.stack.pop().map(|node| {
let next_node = if D::forward() {
&mut node.right
} else {
&mut node.left
};
self.node = deref_mut(next_node);
(&node.key, &mut node.value)
});
}
#[inline]
fn traverse_before(&mut self) {
let node = unsafe { &mut *self.node };
self.node = deref_mut(&mut node.left);
if D::forward() {
self.stack.push(node);
}
}
#[inline]
fn traverse_after(&mut self) {
let node = unsafe { &mut *self.node };
self.node = deref_mut(&mut node.right);
if !D::forward() {
self.stack.push(node);
}
}
#[inline]
fn traverse_complete(&mut self) {
if !self.node.is_null() {
unsafe {
self.stack.push(&mut *self.node);
}
self.node = ptr::null_mut();
}
}
}
impl<'a, K, V, D: Direction> Iterator for IterMut<'a, K, V, D> {
type Item = (&'a K, &'a mut V);
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
self.next_()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(0, None)
}
}
impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
type Item = (&'a K, &'a mut V);
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.empty {
return None;
}
self.start.normalize();
self.end.normalize();
self.empty = match (self.start.stack.last(), self.end.stack.last()) {
(None, _) | (_, None) => true,
(Some(n1), Some(n2)) => *n1 as *const TreeNode<K, V> == *n2 as *const TreeNode<K, V>,
};
self.start.next_node()
}
}
impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> {
fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.empty {
return None;
}
self.start.normalize();
self.end.normalize();
self.empty = match (self.start.stack.last(), self.end.stack.last()) {
(None, _) | (_, None) => true,
(Some(n1), Some(n2)) => *n1 as *const TreeNode<K, V> == *n2 as *const TreeNode<K, V>,
};
self.end.next_node()
}
}
impl<'a, K, V> Iterator for Range<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<(&'a K, &'a V)> {
self.range.next().map(|o| match o {
(k, v) => (k, &*v),
})
}
}
impl<'a, K, V> DoubleEndedIterator for Range<'a, K, V> {
fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
self.range.next_back().map(|o| match o {
(k, v) => (k, &*v),
})
}
}
impl<'a, K, V, D: Direction> Iterator for Iter<'a, K, V, D> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<(&'a K, &'a V)> {
self.iter_mut.next_().map(|o| match o {
(k, v) => (k, &*v),
})
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(0, None)
}
}
fn deref_mut<K, V>(x: &mut Option<Box<TreeNode<K, V>>>) -> *mut TreeNode<K, V> {
match *x {
Some(ref mut n) => &mut **n,
None => ptr::null_mut(),
}
}
pub struct IntoIter<K, V> {
stack: Vec<TreeNode<K, V>>,
remaining: usize,
}
impl<K, V> Iterator for IntoIter<K, V> {
type Item = (K, V);
#[inline]
fn next(&mut self) -> Option<(K, V)> {
while !self.stack.is_empty() {
let TreeNode { key, value, left, right, level } = self.stack.pop().unwrap();
match left {
Some(b_left) => {
let n = TreeNode {
key: key,
value: value,
left: None,
right: right,
level: level,
};
self.stack.push(n);
self.stack.push(*b_left);
}
None => {
match right {
Some(b_right) => self.stack.push(*b_right),
None => (),
}
self.remaining -= 1;
return Some((key, value));
}
}
}
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}
impl<'a, K, V> Iterator for Keys<'a, K, V> {
type Item = &'a K;
#[inline]
fn next(&mut self) -> Option<&'a K> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
impl<'a, K, V> Iterator for Values<'a, K, V> {
type Item = &'a V;
#[inline]
fn next(&mut self) -> Option<&'a V> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
impl<'a, K, V> Iterator for ValuesMut<'a, K, V> {
type Item = &'a mut V;
#[inline]
fn next(&mut self) -> Option<&'a mut V> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
#[derive(Clone)]
struct TreeNode<K, V> {
key: K,
value: V,
left: Option<Box<TreeNode<K, V>>>,
right: Option<Box<TreeNode<K, V>>>,
level: usize,
}
impl<K, V> TreeNode<K, V> {
#[inline]
pub fn new(key: K, value: V) -> TreeNode<K, V> {
TreeNode {
key: key,
value: value,
left: None,
right: None,
level: 1,
}
}
}
fn skew<K, V>(node: &mut Box<TreeNode<K, V>>) {
if node.left.as_ref().map_or(false, |x| x.level == node.level) {
let mut save = node.left.take().unwrap();
swap(&mut node.left, &mut save.right); swap(node, &mut save);
node.right = Some(save);
}
}
fn split<K, V>(node: &mut Box<TreeNode<K, V>>) {
if node.right.as_ref().map_or(false,
|x| x.right.as_ref().map_or(false, |y| y.level == node.level)) {
let mut save = node.right.take().unwrap();
swap(&mut node.right, &mut save.left); save.level += 1;
swap(node, &mut save);
node.left = Some(save);
}
}
fn tree_find_with<K, V, F>(node: &Option<Box<TreeNode<K, V>>>, mut f: F) -> Option<&V>
where F: FnMut(&K) -> Ordering
{
let mut current = node;
loop {
match *current {
Some(ref r) => {
match f(&r.key) {
Less => current = &r.left,
Greater => current = &r.right,
Equal => return Some(&r.value),
}
}
None => return None,
}
}
}
fn tree_find_with_mut<K, V, F>(node: &mut Option<Box<TreeNode<K, V>>>, mut f: F) -> Option<&mut V>
where F: FnMut(&K) -> Ordering
{
let mut current = node;
loop {
let temp = current; match *temp {
Some(ref mut r) => {
match f(&r.key) {
Less => current = &mut r.left,
Greater => current = &mut r.right,
Equal => return Some(&mut r.value),
}
}
None => return None,
}
}
}
fn insert<'a, K, V, F, C>(node: &'a mut Option<Box<TreeNode<K, V>>>,
key: K,
default: F,
cmp: &C)
-> (usize, *mut V)
where C: Compare<K>,
K: 'a,
V: 'a,
F: FnOnce() -> V
{
match *node {
Some(ref mut save) => {
match cmp.compare(&key, &save.key) {
Less => {
let ret = insert(&mut save.left, key, default, cmp);
skew(save);
split(save);
ret
}
Greater => {
let ret = insert(&mut save.right, key, default, cmp);
skew(save);
split(save);
ret
}
Equal => (0, &mut save.value),
}
}
None => {
*node = Some(Box::new(TreeNode::new(key, default())));
(1, &mut node.as_mut().unwrap().value)
}
}
}
fn remove<K, V, C, Q: ?Sized>(node: &mut Option<Box<TreeNode<K, V>>>, key: &Q, cmp: &C) -> Option<V>
where C: Compare<Q, K>
{
fn heir_swap<K, V>(node: &mut Box<TreeNode<K, V>>, child: &mut Option<Box<TreeNode<K, V>>>) {
for x in child.iter_mut() {
if x.right.is_some() {
heir_swap(node, &mut x.right);
} else {
swap(&mut node.key, &mut x.key);
swap(&mut node.value, &mut x.value);
}
}
}
match *node {
None => {
return None; }
Some(ref mut save) => {
let (ret, rebalance) = match cmp.compare(key, &save.key) {
Less => (remove(&mut save.left, key, cmp), true),
Greater => (remove(&mut save.right, key, cmp), true),
Equal => {
if save.left.is_some() {
if save.right.is_some() {
let mut left = save.left.take().unwrap();
if left.right.is_some() {
heir_swap(save, &mut left.right);
} else {
swap(&mut save.key, &mut left.key);
swap(&mut save.value, &mut left.value);
}
save.left = Some(left);
(remove(&mut save.left, key, cmp), true)
} else {
let new = save.left.take().unwrap();
let TreeNode { value, .. } = *replace(save, new);
*save = save.left.take().unwrap();
(Some(value), true)
}
} else if save.right.is_some() {
let new = save.right.take().unwrap();
let TreeNode { value, .. } = *replace(save, new);
(Some(value), true)
} else {
(None, false)
}
}
};
if rebalance {
let left_level = save.left.as_ref().map_or(0, |x| x.level);
let right_level = save.right.as_ref().map_or(0, |x| x.level);
if left_level < save.level - 1 || right_level < save.level - 1 {
save.level -= 1;
if right_level > save.level {
let save_level = save.level;
for x in save.right.iter_mut() {
x.level = save_level
}
}
skew(save);
for right in save.right.iter_mut() {
skew(right);
for x in right.right.iter_mut() {
skew(x)
}
}
split(save);
for x in save.right.iter_mut() {
split(x)
}
}
return ret;
}
}
}
return match node.take() {
Some(b) => {
let TreeNode { value, .. } = *b;
Some(value)
}
None => panic!(),
};
}
impl<K, V, C> iter::FromIterator<(K, V)> for TreeMap<K, V, C>
where C: Compare<K> + Default
{
fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> TreeMap<K, V, C> {
let mut map: TreeMap<K, V, C> = Default::default();
map.extend(iter);
map
}
}
impl<K, V, C> Extend<(K, V)> for TreeMap<K, V, C>
where C: Compare<K>
{
#[inline]
fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
for (k, v) in iter {
self.insert(k, v);
}
}
}
impl<K: Hash, V: Hash, C> Hash for TreeMap<K, V, C>
where C: Compare<K>
{
fn hash<H: Hasher>(&self, state: &mut H) {
for elt in self.iter() {
elt.hash(state);
}
}
}
impl<'a, K, V, C> IntoIterator for &'a TreeMap<K, V, C>
where C: Compare<K>
{
type Item = (&'a K, &'a V);
type IntoIter = Iter<'a, K, V, Forward>;
fn into_iter(self) -> Iter<'a, K, V, Forward> {
self.iter()
}
}
impl<'a, K, V, C> IntoIterator for &'a mut TreeMap<K, V, C>
where C: Compare<K>
{
type Item = (&'a K, &'a mut V);
type IntoIter = IterMut<'a, K, V, Forward>;
fn into_iter(self) -> IterMut<'a, K, V, Forward> {
self.iter_mut()
}
}
impl<K, V, C> IntoIterator for TreeMap<K, V, C>
where C: Compare<K>
{
type Item = (K, V);
type IntoIter = IntoIter<K, V>;
fn into_iter(self) -> IntoIter<K, V> {
self.into_iter()
}
}
#[cfg(feature="ordered_iter")]
impl<'a, K, V> ::ordered_iter::OrderedMapIterator for Iter<'a, K, V, Forward> {
type Key = &'a K;
type Val = &'a V;
}
#[cfg(test)]
mod test_treemap {
use rand::{self, Rng};
use super::{TreeMap, TreeNode, Range, RangeMut};
use super::super::Bound;
#[test]
fn find_empty() {
let m: TreeMap<i32, i32> = TreeMap::new();
assert!(m.get(&5) == None);
}
#[test]
fn find_not_found() {
let mut m = TreeMap::new();
assert!(m.insert(1, 2).is_none());
assert!(m.insert(5, 3).is_none());
assert!(m.insert(9, 3).is_none());
assert_eq!(m.get(&2), None);
}
#[test]
fn find_with_empty() {
let m: TreeMap<&'static str, i32> = TreeMap::new();
assert!(m.find_with(|&k| "test".cmp(k)) == None);
}
#[test]
fn find_with_not_found() {
let mut m = TreeMap::new();
assert!(m.insert("test1", 2).is_none());
assert!(m.insert("test2", 3).is_none());
assert!(m.insert("test3", 3).is_none());
assert_eq!(m.find_with(|&k| "test4".cmp(k)), None);
}
#[test]
fn find_with_found() {
let mut m = TreeMap::new();
assert!(m.insert("test1", 2).is_none());
assert!(m.insert("test2", 3).is_none());
assert!(m.insert("test3", 4).is_none());
assert_eq!(m.find_with(|&k| "test2".cmp(k)), Some(&3));
}
#[test]
fn test_find_mut() {
let mut m = TreeMap::new();
assert!(m.insert(1, 12).is_none());
assert!(m.insert(2, 8).is_none());
assert!(m.insert(5, 14).is_none());
let new = 100;
match m.get_mut(&5) {
None => panic!(),
Some(x) => *x = new,
}
assert_eq!(m.get(&5), Some(&new));
}
#[test]
fn test_find_with_mut() {
let mut m = TreeMap::new();
assert!(m.insert("t1", 12).is_none());
assert!(m.insert("t2", 8).is_none());
assert!(m.insert("t5", 14).is_none());
let new = 100;
match m.find_with_mut(|&k| "t5".cmp(k)) {
None => panic!(),
Some(x) => *x = new,
}
assert_eq!(m.find_with(|&k| "t5".cmp(k)), Some(&new));
}
#[test]
fn insert_replace() {
let mut m = TreeMap::new();
assert!(m.insert(5, 2).is_none());
assert!(m.insert(2, 9).is_none());
assert!(!m.insert(2, 11).is_none());
assert_eq!(m.get(&2).unwrap(), &11);
}
#[test]
fn test_get_or_insert() {
let mut m = TreeMap::new();
assert_eq!(*m.get_or_insert(5, || 2), 2);
assert_eq!(*m.get_or_insert(2, || 9), 9);
assert_eq!(*m.get_or_insert(2, || 7), 9);
*m.get_or_insert(2, || 7) = 8;
assert_eq!(*m.get(&2).unwrap(), 8);
}
#[test]
fn test_clear() {
let mut m = TreeMap::new();
m.clear();
assert!(m.insert(5, 11).is_none());
assert!(m.insert(12, -3).is_none());
assert!(m.insert(19, 2).is_none());
m.clear();
assert!(m.get(&5).is_none());
assert!(m.get(&12).is_none());
assert!(m.get(&19).is_none());
assert!(m.is_empty());
}
#[test]
fn u8_map() {
let mut m = TreeMap::new();
let k1 = "foo".as_bytes();
let k2 = "bar".as_bytes();
let v1 = "baz".as_bytes();
let v2 = "foobar".as_bytes();
m.insert(k1.clone(), v1.clone());
m.insert(k2.clone(), v2.clone());
assert_eq!(m.get(&k2), Some(&v2));
assert_eq!(m.get(&k1), Some(&v1));
}
fn check_equal<K: PartialEq + Ord, V: PartialEq>(ctrl: &[(K, V)], map: &TreeMap<K, V>) {
assert_eq!(ctrl.is_empty(), map.is_empty());
for x in ctrl.iter() {
let &(ref k, ref v) = x;
assert!(map.get(k).unwrap() == v)
}
for (map_k, map_v) in map.iter() {
let mut found = false;
for x in ctrl.iter() {
let &(ref ctrl_k, ref ctrl_v) = x;
if *map_k == *ctrl_k {
assert!(*map_v == *ctrl_v);
found = true;
break;
}
}
assert!(found);
}
}
fn check_left<K: Ord, V>(node: &Option<Box<TreeNode<K, V>>>, parent: &Box<TreeNode<K, V>>) {
match *node {
Some(ref r) => {
assert_eq!(r.key.cmp(&parent.key), ::std::cmp::Ordering::Less);
assert!(r.level == parent.level - 1); check_left(&r.left, r);
check_right(&r.right, r, false);
}
None => assert!(parent.level == 1), }
}
fn check_right<K: Ord, V>(node: &Option<Box<TreeNode<K, V>>>,
parent: &Box<TreeNode<K, V>>,
parent_red: bool) {
match *node {
Some(ref r) => {
assert_eq!(r.key.cmp(&parent.key), ::std::cmp::Ordering::Greater);
let red = r.level == parent.level;
if parent_red {
assert!(!red)
} assert!(red || r.level == parent.level - 1);
check_left(&r.left, r);
check_right(&r.right, r, red);
}
None => assert!(parent.level == 1), }
}
fn check_structure<K: Ord, V>(map: &TreeMap<K, V>) {
match map.root {
Some(ref r) => {
check_left(&r.left, r);
check_right(&r.right, r, false);
}
None => (),
}
}
#[test]
fn test_rand_int() {
let mut map: TreeMap<i32, i32> = TreeMap::new();
let mut ctrl = vec![];
check_equal(&ctrl, &map);
assert!(map.get(&5).is_none());
let seed: &[_] = &[42];
let mut rng: rand::IsaacRng = rand::SeedableRng::from_seed(seed);
for _ in 0..3 {
for _ in 0..90 {
let k = rng.gen();
let v = rng.gen();
if !ctrl.iter().any(|x| x == &(k, v)) {
assert!(map.insert(k, v).is_none());
ctrl.push((k, v));
check_structure(&map);
check_equal(&ctrl, &map);
}
}
for _ in 0..30 {
let r = rng.gen_range(0, ctrl.len());
let (key, _) = ctrl.remove(r);
assert!(map.remove(&key).is_some());
check_structure(&map);
check_equal(&ctrl, &map);
}
}
}
#[test]
fn test_len() {
let mut m = TreeMap::new();
assert!(m.insert(3, 6).is_none());
assert_eq!(m.len(), 1);
assert!(m.insert(0, 0).is_none());
assert_eq!(m.len(), 2);
assert!(m.insert(4, 8).is_none());
assert_eq!(m.len(), 3);
assert!(m.remove(&3).is_some());
assert_eq!(m.len(), 2);
assert!(!m.remove(&5).is_some());
assert_eq!(m.len(), 2);
assert!(m.insert(2, 4).is_none());
assert_eq!(m.len(), 3);
assert!(m.insert(1, 2).is_none());
assert_eq!(m.len(), 4);
}
#[test]
fn test_iterator() {
let mut m = TreeMap::new();
assert!(m.insert(3, 6).is_none());
assert!(m.insert(0, 0).is_none());
assert!(m.insert(4, 8).is_none());
assert!(m.insert(2, 4).is_none());
assert!(m.insert(1, 2).is_none());
let mut n = 0;
for (k, v) in m.iter() {
assert_eq!(*k, n);
assert_eq!(*v, n * 2);
n += 1;
}
assert_eq!(n, 5);
}
#[test]
fn test_interval_iteration() {
let mut m = TreeMap::new();
for i in 1..100 {
assert!(m.insert(i * 2, i * 4).is_none());
}
for i in 1..198 {
let mut lb_it = m.range(Bound::Included(&i), Bound::Unbounded);
let (&k, &v) = lb_it.next().unwrap();
let lb = i + i % 2;
assert_eq!(lb, k);
assert_eq!(lb * 2, v);
let mut ub_it = m.range(Bound::Excluded(&i), Bound::Unbounded);
let (&k, &v) = ub_it.next().unwrap();
let ub = i + 2 - i % 2;
assert_eq!(ub, k);
assert_eq!(ub * 2, v);
}
let mut end_it = m.range(Bound::Included(&199), Bound::Unbounded);
assert_eq!(end_it.next(), None);
}
#[test]
fn test_mut_iter() {
let mut m = TreeMap::new();
for i in 0..10 {
assert!(m.insert(i, 100 * i).is_none());
}
for (i, (&k, v)) in m.iter_mut().enumerate() {
*v += k * 10 + i; }
for (&k, &v) in m.iter() {
assert_eq!(v, 111 * k);
}
}
#[test]
fn test_mut_interval_iter() {
let mut m_lower = TreeMap::new();
let mut m_upper = TreeMap::new();
for i in 1..100 {
assert!(m_lower.insert(i * 2, i * 4).is_none());
assert!(m_upper.insert(i * 2, i * 4).is_none());
}
for i in 1..199 {
let mut lb_it = m_lower.range_mut(Bound::Included(&i), Bound::Unbounded);
let (&k, v) = lb_it.next().unwrap();
let lb = i + i % 2;
assert_eq!(lb, k);
*v -= k;
}
for i in 0..198 {
let mut ub_it = m_upper.range_mut(Bound::Excluded(&i), Bound::Unbounded);
let (&k, v) = ub_it.next().unwrap();
let ub = i + 2 - i % 2;
assert_eq!(ub, k);
*v -= k;
}
assert!(m_lower.range_mut(Bound::Included(&199), Bound::Unbounded).next().is_none());
assert!(m_upper.range_mut(Bound::Excluded(&199), Bound::Unbounded).next().is_none());
assert!(m_lower.iter().all(|(_, &x)| x == 0));
assert!(m_upper.iter().all(|(_, &x)| x == 0));
}
fn to_vec<K, V: Clone>(range: Range<K, V>) -> Vec<V> {
range.map(|o| o.1.clone()).collect::<Vec<V>>()
}
fn skip<K, V>(mut range: Range<K, V>, mut f: u32, mut b: u32) -> Range<K, V> {
while f > 0 {
range.next();
f -= 1;
}
while b > 0 {
range.next_back();
b -= 1;
}
range
}
#[test]
fn test_range() {
let mut m = TreeMap::new();
for i in 0..10 {
assert!(m.insert(i * 10, 100 * i).is_none());
}
assert_eq!(to_vec(m.range(Bound::Unbounded, Bound::Unbounded)),
vec![0, 100, 200, 300, 400, 500, 600, 700, 800, 900]);
assert_eq!(to_vec(m.range(Bound::Included(&50), Bound::Unbounded)),
vec![500, 600, 700, 800, 900]);
assert_eq!(to_vec(m.range(Bound::Excluded(&50), Bound::Unbounded)),
vec![600, 700, 800, 900]);
assert_eq!(to_vec(m.range(Bound::Unbounded, Bound::Included(&70))),
vec![0, 100, 200, 300, 400, 500, 600, 700]);
assert_eq!(to_vec(m.range(Bound::Unbounded, Bound::Excluded(&70))),
vec![0, 100, 200, 300, 400, 500, 600]);
assert_eq!(to_vec(m.range(Bound::Included(&70), Bound::Included(&70))),
vec![700]);
assert_eq!(to_vec(m.range(Bound::Included(&60), Bound::Excluded(&50))),
vec![]);
assert_eq!(to_vec(m.range(Bound::Unbounded, Bound::Excluded(&0))),
vec![]);
assert_eq!(m.range(Bound::Unbounded, Bound::Unbounded).next_back(),
Some((&90, &900)));
assert_eq!(m.range(Bound::Unbounded, Bound::Included(&60)).next_back(),
Some((&60, &600)));
assert_eq!(m.range(Bound::Unbounded, Bound::Excluded(&60)).next_back(),
Some((&50, &500)));
assert_eq!(to_vec(skip(m.range(Bound::Unbounded, Bound::Unbounded), 3, 3)),
vec![300, 400, 500, 600]);
assert_eq!(to_vec(skip(m.range(Bound::Unbounded, Bound::Unbounded), 5, 5)),
vec![]);
assert_eq!(to_vec(skip(m.range(Bound::Unbounded, Bound::Unbounded), 3, 7)),
vec![]);
assert_eq!(to_vec(skip(m.range(Bound::Unbounded, Bound::Unbounded), 7, 3)),
vec![]);
assert_eq!(to_vec(skip(m.range(Bound::Unbounded, Bound::Unbounded), 8, 3)),
vec![]);
}
fn to_vec_mut<K, V: Clone>(range: RangeMut<K, V>) -> Vec<V> {
range.map(|o| o.1.clone()).collect::<Vec<V>>()
}
fn skip_mut<K, V>(mut range: RangeMut<K, V>, mut f: u32, mut b: u32) -> RangeMut<K, V> {
while f > 0 {
range.next();
f -= 1;
}
while b > 0 {
range.next_back();
b -= 1;
}
range
}
#[test]
fn test_range_mut() {
let mut m = TreeMap::new();
for i in 0..10 {
assert!(m.insert(i * 10, 20 * i).is_none());
}
for i in m.range_mut(Bound::Unbounded, Bound::Unbounded) {
match i {
(_, v) => *v *= 5,
}
}
assert_eq!(to_vec_mut(m.range_mut(Bound::Unbounded, Bound::Unbounded)),
vec![0, 100, 200, 300, 400, 500, 600, 700, 800, 900]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Included(&50), Bound::Unbounded)),
vec![500, 600, 700, 800, 900]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Excluded(&50), Bound::Unbounded)),
vec![600, 700, 800, 900]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Unbounded, Bound::Included(&70))),
vec![0, 100, 200, 300, 400, 500, 600, 700]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Unbounded, Bound::Excluded(&70))),
vec![0, 100, 200, 300, 400, 500, 600]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Included(&70), Bound::Included(&70))),
vec![700]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Included(&60), Bound::Excluded(&50))),
vec![]);
assert_eq!(to_vec_mut(m.range_mut(Bound::Unbounded, Bound::Excluded(&0))),
vec![]);
assert_eq!(m.range_mut(Bound::Unbounded, Bound::Unbounded).next_back(),
Some((&90, &mut 900)));
assert_eq!(m.range_mut(Bound::Unbounded, Bound::Included(&60)).next_back(),
Some((&60, &mut 600)));
assert_eq!(m.range_mut(Bound::Unbounded, Bound::Excluded(&60)).next_back(),
Some((&50, &mut 500)));
assert_eq!(to_vec_mut(skip_mut(m.range_mut(Bound::Unbounded, Bound::Unbounded), 3, 3)),
vec![300, 400, 500, 600]);
assert_eq!(to_vec_mut(skip_mut(m.range_mut(Bound::Unbounded, Bound::Unbounded), 5, 5)),
vec![]);
assert_eq!(to_vec_mut(skip_mut(m.range_mut(Bound::Unbounded, Bound::Unbounded), 3, 7)),
vec![]);
assert_eq!(to_vec_mut(skip_mut(m.range_mut(Bound::Unbounded, Bound::Unbounded), 7, 3)),
vec![]);
assert_eq!(to_vec_mut(skip_mut(m.range_mut(Bound::Unbounded, Bound::Unbounded), 8, 3)),
vec![]);
}
#[test]
fn test_keys() {
let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
let map: TreeMap<i32, char> = vec.into_iter().collect();
let keys: Vec<i32> = map.keys().map(|&k| k).collect();
assert_eq!(keys, vec![1, 2, 3]);
}
#[test]
fn test_values() {
let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
let map = vec.into_iter().collect::<TreeMap<i32, char>>();
let values = map.values().map(|&v| v).collect::<Vec<char>>();
assert_eq!(values, vec!['a', 'b', 'c']);
}
#[test]
fn test_values_mut() {
let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
let mut map = vec.into_iter().collect::<TreeMap<i32, char>>();
for ch in map.values_mut() {
*ch = 'x';
}
let values = map.values().map(|&v| v).collect::<Vec<char>>();
assert_eq!(values, vec!['x', 'x', 'x']);
}
#[test]
fn test_eq() {
let mut a = TreeMap::new();
let mut b = TreeMap::new();
assert!(a == b);
assert!(a.insert(0, 5).is_none());
assert!(a != b);
assert!(b.insert(0, 4).is_none());
assert!(a != b);
assert!(a.insert(5, 19).is_none());
assert!(a != b);
assert!(!b.insert(0, 5).is_none());
assert!(a != b);
assert!(b.insert(5, 19).is_none());
assert!(a == b);
}
#[test]
fn test_lt() {
let mut a = TreeMap::new();
let mut b = TreeMap::new();
assert!(!(a < b) && !(b < a));
assert!(b.insert(0, 5).is_none());
assert!(a < b);
assert!(a.insert(0, 7).is_none());
assert!(!(a < b) && b < a);
assert!(b.insert(-2, 0).is_none());
assert!(b < a);
assert!(a.insert(-5, 2).is_none());
assert!(a < b);
assert!(a.insert(6, 2).is_none());
assert!(a < b && !(b < a));
}
#[test]
fn test_ord() {
let mut a = TreeMap::new();
let mut b = TreeMap::new();
assert!(a <= b && a >= b);
assert!(a.insert(1, 1).is_none());
assert!(a > b && a >= b);
assert!(b < a && b <= a);
assert!(b.insert(2, 2).is_none());
assert!(b > a && b >= a);
assert!(a < b && a <= b);
}
#[test]
fn test_debug() {
let mut map = TreeMap::new();
let empty: TreeMap<i32, i32> = TreeMap::new();
map.insert(1, 2);
map.insert(3, 4);
assert_eq!(format!("{:?}", map), "{1: 2, 3: 4}");
assert_eq!(format!("{:?}", empty), "{}");
}
#[test]
fn test_lazy_iterator() {
let mut m = TreeMap::new();
let (x1, y1) = (2, 5);
let (x2, y2) = (9, 12);
let (x3, y3) = (20, -3);
let (x4, y4) = (29, 5);
let (x5, y5) = (103, 3);
assert!(m.insert(x1, y1).is_none());
assert!(m.insert(x2, y2).is_none());
assert!(m.insert(x3, y3).is_none());
assert!(m.insert(x4, y4).is_none());
assert!(m.insert(x5, y5).is_none());
let m = m;
let mut a = m.iter();
assert_eq!(a.next().unwrap(), (&x1, &y1));
assert_eq!(a.next().unwrap(), (&x2, &y2));
assert_eq!(a.next().unwrap(), (&x3, &y3));
assert_eq!(a.next().unwrap(), (&x4, &y4));
assert_eq!(a.next().unwrap(), (&x5, &y5));
assert!(a.next().is_none());
let mut b = m.iter();
let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4), (&x5, &y5)];
let mut i = 0;
for x in b.by_ref() {
assert_eq!(expected[i], x);
i += 1;
if i == 2 {
break;
}
}
for x in b {
assert_eq!(expected[i], x);
i += 1;
}
}
#[test]
fn test_from_iter() {
let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
let map: TreeMap<i32, i32> = xs.iter().map(|&x| x).collect();
for &(k, v) in xs.iter() {
assert_eq!(map.get(&k), Some(&v));
}
}
#[test]
fn test_index() {
let mut map: TreeMap<i32, i32> = TreeMap::new();
map.insert(1, 2);
map.insert(2, 1);
map.insert(3, 4);
assert_eq!(map[&2], 1);
}
#[test]
#[should_panic]
fn test_index_nonexistent() {
let mut map: TreeMap<i32, i32> = TreeMap::new();
map.insert(1, 2);
map.insert(2, 1);
map.insert(3, 4);
map[&4];
}
#[test]
fn test_swap() {
let mut m = TreeMap::new();
assert_eq!(m.insert(1, 2), None);
assert_eq!(m.insert(1, 3), Some(2));
assert_eq!(m.insert(1, 4), Some(3));
}
#[test]
fn test_pop() {
let mut m = TreeMap::new();
m.insert(1, 2);
assert_eq!(m.remove(&1), Some(2));
assert_eq!(m.remove(&1), None);
}
#[test]
fn test_comparator_iterator() {
use compare::{Compare, natural};
let mut m = TreeMap::with_comparator(natural().rev());
assert!(m.insert(3, 6).is_none());
assert!(m.insert(0, 0).is_none());
assert!(m.insert(4, 8).is_none());
assert!(m.insert(2, 4).is_none());
assert!(m.insert(1, 2).is_none());
let mut n = 5;
for (k, v) in m.iter() {
n -= 1;
assert_eq!(*k, n);
assert_eq!(*v, n * 2);
}
assert_eq!(n, 0);
}
#[test]
fn test_comparator_borrowed() {
use compare::{Compare, natural};
let mut m = TreeMap::with_comparator(natural().borrowing());
assert!(m.insert("a".to_string(), 1).is_none());
assert!(m.contains_key("a"));
assert!(m.contains_key(&"a"));
assert!(m.contains_key(&"a".to_string()));
assert_eq!(m.get("a"), Some(&1));
assert_eq!(m.get(&"a"), Some(&1));
assert_eq!(m.get(&"a".to_string()), Some(&1));
m["a"] = 2;
assert_eq!(m["a"], 2);
assert_eq!(m[&"a".to_string()], 2);
m[&"a".to_string()] = 3;
assert_eq!(m.remove("a"), Some(3));
assert!(m.remove(&"a").is_none());
assert!(m.remove(&"a".to_string()).is_none());
}
}
#[cfg(all(test, feature="bench"))]
mod bench {
use rand::{weak_rng, Rng};
use test::{Bencher, black_box};
use super::TreeMap;
map_insert_rand_bench!{insert_rand_100, 100, TreeMap}
map_insert_rand_bench!{insert_rand_10_000, 10_000, TreeMap}
map_insert_seq_bench!{insert_seq_100, 100, TreeMap}
map_insert_seq_bench!{insert_seq_10_000, 10_000, TreeMap}
map_find_rand_bench!{find_rand_100, 100, TreeMap}
map_find_rand_bench!{find_rand_10_000, 10_000, TreeMap}
map_find_seq_bench!{find_seq_100, 100, TreeMap}
map_find_seq_bench!{find_seq_10_000, 10_000, TreeMap}
fn bench_iter(b: &mut Bencher, size: usize) {
let mut map = TreeMap::<u32, u32>::new();
let mut rng = weak_rng();
for _ in 0..size {
map.insert(rng.gen(), rng.gen());
}
b.iter(|| {
for entry in map.iter() {
black_box(entry);
}
});
}
#[bench]
pub fn iter_20(b: &mut Bencher) {
bench_iter(b, 20);
}
#[bench]
pub fn iter_1000(b: &mut Bencher) {
bench_iter(b, 1000);
}
#[bench]
pub fn iter_100000(b: &mut Bencher) {
bench_iter(b, 100000);
}
}