rust/src/libstd/treemap.rs

1203 lines
36 KiB
Rust
Raw Normal View History

2013-01-14 09:27:26 -06:00
// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
2013-01-14 09:27:26 -06:00
//! An ordered map and set implemented as self-balancing binary search
//! trees. The only requirement for the types is that the key implements
2013-03-02 12:27:29 -06:00
//! `TotalOrd`.
2013-01-14 09:27:26 -06:00
use core::prelude::*;
2013-01-14 09:27:26 -06:00
// This is implemented as an AA tree, which is a simplified variation of
// a red-black tree where where red (horizontal) nodes can only be added
// as a right child. The time complexity is the same, and re-balancing
// operations are more frequent but also cheaper.
2012-05-23 19:18:31 -05:00
// Future improvements:
2013-01-15 08:50:51 -06:00
// range search - O(log n) retrieval of an iterator from some key
// (possibly) implement the overloads Python does for sets:
2013-01-14 09:27:26 -06:00
// * intersection: &
// * difference: -
// * symmetric difference: ^
// * union: |
2013-01-15 11:44:43 -06:00
// These would be convenient since the methods work like `each`
2013-01-14 09:27:26 -06:00
pub struct TreeMap<K, V> {
2013-01-14 09:27:26 -06:00
priv root: Option<~TreeNode<K, V>>,
priv length: uint
}
2013-03-02 12:27:29 -06:00
impl<K: Eq + TotalOrd, V: Eq> Eq for TreeMap<K, V> {
2013-01-14 09:27:26 -06:00
pure fn eq(&self, other: &TreeMap<K, V>) -> bool {
if self.len() != other.len() {
2013-01-14 20:12:49 -06:00
false
2013-01-15 13:25:37 -06:00
} else {
2013-01-14 20:03:28 -06:00
let mut x = self.iter();
let mut y = other.iter();
for self.len().times {
unsafe { // unsafe as a purity workaround
if map_next(&mut x).unwrap() !=
map_next(&mut y).unwrap() {
return false
}
2013-01-14 20:03:28 -06:00
}
2013-01-14 09:27:26 -06:00
}
2013-01-14 20:03:28 -06:00
true
2013-01-14 09:27:26 -06:00
}
}
2013-01-21 16:12:35 -06:00
pure fn ne(&self, other: &TreeMap<K, V>) -> bool { !self.eq(other) }
}
// Lexicographical comparison
2013-03-02 12:27:29 -06:00
pure fn lt<K: Ord + TotalOrd, V>(a: &TreeMap<K, V>,
b: &TreeMap<K, V>) -> bool {
let mut x = a.iter();
let mut y = b.iter();
let (a_len, b_len) = (a.len(), b.len());
for uint::min(a_len, b_len).times {
unsafe { // purity workaround
let (key_a,_) = map_next(&mut x).unwrap();
let (key_b,_) = map_next(&mut y).unwrap();
if *key_a < *key_b { return true; }
if *key_a > *key_b { return false; }
}
};
a_len < b_len
}
2013-03-02 12:27:29 -06:00
impl<K: Ord + TotalOrd, V> Ord for TreeMap<K, V> {
#[inline(always)]
2013-03-10 19:11:10 -05:00
pure fn lt(&self, other: &TreeMap<K, V>) -> bool { lt(self, other) }
#[inline(always)]
2013-03-10 19:11:10 -05:00
pure fn le(&self, other: &TreeMap<K, V>) -> bool { !lt(other, self) }
#[inline(always)]
2013-03-10 19:11:10 -05:00
pure fn ge(&self, other: &TreeMap<K, V>) -> bool { !lt(self, other) }
#[inline(always)]
2013-03-10 19:11:10 -05:00
pure fn gt(&self, other: &TreeMap<K, V>) -> bool { lt(other, self) }
}
impl<'self, K: TotalOrd, V> BaseIter<(&'self K, &'self V)> for TreeMap<K, V> {
2013-02-07 17:05:27 -06:00
/// Visit all key-value pairs in order
pure fn each(&self, f: &fn(&(&self/K, &self/V)) -> bool) {
2013-02-07 17:05:27 -06:00
each(&self.root, f)
}
pure fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}
impl<'self, K: TotalOrd, V>
ReverseIter<(&'self K, &'self V)>
for TreeMap<K, V>
{
/// Visit all key-value pairs in reverse order
pure fn each_reverse(&self, f: &fn(&(&self/K, &self/V)) -> bool) {
each_reverse(&self.root, f);
}
}
2013-03-02 12:27:29 -06:00
impl<K: TotalOrd, V> Container for TreeMap<K, V> {
2013-01-21 20:59:19 -06:00
/// Return the number of elements in the map
pure fn len(&self) -> uint { self.length }
/// Return true if the map contains no elements
pure fn is_empty(&self) -> bool { self.root.is_none() }
}
2013-03-02 12:27:29 -06:00
impl<K: TotalOrd, V> Mutable for TreeMap<K, V> {
/// Clear the map, removing all key-value pairs.
fn clear(&mut self) {
self.root = None;
self.length = 0
}
}
2013-03-02 12:27:29 -06:00
impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
2013-01-21 17:22:03 -06:00
/// Return true if the map contains a value for the specified key
pure fn contains_key(&self, key: &K) -> bool {
self.find(key).is_some()
}
/// Visit all keys in order
pure fn each_key(&self, f: &fn(&K) -> bool) { self.each(|&(k, _)| f(k)) }
2013-01-21 17:22:03 -06:00
/// Visit all values in order
pure fn each_value(&self, f: &fn(&V) -> bool) {
self.each(|&(_, v)| f(v))
}
2013-01-21 17:22:03 -06:00
/// Iterate over the map and mutate the contained values
fn mutate_values(&mut self, f: &fn(&'self K, &'self mut V) -> bool) {
mutate_values(&mut self.root, f);
}
/// Return the value corresponding to the key in the map
pure fn find(&self, key: &K) -> Option<&self/V> {
let mut current: &self/Option<~TreeNode<K, V>> = &self.root;
loop {
match *current {
Some(ref r) => {
2013-03-02 12:27:29 -06:00
match key.cmp(&r.key) {
2013-03-10 18:59:41 -05:00
Less => current = &r.left,
Greater => current = &r.right,
Equal => return Some(&r.value)
}
}
None => return None
}
}
}
2013-01-21 17:22:03 -06:00
/// Insert a key-value pair into the map. An existing value for a
/// key is replaced by the new value. Return true if the key did
/// not already exist in the map.
fn insert(&mut self, key: K, value: V) -> bool {
let ret = insert(&mut self.root, key, value);
if ret { self.length += 1 }
ret
}
/// Remove a key-value pair from the map. Return true if the key
/// was present in the map, otherwise false.
fn remove(&mut self, key: &K) -> bool {
let ret = remove(&mut self.root, key);
if ret { self.length -= 1 }
ret
}
}
2013-03-02 12:27:29 -06:00
pub impl<K: TotalOrd, V> TreeMap<K, V> {
2013-01-14 09:27:26 -06:00
/// Create an empty TreeMap
static pure fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} }
2012-05-23 19:18:31 -05:00
2013-01-14 09:27:26 -06:00
/// Visit all keys in reverse order
pure fn each_key_reverse(&self, f: &fn(&K) -> bool) {
2013-02-07 17:05:27 -06:00
self.each_reverse(|&(k, _)| f(k))
2013-01-14 09:27:26 -06:00
}
/// Visit all values in reverse order
pure fn each_value_reverse(&self, f: &fn(&V) -> bool) {
2013-02-07 17:05:27 -06:00
self.each_reverse(|&(_, v)| f(v))
2013-01-14 09:27:26 -06:00
}
2013-01-14 20:03:28 -06:00
/// Get a lazy iterator over the key-value pairs in the map.
/// Requires that it be frozen (immutable).
2013-01-15 13:25:37 -06:00
pure fn iter(&self) -> TreeMapIterator/&self<K, V> {
TreeMapIterator{stack: ~[], node: &self.root}
2013-01-14 19:41:11 -06:00
}
}
/// Lazy forward iterator over a map
pub struct TreeMapIterator<K, V> {
priv stack: ~[&self/~TreeNode<K, V>],
priv node: &self/Option<~TreeNode<K, V>>
}
/// Advance the iterator to the next node (in order) and return a
/// tuple with a reference to the key and value. If there are no
/// more nodes, return `None`.
2013-03-02 12:27:29 -06:00
pub fn map_next<K, V>(iter: &mut TreeMapIterator/&r<K, V>)
-> Option<(&r/K, &r/V)> {
while !iter.stack.is_empty() || iter.node.is_some() {
match *iter.node {
Some(ref x) => {
iter.stack.push(x);
iter.node = &x.left;
}
None => {
let res = iter.stack.pop();
iter.node = &res.right;
return Some((&res.key, &res.value));
}
2013-01-14 19:41:11 -06:00
}
}
None
}
/// Advance the iterator through the map
2013-03-02 12:27:29 -06:00
pub fn map_advance<K, V>(iter: &mut TreeMapIterator/&r<K, V>,
f: &fn((&r/K, &r/V)) -> bool) {
loop {
match map_next(iter) {
Some(x) => {
if !f(x) { return }
}
None => return
}
}
2013-01-14 19:41:11 -06:00
}
pub struct TreeSet<T> {
2013-01-14 09:27:26 -06:00
priv map: TreeMap<T, ()>
}
2013-03-02 12:27:29 -06:00
impl<T: TotalOrd> BaseIter<T> for TreeSet<T> {
2013-01-14 09:27:26 -06:00
/// Visit all values in order
2013-03-10 19:11:10 -05:00
#[inline(always)]
pure fn each(&self, f: &fn(&T) -> bool) { self.map.each_key(f) }
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-01-14 09:27:26 -06:00
pure fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}
2013-03-02 12:27:29 -06:00
impl<T: TotalOrd> ReverseIter<T> for TreeSet<T> {
/// Visit all values in reverse order
2013-03-10 19:11:10 -05:00
#[inline(always)]
pure fn each_reverse(&self, f: &fn(&T) -> bool) {
self.map.each_key_reverse(f)
}
}
2013-03-02 12:27:29 -06:00
impl<T: Eq + TotalOrd> Eq for TreeSet<T> {
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-01-14 09:27:26 -06:00
pure fn eq(&self, other: &TreeSet<T>) -> bool { self.map == other.map }
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-01-14 09:27:26 -06:00
pure fn ne(&self, other: &TreeSet<T>) -> bool { self.map != other.map }
}
2013-03-02 12:27:29 -06:00
impl<T: Ord + TotalOrd> Ord for TreeSet<T> {
#[inline(always)]
pure fn lt(&self, other: &TreeSet<T>) -> bool { self.map < other.map }
#[inline(always)]
pure fn le(&self, other: &TreeSet<T>) -> bool { self.map <= other.map }
#[inline(always)]
pure fn ge(&self, other: &TreeSet<T>) -> bool { self.map >= other.map }
#[inline(always)]
pure fn gt(&self, other: &TreeSet<T>) -> bool { self.map > other.map }
}
2013-03-02 12:27:29 -06:00
impl<T: TotalOrd> Container for TreeSet<T> {
2013-01-23 15:47:27 -06:00
/// Return the number of elements in the set
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-01-21 20:59:19 -06:00
pure fn len(&self) -> uint { self.map.len() }
2013-01-23 15:47:27 -06:00
/// Return true if the set contains no elements
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-01-21 20:59:19 -06:00
pure fn is_empty(&self) -> bool { self.map.is_empty() }
}
2013-03-02 12:27:29 -06:00
impl<T: TotalOrd> Mutable for TreeSet<T> {
/// Clear the set, removing all values.
2013-03-10 19:11:10 -05:00
#[inline(always)]
fn clear(&mut self) { self.map.clear() }
}
2013-03-02 12:27:29 -06:00
impl<T: TotalOrd> Set<T> for TreeSet<T> {
/// Return true if the set contains a value
2013-03-10 19:11:10 -05:00
#[inline(always)]
pure fn contains(&self, value: &T) -> bool {
self.map.contains_key(value)
}
/// Add a value to the set. Return true if the value was not already
/// present in the set.
2013-03-10 19:11:10 -05:00
#[inline(always)]
fn insert(&mut self, value: T) -> bool { self.map.insert(value, ()) }
/// Remove a value from the set. Return true if the value was
/// present in the set.
2013-03-10 19:11:10 -05:00
#[inline(always)]
fn remove(&mut self, value: &T) -> bool { self.map.remove(value) }
2013-01-29 16:04:25 -06:00
/// Return true if the set has no elements in common with `other`.
/// This is equivalent to checking for an empty intersection.
pure fn is_disjoint(&self, other: &TreeSet<T>) -> bool {
let mut x = self.iter();
let mut y = other.iter();
unsafe { // purity workaround
let mut a = set_next(&mut x);
let mut b = set_next(&mut y);
2013-01-29 16:04:25 -06:00
while a.is_some() && b.is_some() {
let a1 = a.unwrap();
let b1 = b.unwrap();
2013-03-02 12:27:29 -06:00
match a1.cmp(b1) {
Less => a = set_next(&mut x),
Greater => b = set_next(&mut y),
Equal => return false
2013-01-29 16:04:25 -06:00
}
}
}
true
}
/// Return true if the set is a subset of another
2013-03-10 19:11:10 -05:00
#[inline(always)]
pure fn is_subset(&self, other: &TreeSet<T>) -> bool {
other.is_superset(self)
}
/// Return true if the set is a superset of another
pure fn is_superset(&self, other: &TreeSet<T>) -> bool {
let mut x = self.iter();
let mut y = other.iter();
unsafe { // purity workaround
let mut a = set_next(&mut x);
let mut b = set_next(&mut y);
while b.is_some() {
if a.is_none() {
return false
}
let a1 = a.unwrap();
let b1 = b.unwrap();
2013-03-02 12:27:29 -06:00
match a1.cmp(b1) {
Less => (),
Greater => return false,
Equal => b = set_next(&mut y),
}
a = set_next(&mut x);
}
}
true
}
2013-01-15 13:25:37 -06:00
2013-01-14 09:27:26 -06:00
/// Visit the values (in-order) representing the difference
pure fn difference(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) {
2013-01-15 13:25:37 -06:00
let mut x = self.iter();
let mut y = other.iter();
2013-01-15 07:55:13 -06:00
2013-01-15 13:25:37 -06:00
unsafe { // purity workaround
let mut a = set_next(&mut x);
let mut b = set_next(&mut y);
2013-01-15 07:55:13 -06:00
while a.is_some() {
if b.is_none() {
2013-01-15 13:25:37 -06:00
return do a.while_some() |a1| {
if f(a1) { set_next(&mut x) } else { None }
2013-01-15 07:55:13 -06:00
}
}
2013-01-15 13:25:37 -06:00
let a1 = a.unwrap();
let b1 = b.unwrap();
2013-01-15 07:55:13 -06:00
2013-03-02 12:27:29 -06:00
let cmp = a1.cmp(b1);
if cmp == Less {
2013-01-15 07:55:13 -06:00
if !f(a1) { return }
a = set_next(&mut x);
2013-01-15 07:55:13 -06:00
} else {
2013-03-02 12:27:29 -06:00
if cmp == Equal { a = set_next(&mut x) }
b = set_next(&mut y);
2013-01-15 07:55:13 -06:00
}
}
}
2013-01-14 09:27:26 -06:00
}
/// Visit the values (in-order) representing the symmetric difference
pure fn symmetric_difference(&self, other: &TreeSet<T>,
f: &fn(&T) -> bool) {
2013-01-15 13:25:37 -06:00
let mut x = self.iter();
let mut y = other.iter();
2013-01-15 13:25:37 -06:00
unsafe { // purity workaround
let mut a = set_next(&mut x);
let mut b = set_next(&mut y);
2013-01-15 10:41:47 -06:00
while a.is_some() {
if b.is_none() {
2013-01-15 13:25:37 -06:00
return do a.while_some() |a1| {
if f(a1) { set_next(&mut x) } else { None }
2013-01-15 10:41:47 -06:00
}
}
2013-01-15 13:25:37 -06:00
let a1 = a.unwrap();
let b1 = b.unwrap();
2013-01-15 10:41:47 -06:00
2013-03-02 12:27:29 -06:00
let cmp = a1.cmp(b1);
if cmp == Less {
2013-01-15 10:41:47 -06:00
if !f(a1) { return }
a = set_next(&mut x);
2013-01-15 10:41:47 -06:00
} else {
2013-03-02 12:27:29 -06:00
if cmp == Greater {
2013-01-15 10:41:47 -06:00
if !f(b1) { return }
} else {
a = set_next(&mut x);
2013-01-15 10:41:47 -06:00
}
b = set_next(&mut y);
2013-01-15 10:41:47 -06:00
}
}
2013-01-15 13:25:37 -06:00
do b.while_some |b1| {
if f(b1) { set_next(&mut y) } else { None }
2013-01-15 10:41:47 -06:00
}
}
2013-01-14 09:27:26 -06:00
}
/// Visit the values (in-order) representing the intersection
pure fn intersection(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) {
let mut x = self.iter();
let mut y = other.iter();
unsafe { // purity workaround
let mut a = set_next(&mut x);
let mut b = set_next(&mut y);
while a.is_some() && b.is_some() {
let a1 = a.unwrap();
let b1 = b.unwrap();
2013-03-02 12:27:29 -06:00
let cmp = a1.cmp(b1);
if cmp == Less {
a = set_next(&mut x);
} else {
2013-03-02 12:27:29 -06:00
if cmp == Equal {
if !f(a1) { return }
}
b = set_next(&mut y);
}
2013-01-14 09:27:26 -06:00
}
}
}
/// Visit the values (in-order) representing the union
pure fn union(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) {
2013-01-15 13:25:37 -06:00
let mut x = self.iter();
let mut y = other.iter();
2013-01-15 13:25:37 -06:00
unsafe { // purity workaround
let mut a = set_next(&mut x);
let mut b = set_next(&mut y);
2013-01-15 09:21:45 -06:00
while a.is_some() {
if b.is_none() {
2013-01-15 13:25:37 -06:00
return do a.while_some() |a1| {
if f(a1) { set_next(&mut x) } else { None }
2013-01-15 09:21:45 -06:00
}
}
2013-01-15 13:25:37 -06:00
let a1 = a.unwrap();
let b1 = b.unwrap();
2013-01-15 09:21:45 -06:00
2013-03-02 12:27:29 -06:00
let cmp = a1.cmp(b1);
if cmp == Greater {
2013-01-15 09:21:45 -06:00
if !f(b1) { return }
b = set_next(&mut y);
2013-01-15 09:21:45 -06:00
} else {
if !f(a1) { return }
2013-03-02 12:27:29 -06:00
if cmp == Equal {
b = set_next(&mut y);
2013-01-15 09:21:45 -06:00
}
a = set_next(&mut x);
2013-01-15 09:21:45 -06:00
}
}
do b.while_some |b1| {
if f(b1) { set_next(&mut y) } else { None }
}
}
}
}
2012-01-17 21:05:07 -06:00
2013-03-02 12:27:29 -06:00
pub impl <T: TotalOrd> TreeSet<T> {
/// Create an empty TreeSet
2013-03-10 19:11:10 -05:00
#[inline(always)]
static pure fn new() -> TreeSet<T> { TreeSet{map: TreeMap::new()} }
/// Get a lazy iterator over the values in the set.
/// Requires that it be frozen (immutable).
2013-03-10 19:11:10 -05:00
#[inline(always)]
pure fn iter(&self) -> TreeSetIterator/&self<T> {
TreeSetIterator{iter: self.map.iter()}
}
}
2013-01-15 13:25:37 -06:00
/// Lazy forward iterator over a set
pub struct TreeSetIterator<T> {
priv iter: TreeMapIterator/&self<T, ()>
2013-01-15 13:25:37 -06:00
}
/// Advance the iterator to the next node (in order). If this iterator is
/// finished, does nothing.
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-03-02 12:27:29 -06:00
pub fn set_next<T>(iter: &mut TreeSetIterator/&r<T>) -> Option<&r/T> {
do map_next(&mut iter.iter).map |&(value, _)| { value }
}
/// Advance the iterator through the set
2013-03-10 19:11:10 -05:00
#[inline(always)]
2013-03-10 18:57:41 -05:00
pub fn set_advance<T>(iter: &mut TreeSetIterator/&r<T>,
f: &fn(&r/T) -> bool) {
do map_advance(&mut iter.iter) |(k, _)| { f(k) }
2013-01-15 13:25:37 -06:00
}
2013-01-14 09:27:26 -06:00
// Nodes keep track of their level in the tree, starting at 1 in the
// leaves and with a red child sharing the level of the parent.
struct TreeNode<K, V> {
2013-01-14 09:27:26 -06:00
key: K,
value: V,
left: Option<~TreeNode<K, V>>,
right: Option<~TreeNode<K, V>>,
level: uint
}
2013-03-02 12:27:29 -06:00
pub impl<K: TotalOrd, V> TreeNode<K, V> {
2013-01-14 09:27:26 -06:00
#[inline(always)]
static pure fn new(key: K, value: V) -> TreeNode<K, V> {
TreeNode{key: key, value: value, left: None, right: None, level: 1}
}
}
2013-03-02 12:27:29 -06:00
pure fn each<K: TotalOrd, V>(node: &r/Option<~TreeNode<K, V>>,
2013-03-10 18:59:41 -05:00
f: &fn(&(&r/K, &r/V)) -> bool) {
for node.each |x| {
2013-01-14 09:27:26 -06:00
each(&x.left, f);
2013-02-07 17:05:27 -06:00
if f(&(&x.key, &x.value)) { each(&x.right, f) }
}
2013-01-14 09:27:26 -06:00
}
2013-03-02 12:27:29 -06:00
pure fn each_reverse<K: TotalOrd, V>(node: &r/Option<~TreeNode<K, V>>,
2013-03-10 18:59:41 -05:00
f: &fn(&(&r/K, &r/V)) -> bool) {
for node.each |x| {
2013-01-14 09:27:26 -06:00
each_reverse(&x.right, f);
2013-02-07 17:05:27 -06:00
if f(&(&x.key, &x.value)) { each_reverse(&x.left, f) }
}
2013-01-14 09:27:26 -06:00
}
fn mutate_values<K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
f: &fn(&'r K, &'r mut V) -> bool) -> bool {
match *node {
Some(~TreeNode{key: ref key, value: ref mut value, left: ref mut left,
right: ref mut right, _}) => {
if !mutate_values(left, f) { return false }
if !f(key, value) { return false }
if !mutate_values(right, f) { return false }
}
None => return false
}
true
}
2013-01-14 09:27:26 -06:00
// Remove left horizontal link by rotating right
2013-03-02 12:27:29 -06:00
fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
2013-01-14 09:27:26 -06:00
if node.left.map_default(false, |x| x.level == node.level) {
let mut save = node.left.swap_unwrap();
node.left <-> save.right; // save.right now None
*node <-> save;
node.right = Some(save);
2013-01-14 09:27:26 -06:00
}
}
// Remove dual horizontal link by rotating left and increasing level of
// the parent
2013-03-02 12:27:29 -06:00
fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) {
2013-01-15 14:34:46 -06:00
if node.right.map_default(false,
|x| x.right.map_default(false, |y| y.level == node.level)) {
2013-01-14 09:27:26 -06:00
let mut save = node.right.swap_unwrap();
node.right <-> save.left; // save.left now None
save.level += 1;
*node <-> save;
node.left = Some(save);
2013-01-14 09:27:26 -06:00
}
}
2013-03-02 12:27:29 -06:00
fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K,
value: V) -> bool {
match *node {
Some(ref mut save) => {
2013-03-02 12:27:29 -06:00
match key.cmp(&save.key) {
Less => {
2013-01-14 09:27:26 -06:00
let inserted = insert(&mut save.left, key, value);
skew(save);
split(save);
2013-01-14 09:27:26 -06:00
inserted
2013-03-02 12:27:29 -06:00
}
Greater => {
2013-01-14 09:27:26 -06:00
let inserted = insert(&mut save.right, key, value);
skew(save);
split(save);
2013-01-14 09:27:26 -06:00
inserted
2013-03-02 12:27:29 -06:00
}
Equal => {
2013-01-14 09:27:26 -06:00
save.key = key;
save.value = value;
false
2013-03-02 12:27:29 -06:00
}
2013-01-14 09:27:26 -06:00
}
}
None => {
*node = Some(~TreeNode::new(key, value));
true
}
2013-01-14 09:27:26 -06:00
}
}
2013-03-02 12:27:29 -06:00
fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
key: &K) -> bool {
fn heir_swap<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>,
2013-01-14 09:27:26 -06:00
child: &mut Option<~TreeNode<K, V>>) {
// *could* be done without recursion, but it won't borrow check
do child.mutate |mut child| {
2013-01-14 09:27:26 -06:00
if child.right.is_some() {
2013-02-10 20:11:33 -06:00
heir_swap(node, &mut child.right);
2013-01-14 09:27:26 -06:00
} else {
node.key <-> child.key;
node.value <-> child.value;
}
child
}
}
match *node {
None => {
2013-01-14 09:27:26 -06:00
return false // bottom of tree
}
Some(ref mut save) => {
2013-03-02 12:27:29 -06:00
let (removed, this) = match key.cmp(&save.key) {
Less => (remove(&mut save.left, key), false),
Greater => (remove(&mut save.right, key), false),
Equal => {
2013-01-14 09:27:26 -06:00
if save.left.is_some() {
if save.right.is_some() {
let mut left = save.left.swap_unwrap();
if left.right.is_some() {
heir_swap(save, &mut left.right);
} else {
save.key <-> left.key;
save.value <-> left.value;
}
2013-01-29 14:48:18 -06:00
save.left = Some(left);
remove(&mut save.left, key);
2013-01-14 09:27:26 -06:00
} else {
*save = save.left.swap_unwrap();
2013-01-14 09:27:26 -06:00
}
(true, false)
2013-01-14 09:27:26 -06:00
} else if save.right.is_some() {
*save = save.right.swap_unwrap();
(true, false)
2013-01-14 09:27:26 -06:00
} else {
(true, true)
2013-01-14 09:27:26 -06:00
}
2013-03-02 12:27:29 -06:00
}
2013-01-14 09:27:26 -06:00
};
if !this {
let left_level = save.left.map_default(0, |x| x.level);
let right_level = save.right.map_default(0, |x| x.level);
// re-balance, if necessary
if left_level < save.level - 1 || right_level < save.level - 1 {
save.level -= 1;
2013-01-14 09:27:26 -06:00
if right_level > save.level {
do save.right.mutate |mut x| { x.level = save.level; x }
}
2013-01-14 09:27:26 -06:00
skew(save);
2013-01-14 09:27:26 -06:00
for save.right.each_mut |right| {
2013-03-10 18:59:41 -05:00
skew(right);
for right.right.each_mut |x| { skew(x) }
}
split(save);
for save.right.each_mut |x| { split(x) }
}
2013-01-14 09:27:26 -06:00
return removed;
2013-01-14 09:27:26 -06:00
}
}
2013-01-14 09:27:26 -06:00
}
*node = None;
true
2013-01-14 09:27:26 -06:00
}
2013-01-14 09:27:26 -06:00
#[cfg(test)]
mod test_treemap {
2013-03-02 12:27:29 -06:00
use core::prelude::*;
2013-01-14 09:27:26 -06:00
use super::*;
use core::rand::RngUtil;
use core::rand;
2012-01-17 21:05:07 -06:00
#[test]
2013-01-14 09:27:26 -06:00
fn find_empty() {
let m = TreeMap::new::<int, int>(); fail_unless!(m.find(&5) == None);
2013-01-14 09:27:26 -06:00
}
2012-01-17 21:05:07 -06:00
#[test]
2013-01-14 09:27:26 -06:00
fn find_not_found() {
let mut m = TreeMap::new();
fail_unless!(m.insert(1, 2));
fail_unless!(m.insert(5, 3));
fail_unless!(m.insert(9, 3));
fail_unless!(m.find(&2) == None);
2013-01-14 09:27:26 -06:00
}
2012-01-17 21:05:07 -06:00
#[test]
2013-01-14 09:27:26 -06:00
fn insert_replace() {
let mut m = TreeMap::new();
fail_unless!(m.insert(5, 2));
fail_unless!(m.insert(2, 9));
fail_unless!(!m.insert(2, 11));
fail_unless!(m.find(&2).unwrap() == &11);
2013-01-14 09:27:26 -06:00
}
2012-01-17 21:05:07 -06:00
#[test]
fn test_clear() {
let mut m = TreeMap::new();
m.clear();
fail_unless!(m.insert(5, 11));
fail_unless!(m.insert(12, -3));
fail_unless!(m.insert(19, 2));
m.clear();
fail_unless!(m.find(&5).is_none());
fail_unless!(m.find(&12).is_none());
fail_unless!(m.find(&19).is_none());
fail_unless!(m.is_empty());
}
2012-01-17 21:05:07 -06:00
#[test]
2013-01-14 09:27:26 -06:00
fn u8_map() {
let mut m = TreeMap::new();
let k1 = str::to_bytes(~"foo");
let k2 = str::to_bytes(~"bar");
let v1 = str::to_bytes(~"baz");
let v2 = str::to_bytes(~"foobar");
m.insert(copy k1, copy v1);
m.insert(copy k2, copy v2);
2013-01-14 09:27:26 -06:00
fail_unless!(m.find(&k2) == Some(&v2));
fail_unless!(m.find(&k1) == Some(&v1));
2013-01-14 09:27:26 -06:00
}
2013-03-02 12:27:29 -06:00
fn check_equal<K: Eq + TotalOrd, V: Eq>(ctrl: &[(K, V)],
map: &TreeMap<K, V>) {
fail_unless!(ctrl.is_empty() == map.is_empty());
2013-01-14 09:27:26 -06:00
for ctrl.each |x| {
let &(k, v) = x;
fail_unless!(map.find(&k).unwrap() == &v)
2013-01-14 09:27:26 -06:00
}
2013-02-07 21:58:16 -06:00
for map.each |&(map_k, map_v)| {
2013-01-14 09:27:26 -06:00
let mut found = false;
for ctrl.each |x| {
let &(ctrl_k, ctrl_v) = x;
if *map_k == ctrl_k {
fail_unless!(*map_v == ctrl_v);
2013-01-14 09:27:26 -06:00
found = true;
break;
}
}
fail_unless!(found);
2013-01-14 09:27:26 -06:00
}
}
2013-03-02 12:27:29 -06:00
fn check_left<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>,
parent: &~TreeNode<K, V>) {
2013-01-14 09:27:26 -06:00
match *node {
Some(ref r) => {
fail_unless!(r.key.cmp(&parent.key) == Less);
fail_unless!(r.level == parent.level - 1); // left is black
2013-01-14 09:27:26 -06:00
check_left(&r.left, r);
check_right(&r.right, r, false);
}
None => fail_unless!(parent.level == 1) // parent is leaf
2013-01-14 09:27:26 -06:00
}
}
2013-03-02 12:27:29 -06:00
fn check_right<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>,
parent: &~TreeNode<K, V>,
parent_red: bool) {
2013-01-14 09:27:26 -06:00
match *node {
Some(ref r) => {
fail_unless!(r.key.cmp(&parent.key) == Greater);
2013-01-14 09:27:26 -06:00
let red = r.level == parent.level;
if parent_red { fail_unless!(!red) } // no dual horizontal links
2013-03-06 21:09:17 -06:00
// Right red or black
fail_unless!(red || r.level == parent.level - 1);
2013-01-14 09:27:26 -06:00
check_left(&r.left, r);
check_right(&r.right, r, red);
}
None => fail_unless!(parent.level == 1) // parent is leaf
2013-01-14 09:27:26 -06:00
}
}
2013-03-02 12:27:29 -06:00
fn check_structure<K: TotalOrd, V>(map: &TreeMap<K, V>) {
2013-01-14 09:27:26 -06:00
match map.root {
Some(ref r) => {
check_left(&r.left, r);
check_right(&r.right, r, false);
}
None => ()
}
2012-01-17 21:05:07 -06:00
}
#[test]
2013-01-14 09:27:26 -06:00
fn test_rand_int() {
let mut map = TreeMap::new::<int, int>();
let mut ctrl = ~[];
check_equal(ctrl, &map);
fail_unless!(map.find(&5).is_none());
2013-01-14 09:27:26 -06:00
let rng = rand::seeded_rng(&[42]);
2013-01-14 09:27:26 -06:00
for 3.times {
for 90.times {
let k = rng.gen_int();
let v = rng.gen_int();
if !ctrl.contains(&(k, v)) {
fail_unless!(map.insert(k, v));
2013-01-14 09:27:26 -06:00
ctrl.push((k, v));
check_structure(&map);
check_equal(ctrl, &map);
}
}
for 30.times {
let r = rng.gen_uint_range(0, ctrl.len());
let (key, _) = vec::remove(&mut ctrl, r);
fail_unless!(map.remove(&key));
2013-01-14 09:27:26 -06:00
check_structure(&map);
check_equal(ctrl, &map);
}
}
2012-01-17 21:05:07 -06:00
}
#[test]
2013-01-14 09:27:26 -06:00
fn test_len() {
let mut m = TreeMap::new();
fail_unless!(m.insert(3, 6));
fail_unless!(m.len() == 1);
fail_unless!(m.insert(0, 0));
fail_unless!(m.len() == 2);
fail_unless!(m.insert(4, 8));
fail_unless!(m.len() == 3);
fail_unless!(m.remove(&3));
fail_unless!(m.len() == 2);
fail_unless!(!m.remove(&5));
fail_unless!(m.len() == 2);
fail_unless!(m.insert(2, 4));
fail_unless!(m.len() == 3);
fail_unless!(m.insert(1, 2));
fail_unless!(m.len() == 4);
2013-01-14 09:27:26 -06:00
}
#[test]
fn test_each() {
let mut m = TreeMap::new();
fail_unless!(m.insert(3, 6));
fail_unless!(m.insert(0, 0));
fail_unless!(m.insert(4, 8));
fail_unless!(m.insert(2, 4));
fail_unless!(m.insert(1, 2));
2013-01-14 09:27:26 -06:00
let mut n = 0;
2013-02-07 21:58:16 -06:00
for m.each |&(k, v)| {
fail_unless!(*k == n);
fail_unless!(*v == n * 2);
2013-01-14 09:27:26 -06:00
n += 1;
}
2012-01-17 21:05:07 -06:00
}
#[test]
2013-01-14 09:27:26 -06:00
fn test_each_reverse() {
let mut m = TreeMap::new();
fail_unless!(m.insert(3, 6));
fail_unless!(m.insert(0, 0));
fail_unless!(m.insert(4, 8));
fail_unless!(m.insert(2, 4));
fail_unless!(m.insert(1, 2));
2012-01-17 21:05:07 -06:00
2013-01-14 09:27:26 -06:00
let mut n = 4;
2013-02-07 21:58:16 -06:00
for m.each_reverse |&(k, v)| {
fail_unless!(*k == n);
fail_unless!(*v == n * 2);
2013-01-14 09:27:26 -06:00
n -= 1;
2012-01-17 21:05:07 -06:00
}
}
#[test]
2013-01-14 09:27:26 -06:00
fn test_eq() {
let mut a = TreeMap::new();
let mut b = TreeMap::new();
fail_unless!(a == b);
fail_unless!(a.insert(0, 5));
fail_unless!(a != b);
fail_unless!(b.insert(0, 4));
fail_unless!(a != b);
fail_unless!(a.insert(5, 19));
fail_unless!(a != b);
fail_unless!(!b.insert(0, 5));
fail_unless!(a != b);
fail_unless!(b.insert(5, 19));
fail_unless!(a == b);
2013-01-14 09:27:26 -06:00
}
#[test]
fn test_lt() {
let mut a = TreeMap::new();
let mut b = TreeMap::new();
fail_unless!(!(a < b) && !(b < a));
fail_unless!(b.insert(0, 5));
fail_unless!(a < b);
fail_unless!(a.insert(0, 7));
fail_unless!(!(a < b) && !(b < a));
fail_unless!(b.insert(-2, 0));
fail_unless!(b < a);
fail_unless!(a.insert(-5, 2));
fail_unless!(a < b);
fail_unless!(a.insert(6, 2));
fail_unless!(a < b && !(b < a));
}
#[test]
fn test_ord() {
let mut a = TreeMap::new();
let mut b = TreeMap::new();
fail_unless!(a <= b && a >= b);
fail_unless!(a.insert(1, 1));
fail_unless!(a > b && a >= b);
fail_unless!(b < a && b <= a);
fail_unless!(b.insert(2, 2));
fail_unless!(b > a && b >= a);
fail_unless!(a < b && a <= b);
}
2013-01-14 19:41:11 -06:00
#[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);
fail_unless!(m.insert(x1, y1));
fail_unless!(m.insert(x2, y2));
fail_unless!(m.insert(x3, y3));
fail_unless!(m.insert(x4, y4));
fail_unless!(m.insert(x5, y5));
2013-01-14 19:41:11 -06:00
let m = m;
let mut a = m.iter();
2013-01-14 19:41:11 -06:00
fail_unless!(map_next(&mut a).unwrap() == (&x1, &y1));
fail_unless!(map_next(&mut a).unwrap() == (&x2, &y2));
fail_unless!(map_next(&mut a).unwrap() == (&x3, &y3));
fail_unless!(map_next(&mut a).unwrap() == (&x4, &y4));
fail_unless!(map_next(&mut a).unwrap() == (&x5, &y5));
fail_unless!(map_next(&mut a).is_none());
let mut b = m.iter();
let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4),
(&x5, &y5)];
let mut i = 0;
for map_advance(&mut b) |x| {
fail_unless!(expected[i] == x);
i += 1;
if i == 2 {
break
}
}
for map_advance(&mut b) |x| {
fail_unless!(expected[i] == x);
i += 1;
}
2013-01-14 19:41:11 -06:00
}
2013-01-14 09:27:26 -06:00
}
2013-01-14 09:27:26 -06:00
#[cfg(test)]
mod test_set {
use super::*;
#[test]
fn test_clear() {
let mut s = TreeSet::new();
s.clear();
fail_unless!(s.insert(5));
fail_unless!(s.insert(12));
fail_unless!(s.insert(19));
s.clear();
fail_unless!(!s.contains(&5));
fail_unless!(!s.contains(&12));
fail_unless!(!s.contains(&19));
fail_unless!(s.is_empty());
}
2013-01-14 09:27:26 -06:00
#[test]
fn test_disjoint() {
let mut xs = TreeSet::new();
let mut ys = TreeSet::new();
fail_unless!(xs.is_disjoint(&ys));
fail_unless!(ys.is_disjoint(&xs));
fail_unless!(xs.insert(5));
fail_unless!(ys.insert(11));
fail_unless!(xs.is_disjoint(&ys));
fail_unless!(ys.is_disjoint(&xs));
fail_unless!(xs.insert(7));
fail_unless!(xs.insert(19));
fail_unless!(xs.insert(4));
fail_unless!(ys.insert(2));
fail_unless!(ys.insert(-11));
fail_unless!(xs.is_disjoint(&ys));
fail_unless!(ys.is_disjoint(&xs));
fail_unless!(ys.insert(7));
fail_unless!(!xs.is_disjoint(&ys));
fail_unless!(!ys.is_disjoint(&xs));
}
2012-01-17 21:05:07 -06:00
#[test]
2013-01-14 09:27:26 -06:00
fn test_subset_and_superset() {
let mut a = TreeSet::new();
fail_unless!(a.insert(0));
fail_unless!(a.insert(5));
fail_unless!(a.insert(11));
fail_unless!(a.insert(7));
2012-01-17 21:05:07 -06:00
2013-01-14 09:27:26 -06:00
let mut b = TreeSet::new();
fail_unless!(b.insert(0));
fail_unless!(b.insert(7));
fail_unless!(b.insert(19));
fail_unless!(b.insert(250));
fail_unless!(b.insert(11));
fail_unless!(b.insert(200));
fail_unless!(!a.is_subset(&b));
fail_unless!(!a.is_superset(&b));
fail_unless!(!b.is_subset(&a));
fail_unless!(!b.is_superset(&a));
fail_unless!(b.insert(5));
fail_unless!(a.is_subset(&b));
fail_unless!(!a.is_superset(&b));
fail_unless!(!b.is_subset(&a));
fail_unless!(b.is_superset(&a));
2013-01-14 09:27:26 -06:00
}
#[test]
fn test_each() {
let mut m = TreeSet::new();
fail_unless!(m.insert(3));
fail_unless!(m.insert(0));
fail_unless!(m.insert(4));
fail_unless!(m.insert(2));
fail_unless!(m.insert(1));
2013-01-14 09:27:26 -06:00
let mut n = 0;
for m.each |x| {
fail_unless!(*x == n);
2013-01-14 09:27:26 -06:00
n += 1
}
}
#[test]
fn test_each_reverse() {
let mut m = TreeSet::new();
fail_unless!(m.insert(3));
fail_unless!(m.insert(0));
fail_unless!(m.insert(4));
fail_unless!(m.insert(2));
fail_unless!(m.insert(1));
2013-01-14 09:27:26 -06:00
let mut n = 4;
for m.each_reverse |x| {
fail_unless!(*x == n);
2013-01-14 09:27:26 -06:00
n -= 1
}
}
fn check(a: &[int], b: &[int], expected: &[int],
f: &fn(&TreeSet<int>, &TreeSet<int>, f: &fn(&int) -> bool)) {
let mut set_a = TreeSet::new();
let mut set_b = TreeSet::new();
2013-01-14 09:27:26 -06:00
for a.each |x| { fail_unless!(set_a.insert(*x)) }
for b.each |y| { fail_unless!(set_b.insert(*y)) }
2013-01-14 09:27:26 -06:00
let mut i = 0;
for f(&set_a, &set_b) |x| {
fail_unless!(*x == expected[i]);
i += 1;
2013-01-14 09:27:26 -06:00
}
fail_unless!(i == expected.len());
2012-01-17 21:05:07 -06:00
}
2013-01-15 07:55:13 -06:00
#[test]
fn test_intersection() {
fn check_intersection(a: &[int], b: &[int], expected: &[int]) {
check(a, b, expected, |x, y, z| x.intersection(y, z))
}
2013-01-15 07:55:13 -06:00
2013-03-10 19:41:50 -05:00
check_intersection([], [], []);
check_intersection([1, 2, 3], [], []);
check_intersection([], [1, 2, 3], []);
check_intersection([2], [1, 2, 3], [2]);
check_intersection([1, 2, 3], [2], [2]);
check_intersection([11, 1, 3, 77, 103, 5, -5],
[2, 11, 77, -9, -42, 5, 3],
[3, 5, 11, 77]);
}
2013-01-15 07:55:13 -06:00
#[test]
fn test_difference() {
fn check_difference(a: &[int], b: &[int], expected: &[int]) {
check(a, b, expected, |x, y, z| x.difference(y, z))
2013-01-15 09:21:45 -06:00
}
check_difference([], [], []);
check_difference([1, 12], [], [1, 12]);
check_difference([], [1, 2, 3, 9], []);
check_difference([1, 3, 5, 9, 11],
[3, 9],
[1, 5, 11]);
check_difference([-5, 11, 22, 33, 40, 42],
[-12, -5, 14, 23, 34, 38, 39, 50],
[11, 22, 33, 40, 42]);
2013-01-15 09:21:45 -06:00
}
2013-01-15 10:41:47 -06:00
#[test]
fn test_symmetric_difference() {
fn check_symmetric_difference(a: &[int], b: &[int],
expected: &[int]) {
check(a, b, expected, |x, y, z| x.symmetric_difference(y, z))
2013-01-15 10:41:47 -06:00
}
2013-03-10 19:41:50 -05:00
check_symmetric_difference([], [], []);
check_symmetric_difference([1, 2, 3], [2], [1, 3]);
check_symmetric_difference([2], [1, 2, 3], [1, 3]);
check_symmetric_difference([1, 3, 5, 9, 11],
[-2, 3, 9, 14, 22],
[-2, 1, 5, 11, 14, 22]);
2013-01-15 10:41:47 -06:00
}
2013-01-15 09:21:45 -06:00
#[test]
fn test_union() {
fn check_union(a: &[int], b: &[int],
expected: &[int]) {
check(a, b, expected, |x, y, z| x.union(y, z))
2013-01-15 07:55:13 -06:00
}
2013-03-10 19:41:50 -05:00
check_union([], [], []);
check_union([1, 2, 3], [2], [1, 2, 3]);
check_union([2], [1, 2, 3], [1, 2, 3]);
check_union([1, 3, 5, 9, 11, 16, 19, 24],
[-2, 1, 5, 9, 13, 19],
[-2, 1, 3, 5, 9, 11, 13, 16, 19, 24]);
2013-01-15 07:55:13 -06:00
}
2012-05-23 19:18:31 -05:00
}