auto merge of #6236 : alexcrichton/rust/more-map-methods, r=thestinger

Closes #5392 and #5393

I implemented the pop/swap methods for TrieMap/TreeMap/SmallIntMap, and I also updated all of them such that pop isn't just a remove/insert, but rather it's all done in one operation.

One thing I did notice is that with default methods it'd be really nice to define `insert` and `remove` in terms of `pop` and `swap` (or vice versa, just to have them available).
This commit is contained in:
bors 2013-05-06 23:06:36 -07:00
commit 3225870191
6 changed files with 188 additions and 84 deletions

View File

@ -55,6 +55,14 @@ pub trait Map<K, V>: Mutable {
/// 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;
/// Insert a key-value pair from the map. If the key already had a value
/// present in the map, that value is returned. Otherwise None is returned.
fn swap(&mut self, k: K, v: V) -> Option<V>;
/// Removes a key from the map, returning the value at the key if the key
/// was previously in the map.
fn pop(&mut self, k: &K) -> Option<V>;
}
pub trait Set<T>: Mutable {

View File

@ -24,8 +24,8 @@ use rand::RngUtil;
use rand;
use uint;
use vec;
use util::unreachable;
use kinds::Copy;
use util::{replace, unreachable};
static INITIAL_CAPACITY: uint = 32u; // 2^5
@ -204,7 +204,7 @@ priv impl<K:Hash + Eq,V> HashMap<K, V> {
/// Inserts the key value pair into the buckets.
/// Assumes that there will be a bucket.
/// True if there was no previous entry with that key
fn insert_internal(&mut self, hash: uint, k: K, v: V) -> bool {
fn insert_internal(&mut self, hash: uint, k: K, v: V) -> Option<V> {
match self.bucket_for_key_with_hash(hash, &k) {
TableFull => { fail!(~"Internal logic error"); }
FoundHole(idx) => {
@ -213,14 +213,19 @@ priv impl<K:Hash + Eq,V> HashMap<K, V> {
self.buckets[idx] = Some(Bucket{hash: hash, key: k,
value: v});
self.size += 1;
true
None
}
FoundEntry(idx) => {
debug!("insert overwrite (%?->%?) at idx %?, hash %?",
k, v, idx, hash);
self.buckets[idx] = Some(Bucket{hash: hash, key: k,
value: v});
false
match self.buckets[idx] {
None => { fail!(~"insert_internal: Internal logic error") }
Some(ref mut b) => {
b.hash = hash;
b.key = k;
Some(replace(&mut b.value, v))
}
}
}
}
}
@ -361,6 +366,20 @@ impl<K:Hash + Eq,V> Map<K, V> for HashMap<K, V> {
/// key is replaced by the new value. Return true if the key did
/// not already exist in the map.
fn insert(&mut self, k: K, v: V) -> bool {
self.swap(k, v).is_none()
}
/// Remove a key-value pair from the map. Return true if the key
/// was present in the map, otherwise false.
fn remove(&mut self, k: &K) -> bool {
self.pop(k).is_some()
}
/// Insert a key-value pair from the map. If the key already had a value
/// present in the map, that value is returned. Otherwise None is returned.
fn swap(&mut self, k: K, v: V) -> Option<V> {
// this could be faster.
if self.size >= self.resize_at {
// n.b.: We could also do this after searching, so
// that we do not resize if this call to insert is
@ -375,10 +394,11 @@ impl<K:Hash + Eq,V> Map<K, V> for HashMap<K, V> {
self.insert_internal(hash, k, v)
}
/// Remove a key-value pair from the map. Return true if the key
/// was present in the map, otherwise false.
fn remove(&mut self, k: &K) -> bool {
self.pop(k).is_some()
/// Removes a key from the map, returning the value at the key if the key
/// was previously in the map.
fn pop(&mut self, k: &K) -> Option<V> {
let hash = k.hash_keyed(self.k0, self.k1) as uint;
self.pop_internal(hash, k)
}
}
@ -402,31 +422,6 @@ pub impl<K: Hash + Eq, V> HashMap<K, V> {
}
}
fn pop(&mut self, k: &K) -> Option<V> {
let hash = k.hash_keyed(self.k0, self.k1) as uint;
self.pop_internal(hash, k)
}
fn swap(&mut self, k: K, v: V) -> Option<V> {
// this could be faster.
let hash = k.hash_keyed(self.k0, self.k1) as uint;
let old_value = self.pop_internal(hash, &k);
if self.size >= self.resize_at {
// n.b.: We could also do this after searching, so
// that we do not resize if this call to insert is
// simply going to update a key in place. My sense
// though is that it's worse to have to search through
// buckets to find the right spot twice than to just
// resize in this corner case.
self.expand();
}
self.insert_internal(hash, k, v);
old_value
}
/// Return the value corresponding to the key in the map, or insert
/// and return the value if it doesn't exist.
fn find_or_insert<'a>(&'a mut self, k: K, v: V) -> &'a V {

View File

@ -11,6 +11,7 @@
//! An ordered map and set for integer keys implemented as a radix trie
use prelude::*;
use util::{swap, replace};
// FIXME: #5244: need to manually update the TrieNode constructor
static SHIFT: uint = 4;
@ -110,21 +111,33 @@ impl<T> Map<uint, T> for TrieMap<T> {
/// not already exist in the map.
#[inline(always)]
fn insert(&mut self, key: uint, value: T) -> bool {
let ret = insert(&mut self.root.count,
&mut self.root.children[chunk(key, 0)],
key, value, 1);
if ret { self.length += 1 }
ret
self.swap(key, value).is_none()
}
/// Remove a key-value pair from the map. Return true if the key
/// was present in the map, otherwise false.
#[inline(always)]
fn remove(&mut self, key: &uint) -> bool {
self.pop(key).is_some()
}
/// Insert a key-value pair from the map. If the key already had a value
/// present in the map, that value is returned. Otherwise None is returned.
fn swap(&mut self, key: uint, value: T) -> Option<T> {
let ret = insert(&mut self.root.count,
&mut self.root.children[chunk(key, 0)],
key, value, 1);
if ret.is_none() { self.length += 1 }
ret
}
/// Removes a key from the map, returning the value at the key if the key
/// was previously in the map.
fn pop(&mut self, key: &uint) -> Option<T> {
let ret = remove(&mut self.root.count,
&mut self.root.children[chunk(*key, 0)],
*key, 1);
if ret { self.length -= 1 }
if ret.is_some() { self.length -= 1 }
ret
}
}
@ -289,14 +302,15 @@ fn find_mut<'r, T>(child: &'r mut Child<T>, key: uint, idx: uint)
}
fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
idx: uint) -> bool {
idx: uint) -> Option<T> {
let mut tmp = Nothing;
tmp <-> *child;
let mut added = false;
let ret;
swap(&mut tmp, child);
*child = match tmp {
External(stored_key, stored_value) => {
if stored_key == key {
ret = Some(stored_value);
External(stored_key, value)
} else {
// conflict - split the node
@ -304,46 +318,49 @@ fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
insert(&mut new.count,
&mut new.children[chunk(stored_key, idx)],
stored_key, stored_value, idx + 1);
insert(&mut new.count, &mut new.children[chunk(key, idx)], key,
value, idx + 1);
added = true;
ret = insert(&mut new.count, &mut new.children[chunk(key, idx)],
key, value, idx + 1);
Internal(new)
}
}
Internal(x) => {
let mut x = x;
added = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
value, idx + 1);
ret = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
value, idx + 1);
Internal(x)
}
Nothing => {
*count += 1;
added = true;
ret = None;
External(key, value)
}
};
added
return ret;
}
fn remove<T>(count: &mut uint, child: &mut Child<T>, key: uint,
idx: uint) -> bool {
idx: uint) -> Option<T> {
let (ret, this) = match *child {
External(stored, _) => {
if stored == key { (true, true) } else { (false, false) }
External(stored, _) if stored == key => {
match replace(child, Nothing) {
External(_, value) => (Some(value), true),
_ => fail!()
}
}
External(*) => (None, false),
Internal(ref mut x) => {
let ret = remove(&mut x.count, &mut x.children[chunk(key, idx)],
key, idx + 1);
(ret, x.count == 0)
}
Nothing => (false, false)
Nothing => (None, false)
};
if this {
*child = Nothing;
*count -= 1;
}
ret
return ret;
}
#[cfg(test)]
@ -516,4 +533,20 @@ mod tests {
i += 1;
}
}
#[test]
fn test_swap() {
let mut m = TrieMap::new();
assert!(m.swap(1, 2) == None);
assert!(m.swap(1, 3) == Some(2));
assert!(m.swap(1, 4) == Some(3));
}
#[test]
fn test_pop() {
let mut m = TrieMap::new();
m.insert(1, 2);
assert!(m.pop(&1) == Some(2));
assert!(m.pop(&1) == None);
}
}

View File

@ -16,6 +16,7 @@
use core::container::{Container, Mutable, Map, Set};
use core::old_iter::{BaseIter};
use core::option::{Some, None};
use core::util::replace;
pub struct SmallIntMap<T> {
priv v: ~[Option<T>],
@ -119,12 +120,27 @@ impl<V> Map<uint, V> for SmallIntMap<V> {
/// 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: &uint) -> bool {
if *key >= self.v.len() {
return false;
self.pop(key).is_some()
}
/// Insert a key-value pair from the map. If the key already had a value
/// present in the map, that value is returned. Otherwise None is returned.
fn swap(&mut self, key: uint, value: V) -> Option<V> {
match self.find_mut(&key) {
Some(loc) => { return Some(replace(loc, value)); }
None => ()
}
let removed = self.v[*key].is_some();
self.v[*key] = None;
removed
self.insert(key, value);
return None;
}
/// Removes a key from the map, returning the value at the key if the key
/// was previously in the map.
fn pop(&mut self, key: &uint) -> Option<V> {
if *key >= self.v.len() {
return None;
}
replace(&mut self.v[*key], None)
}
}
@ -237,4 +253,20 @@ mod tests {
// sadly, no sevens were counted
assert!(map.find(&7).is_none());
}
#[test]
fn test_swap() {
let mut m = SmallIntMap::new();
assert!(m.swap(1, 2) == None);
assert!(m.swap(1, 3) == Some(2));
assert!(m.swap(1, 4) == Some(3));
}
#[test]
fn test_pop() {
let mut m = SmallIntMap::new();
m.insert(1, 2);
assert!(m.pop(&1) == Some(2));
assert!(m.pop(&1) == None);
}
}

View File

@ -13,6 +13,7 @@
//! `TotalOrd`.
use core::iterator::*;
use core::util::replace;
// 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
@ -150,16 +151,28 @@ impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> {
/// 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
self.swap(key, value).is_none()
}
/// 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 {
self.pop(key).is_some()
}
/// Insert a key-value pair from the map. If the key already had a value
/// present in the map, that value is returned. Otherwise None is returned.
fn swap(&mut self, key: K, value: V) -> Option<V> {
let ret = insert(&mut self.root, key, value);
if ret.is_none() { self.length += 1 }
ret
}
/// Removes a key from the map, returning the value at the key if the key
/// was previously in the map.
fn pop(&mut self, key: &K) -> Option<V> {
let ret = remove(&mut self.root, key);
if ret { self.length -= 1 }
if ret.is_some() { self.length -= 1 }
ret
}
}
@ -581,7 +594,8 @@ fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>,
}
}
fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K, value: V) -> bool {
fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
key: K, value: V) -> Option<V> {
match *node {
Some(ref mut save) => {
match key.cmp(&save.key) {
@ -599,20 +613,19 @@ fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K, value: V)
}
Equal => {
save.key = key;
save.value = value;
false
Some(replace(&mut save.value, value))
}
}
}
None => {
*node = Some(~TreeNode::new(key, value));
true
None
}
}
}
fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
key: &K) -> bool {
key: &K) -> Option<V> {
fn heir_swap<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>,
child: &mut Option<~TreeNode<K, V>>) {
// *could* be done without recursion, but it won't borrow check
@ -628,12 +641,12 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
match *node {
None => {
return false // bottom of tree
return None; // bottom of tree
}
Some(ref mut save) => {
let (removed, this) = match key.cmp(&save.key) {
Less => (remove(&mut save.left, key), false),
Greater => (remove(&mut save.right, key), false),
let (ret, rebalance) = match key.cmp(&save.key) {
Less => (remove(&mut save.left, key), true),
Greater => (remove(&mut save.right, key), true),
Equal => {
if save.left.is_some() {
if save.right.is_some() {
@ -645,21 +658,24 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
save.value <-> left.value;
}
save.left = Some(left);
remove(&mut save.left, key);
(remove(&mut save.left, key), true)
} else {
let new = save.left.swap_unwrap();
let ~TreeNode{value, _} = replace(save, new);
*save = save.left.swap_unwrap();
(Some(value), true)
}
(true, false)
} else if save.right.is_some() {
*save = save.right.swap_unwrap();
(true, false)
let new = save.right.swap_unwrap();
let ~TreeNode{value, _} = replace(save, new);
(Some(value), true)
} else {
(true, true)
(None, false)
}
}
};
if !this {
if rebalance {
let left_level = save.left.map_default(0, |x| x.level);
let right_level = save.right.map_default(0, |x| x.level);
@ -682,13 +698,13 @@ fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>,
for save.right.each_mut |x| { split(x) }
}
return removed;
return ret;
}
}
}
*node = None;
true
return match replace(node, None) {
Some(~TreeNode{value, _}) => Some(value), None => fail!()
};
}
#[cfg(test)]
@ -1217,4 +1233,20 @@ mod test_set {
let result: Option<(&uint, & &'static str)> = z.next();
assert!(result.is_none());
}
#[test]
fn test_swap() {
let mut m = TreeMap::new();
assert!(m.swap(1, 2) == None);
assert!(m.swap(1, 3) == Some(2));
assert!(m.swap(1, 4) == Some(3));
}
#[test]
fn test_pop() {
let mut m = TreeMap::new();
m.insert(1, 2);
assert!(m.pop(&1) == Some(2));
assert!(m.pop(&1) == None);
}
}

View File

@ -103,6 +103,10 @@ impl<T> Map<int, T> for cat<T> {
false
}
}
fn pop(&mut self, _k: &int) -> Option<T> { fail!() }
fn swap(&mut self, _k: int, _v: T) -> Option<T> { fail!() }
}
pub impl<T> cat<T> {