From 393a409b5d418be30f4e959cce74daacad112f75 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Sat, 4 May 2013 09:54:58 -0400 Subject: [PATCH] Add pop() and swap() to the Map trait --- src/libcore/container.rs | 8 ++ src/libcore/hashmap.rs | 65 +++++++--------- src/libcore/trie.rs | 75 +++++++++++++----- src/libstd/smallintmap.rs | 42 ++++++++-- src/libstd/treemap.rs | 78 +++++++++++++------ .../class-impl-very-parameterized-trait.rs | 4 + 6 files changed, 188 insertions(+), 84 deletions(-) diff --git a/src/libcore/container.rs b/src/libcore/container.rs index 00ea4a93221..37b904bbe63 100644 --- a/src/libcore/container.rs +++ b/src/libcore/container.rs @@ -55,6 +55,14 @@ pub trait Map: 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; + + /// 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; } pub trait Set: Mutable { diff --git a/src/libcore/hashmap.rs b/src/libcore/hashmap.rs index 9b01c1dad06..33fdc98137e 100644 --- a/src/libcore/hashmap.rs +++ b/src/libcore/hashmap.rs @@ -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 HashMap { /// 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 { match self.bucket_for_key_with_hash(hash, &k) { TableFull => { fail!(~"Internal logic error"); } FoundHole(idx) => { @@ -213,14 +213,19 @@ priv impl HashMap { 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 Map for HashMap { /// 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 { + // 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 Map for HashMap { 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 { + let hash = k.hash_keyed(self.k0, self.k1) as uint; + self.pop_internal(hash, k) } } @@ -402,31 +422,6 @@ pub impl HashMap { } } - fn pop(&mut self, k: &K) -> Option { - 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 { - // 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 { diff --git a/src/libcore/trie.rs b/src/libcore/trie.rs index f0756be9944..380ad11ae92 100644 --- a/src/libcore/trie.rs +++ b/src/libcore/trie.rs @@ -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 Map for TrieMap { /// 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 { + 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 { 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, key: uint, idx: uint) } fn insert(count: &mut uint, child: &mut Child, key: uint, value: T, - idx: uint) -> bool { + idx: uint) -> Option { 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(count: &mut uint, child: &mut Child, 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(count: &mut uint, child: &mut Child, key: uint, - idx: uint) -> bool { + idx: uint) -> Option { 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); + } } diff --git a/src/libstd/smallintmap.rs b/src/libstd/smallintmap.rs index 1b72300a178..fc83a39cacf 100644 --- a/src/libstd/smallintmap.rs +++ b/src/libstd/smallintmap.rs @@ -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 { priv v: ~[Option], @@ -119,12 +120,27 @@ impl Map for SmallIntMap { /// 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 { + 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 { + 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); + } } diff --git a/src/libstd/treemap.rs b/src/libstd/treemap.rs index 51695f2fa7d..c8ab48e65c0 100644 --- a/src/libstd/treemap.rs +++ b/src/libstd/treemap.rs @@ -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 Map for TreeMap { /// 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 { + 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 { 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>, } } -fn insert(node: &mut Option<~TreeNode>, key: K, value: V) -> bool { +fn insert(node: &mut Option<~TreeNode>, + key: K, value: V) -> Option { match *node { Some(ref mut save) => { match key.cmp(&save.key) { @@ -599,20 +613,19 @@ fn insert(node: &mut Option<~TreeNode>, 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(node: &mut Option<~TreeNode>, - key: &K) -> bool { + key: &K) -> Option { fn heir_swap(node: &mut ~TreeNode, child: &mut Option<~TreeNode>) { // *could* be done without recursion, but it won't borrow check @@ -628,12 +641,12 @@ fn remove(node: &mut Option<~TreeNode>, 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(node: &mut Option<~TreeNode>, 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(node: &mut Option<~TreeNode>, 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); + } } diff --git a/src/test/run-pass/class-impl-very-parameterized-trait.rs b/src/test/run-pass/class-impl-very-parameterized-trait.rs index cf887758bff..b89bf064092 100644 --- a/src/test/run-pass/class-impl-very-parameterized-trait.rs +++ b/src/test/run-pass/class-impl-very-parameterized-trait.rs @@ -103,6 +103,10 @@ impl Map for cat { false } } + + fn pop(&mut self, _k: &int) -> Option { fail!() } + + fn swap(&mut self, _k: int, _v: T) -> Option { fail!() } } pub impl cat {