rust/src/libstd/treemap.rs

149 lines
3.4 KiB
Rust
Raw Normal View History

/*!
* A key,value store that works on anything.
*
* This works using a binary search tree. In the first version, it's a
* very naive algorithm, but it will probably be updated to be a
* red-black tree or something else.
*/
2012-09-02 20:09:48 -05:00
#[forbid(deprecated_mode)];
#[forbid(deprecated_pattern)];
2012-09-04 13:23:53 -05:00
use core::cmp::{Eq, Ord};
use core::option::{Some, None};
use Option = core::Option;
2012-09-04 16:35:43 -05:00
export TreeMap;
export insert;
export find;
export traverse;
2012-09-04 16:35:43 -05:00
type TreeMap<K, V> = @mut TreeEdge<K, V>;
2012-05-23 19:18:31 -05:00
2012-09-04 16:35:43 -05:00
type TreeEdge<K, V> = Option<@TreeNode<K, V>>;
2012-09-04 16:35:43 -05:00
enum TreeNode<K, V> = {
2012-05-23 19:18:31 -05:00
key: K,
mut value: V,
2012-09-04 16:35:43 -05:00
mut left: TreeEdge<K, V>,
mut right: TreeEdge<K, V>
2012-05-23 19:18:31 -05:00
};
/// Create a treemap
2012-09-04 16:35:43 -05:00
fn TreeMap<K, V>() -> TreeMap<K, V> { @mut None }
/// Insert a value into the map
2012-09-04 16:35:43 -05:00
fn insert<K: copy Eq Ord, V: copy>(m: &mut TreeEdge<K, V>, +k: K, +v: V) {
2012-08-06 14:34:08 -05:00
match copy *m {
2012-08-20 14:23:37 -05:00
None => {
2012-09-04 16:35:43 -05:00
*m = Some(@TreeNode({key: k,
2012-05-23 19:18:31 -05:00
mut value: v,
2012-08-20 14:23:37 -05:00
mut left: None,
mut right: None}));
2012-08-01 19:30:05 -05:00
return;
2012-05-23 19:18:31 -05:00
}
2012-08-20 14:23:37 -05:00
Some(node) => {
2012-05-23 19:18:31 -05:00
if k == node.key {
node.value = v;
} else if k < node.key {
insert(&mut node.left, k, v);
} else {
2012-05-23 19:18:31 -05:00
insert(&mut node.right, k, v);
}
}
2012-05-23 19:18:31 -05:00
};
}
/// Find a value based on the key
2012-09-04 16:35:43 -05:00
fn find<K: copy Eq Ord, V: copy>(m: &const TreeEdge<K, V>, +k: K)
2012-08-27 18:26:35 -05:00
-> Option<V> {
2012-08-06 14:34:08 -05:00
match copy *m {
2012-08-20 14:23:37 -05:00
None => None,
2012-05-23 19:18:31 -05:00
// FIXME (#2808): was that an optimization?
2012-08-20 14:23:37 -05:00
Some(node) => {
2012-05-23 19:18:31 -05:00
if k == node.key {
2012-08-20 14:23:37 -05:00
Some(node.value)
2012-05-23 19:18:31 -05:00
} else if k < node.key {
find(&const node.left, k)
} else {
find(&const node.right, k)
}
}
}
}
/// Visit all pairs in the map in order.
2012-09-04 16:35:43 -05:00
fn traverse<K, V: copy>(m: &const TreeEdge<K, V>, f: fn(K, V)) {
2012-08-06 14:34:08 -05:00
match copy *m {
2012-08-20 14:23:37 -05:00
None => (),
Some(node) => {
2012-05-23 19:18:31 -05:00
traverse(&const node.left, f);
// copy of value is req'd as f() requires an immutable ptr
f(node.key, copy node.value);
traverse(&const node.right, f);
}
}
}
2012-01-17 21:05:07 -06:00
#[cfg(test)]
mod tests {
#[test]
2012-09-04 16:35:43 -05:00
fn init_treemap() { let _m = TreeMap::<int, int>(); }
2012-01-17 21:05:07 -06:00
#[test]
2012-09-04 16:35:43 -05:00
fn insert_one() { let m = TreeMap(); insert(m, 1, 2); }
2012-01-17 21:05:07 -06:00
#[test]
2012-09-04 16:35:43 -05:00
fn insert_two() { let m = TreeMap(); insert(m, 1, 2); insert(m, 3, 4); }
2012-01-17 21:05:07 -06:00
#[test]
fn insert_find() {
2012-09-04 16:35:43 -05:00
let m = TreeMap();
2012-01-17 21:05:07 -06:00
insert(m, 1, 2);
2012-08-20 14:23:37 -05:00
assert (find(m, 1) == Some(2));
2012-01-17 21:05:07 -06:00
}
#[test]
fn find_empty() {
2012-09-04 16:35:43 -05:00
let m = TreeMap::<int, int>(); assert (find(m, 1) == None);
2012-01-17 21:05:07 -06:00
}
#[test]
fn find_not_found() {
2012-09-04 16:35:43 -05:00
let m = TreeMap();
2012-01-17 21:05:07 -06:00
insert(m, 1, 2);
2012-08-20 14:23:37 -05:00
assert (find(m, 2) == None);
2012-01-17 21:05:07 -06:00
}
#[test]
fn traverse_in_order() {
2012-09-04 16:35:43 -05:00
let m = TreeMap();
2012-01-17 21:05:07 -06:00
insert(m, 3, ());
insert(m, 0, ());
insert(m, 4, ());
insert(m, 2, ());
insert(m, 1, ());
2012-03-26 20:35:18 -05:00
let n = @mut 0;
2012-09-02 20:09:48 -05:00
fn t(n: @mut int, +k: int, +_v: ()) {
2012-01-17 21:05:07 -06:00
assert (*n == k); *n += 1;
}
2012-06-30 18:19:07 -05:00
traverse(m, |x,y| t(n, x, y));
2012-01-17 21:05:07 -06:00
}
#[test]
fn u8_map() {
2012-09-04 16:35:43 -05:00
let m = TreeMap();
2012-01-17 21:05:07 -06:00
let k1 = str::to_bytes(~"foo");
let k2 = str::to_bytes(~"bar");
2012-01-17 21:05:07 -06:00
insert(m, k1, ~"foo");
insert(m, k2, ~"bar");
2012-01-17 21:05:07 -06:00
2012-08-20 14:23:37 -05:00
assert (find(m, k2) == Some(~"bar"));
assert (find(m, k1) == Some(~"foo"));
2012-01-17 21:05:07 -06:00
}
2012-05-23 19:18:31 -05:00
}