From f9b63d39730740e002cc31a38a7a8a3a18d3f46d Mon Sep 17 00:00:00 2001 From: Andrew Paseltiner Date: Thu, 27 Aug 2015 11:10:28 -0400 Subject: [PATCH] implement RFC 1194 --- src/libcollections/btree/map.rs | 101 +++++++++++++++++++++++++--- src/libcollections/btree/mod.rs | 8 +++ src/libcollections/btree/set.rs | 28 ++++++++ src/libcollectionstest/btree/set.rs | 49 ++++++++++++++ src/libcollectionstest/lib.rs | 1 + src/libstd/collections/hash/map.rs | 37 ++++++++-- src/libstd/collections/hash/mod.rs | 8 +++ src/libstd/collections/hash/set.rs | 64 ++++++++++++++++++ 8 files changed, 283 insertions(+), 13 deletions(-) diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs index 11d389d85ba..3ed4deccb5d 100644 --- a/src/libcollections/btree/map.rs +++ b/src/libcollections/btree/map.rs @@ -459,7 +459,7 @@ impl BTreeMap { } }); match result { - Finished(ret) => return ret, + Finished(ret) => return ret.map(|(_, v)| v), Continue(new_stack) => stack = new_stack } } @@ -693,16 +693,16 @@ mod stack { impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::Leaf> { /// Removes the key and value in the top element of the stack, then handles underflows as /// described in BTree's pop function. - fn remove_leaf(mut self) -> V { + fn remove_leaf(mut self) -> (K, V) { self.map.length -= 1; // Remove the key-value pair from the leaf that this search stack points to. // Then, note if the leaf is underfull, and promptly forget the leaf and its ptr // to avoid ownership issues. - let (value, mut underflow) = unsafe { - let (_, value) = self.top.from_raw_mut().remove_as_leaf(); + let (key_val, mut underflow) = unsafe { + let key_val = self.top.from_raw_mut().remove_as_leaf(); let underflow = self.top.from_raw().node().is_underfull(); - (value, underflow) + (key_val, underflow) }; loop { @@ -717,7 +717,7 @@ mod stack { self.map.depth -= 1; self.map.root.hoist_lone_child(); } - return value; + return key_val; } Some(mut handle) => { if underflow { @@ -728,7 +728,7 @@ mod stack { } } else { // All done! - return value; + return key_val; } } } @@ -739,7 +739,7 @@ mod stack { impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::LeafOrInternal> { /// Removes the key and value in the top element of the stack, then handles underflows as /// described in BTree's pop function. - pub fn remove(self) -> V { + pub fn remove(self) -> (K, V) { // Ensure that the search stack goes to a leaf. This is necessary to perform deletion // in a BTree. Note that this may put the tree in an inconsistent state (further // described in into_leaf's comments), but this is immediately fixed by the @@ -1208,7 +1208,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { /// Takes the value of the entry out of the map, and returns it. #[stable(feature = "rust1", since = "1.0.0")] pub fn remove(self) -> V { - self.stack.remove() + self.stack.remove().1 } } @@ -1609,3 +1609,86 @@ impl BTreeMap { } } } + +impl super::Recover for BTreeMap where K: Borrow + Ord, Q: Ord { + type Key = K; + + fn get(&self, key: &Q) -> Option<&K> { + let mut cur_node = &self.root; + loop { + match Node::search(cur_node, key) { + Found(handle) => return Some(handle.into_kv().0), + GoDown(handle) => match handle.force() { + Leaf(_) => return None, + Internal(internal_handle) => { + cur_node = internal_handle.into_edge(); + continue; + } + } + } + } + } + + fn take(&mut self, key: &Q) -> Option { + // See `remove` for an explanation of this. + + let mut stack = stack::PartialSearchStack::new(self); + loop { + let result = stack.with(move |pusher, node| { + match Node::search(node, key) { + Found(handle) => { + // Perfect match. Terminate the stack here, and remove the entry + Finished(Some(pusher.seal(handle).remove())) + }, + GoDown(handle) => { + // We need to keep searching, try to go down the next edge + match handle.force() { + // We're at a leaf; the key isn't in here + Leaf(_) => Finished(None), + Internal(internal_handle) => Continue(pusher.push(internal_handle)) + } + } + } + }); + match result { + Finished(ret) => return ret.map(|(k, _)| k), + Continue(new_stack) => stack = new_stack + } + } + } + + fn replace(&mut self, mut key: K) -> Option { + // See `insert` for an explanation of this. + + let mut stack = stack::PartialSearchStack::new(self); + + loop { + let result = stack.with(move |pusher, node| { + match Node::search::(node, &key) { + Found(mut handle) => { + mem::swap(handle.key_mut(), &mut key); + Finished(Some(key)) + }, + GoDown(handle) => { + match handle.force() { + Leaf(leaf_handle) => { + pusher.seal(leaf_handle).insert(key, ()); + Finished(None) + } + Internal(internal_handle) => { + Continue((pusher.push(internal_handle), key, ())) + } + } + } + } + }); + match result { + Finished(ret) => return ret, + Continue((new_stack, renewed_key, _)) => { + stack = new_stack; + key = renewed_key; + } + } + } + } +} diff --git a/src/libcollections/btree/mod.rs b/src/libcollections/btree/mod.rs index 282128099da..a235021e639 100644 --- a/src/libcollections/btree/mod.rs +++ b/src/libcollections/btree/mod.rs @@ -11,3 +11,11 @@ mod node; pub mod map; pub mod set; + +trait Recover { + type Key; + + fn get(&self, key: &Q) -> Option<&Self::Key>; + fn take(&mut self, key: &Q) -> Option; + fn replace(&mut self, key: Self::Key) -> Option; +} diff --git a/src/libcollections/btree/set.rs b/src/libcollections/btree/set.rs index a942d6aa669..fe680c72615 100644 --- a/src/libcollections/btree/set.rs +++ b/src/libcollections/btree/set.rs @@ -19,6 +19,7 @@ use core::ops::{BitOr, BitAnd, BitXor, Sub}; use borrow::Borrow; use btree_map::{BTreeMap, Keys}; +use super::Recover; use Bound; // FIXME(conventions): implement bounded iterators @@ -329,6 +330,16 @@ impl BTreeSet { self.map.contains_key(value) } + /// Returns a reference to the value in the set, if any, that is equal to the given value. + /// + /// The value may be any borrowed form of the set's value type, + /// but the ordering on the borrowed form *must* match the + /// ordering on the value type. + #[unstable(feature = "set_recovery", issue = "28050")] + pub fn get(&self, value: &Q) -> Option<&T> where T: Borrow, Q: Ord { + Recover::get(&self.map, value) + } + /// Returns `true` if the set has no elements in common with `other`. /// This is equivalent to checking for an empty intersection. /// @@ -436,6 +447,13 @@ impl BTreeSet { self.map.insert(value, ()).is_none() } + /// Adds a value to the set, replacing the existing value, if any, that is equal to the given + /// one. Returns the replaced value. + #[unstable(feature = "set_recovery", issue = "28050")] + pub fn replace(&mut self, value: T) -> Option { + Recover::replace(&mut self.map, value) + } + /// Removes a value from the set. Returns `true` if the value was /// present in the set. /// @@ -458,6 +476,16 @@ impl BTreeSet { pub fn remove(&mut self, value: &Q) -> bool where T: Borrow, Q: Ord { self.map.remove(value).is_some() } + + /// Removes and returns the value in the set, if any, that is equal to the given one. + /// + /// The value may be any borrowed form of the set's value type, + /// but the ordering on the borrowed form *must* match the + /// ordering on the value type. + #[unstable(feature = "set_recovery", issue = "28050")] + pub fn take(&mut self, value: &Q) -> Option where T: Borrow, Q: Ord { + Recover::take(&mut self.map, value) + } } #[stable(feature = "rust1", since = "1.0.0")] diff --git a/src/libcollectionstest/btree/set.rs b/src/libcollectionstest/btree/set.rs index b0d7390a49c..e7593bfcfe5 100644 --- a/src/libcollectionstest/btree/set.rs +++ b/src/libcollectionstest/btree/set.rs @@ -211,3 +211,52 @@ fn test_extend_ref() { assert!(a.contains(&5)); assert!(a.contains(&6)); } + +#[test] +fn test_recovery() { + use std::cmp::Ordering; + + #[derive(Debug)] + struct Foo(&'static str, i32); + + impl PartialEq for Foo { + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } + } + + impl Eq for Foo {} + + impl PartialOrd for Foo { + fn partial_cmp(&self, other: &Self) -> Option { + self.0.partial_cmp(&other.0) + } + } + + impl Ord for Foo { + fn cmp(&self, other: &Self) -> Ordering { + self.0.cmp(&other.0) + } + } + + let mut s = BTreeSet::new(); + assert_eq!(s.replace(Foo("a", 1)), None); + assert_eq!(s.len(), 1); + assert_eq!(s.replace(Foo("a", 2)), Some(Foo("a", 1))); + assert_eq!(s.len(), 1); + + { + let mut it = s.iter(); + assert_eq!(it.next(), Some(&Foo("a", 2))); + assert_eq!(it.next(), None); + } + + assert_eq!(s.get(&Foo("a", 1)), Some(&Foo("a", 2))); + assert_eq!(s.take(&Foo("a", 1)), Some(Foo("a", 2))); + assert_eq!(s.len(), 0); + + assert_eq!(s.get(&Foo("a", 1)), None); + assert_eq!(s.take(&Foo("a", 1)), None); + + assert_eq!(s.iter().next(), None); +} diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs index ca96f27d80d..5748c105eed 100644 --- a/src/libcollectionstest/lib.rs +++ b/src/libcollectionstest/lib.rs @@ -25,6 +25,7 @@ #![feature(rand)] #![feature(range_inclusive)] #![feature(rustc_private)] +#![feature(set_recovery)] #![feature(slice_bytes)] #![feature(slice_splits)] #![feature(split_off)] diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs index d5638bdac69..1e5c012e7d8 100644 --- a/src/libstd/collections/hash/map.rs +++ b/src/libstd/collections/hash/map.rs @@ -769,7 +769,7 @@ impl HashMap /// If the key already exists, the hashtable will be returned untouched /// and a reference to the existing element will be returned. fn insert_hashed_nocheck(&mut self, hash: SafeHash, k: K, v: V) -> &mut V { - self.insert_or_replace_with(hash, k, v, |_, _, _| ()) + self.insert_or_replace_with(hash, k, v, |_, _, _, _| ()) } fn insert_or_replace_with<'a, F>(&'a mut self, @@ -778,7 +778,7 @@ impl HashMap v: V, mut found_existing: F) -> &'a mut V where - F: FnMut(&mut K, &mut V, V), + F: FnMut(&mut K, &mut V, K, V), { // Worst case, we'll find one empty bucket among `size + 1` buckets. let size = self.table.size(); @@ -801,7 +801,7 @@ impl HashMap let (bucket_k, bucket_v) = bucket.into_mut_refs(); debug_assert!(k == *bucket_k); // Key already exists. Get its reference. - found_existing(bucket_k, bucket_v, v); + found_existing(bucket_k, bucket_v, k, v); return bucket_v; } } @@ -1123,7 +1123,7 @@ impl HashMap self.reserve(1); let mut retval = None; - self.insert_or_replace_with(hash, k, v, |_, val_ref, val| { + self.insert_or_replace_with(hash, k, v, |_, val_ref, _, val| { retval = Some(replace(val_ref, val)); }); retval @@ -1630,6 +1630,35 @@ impl Default for RandomState { } } +impl super::Recover for HashMap + where K: Eq + Hash + Borrow, S: HashState, Q: Eq + Hash +{ + type Key = K; + + fn get(&self, key: &Q) -> Option<&K> { + self.search(key).map(|bucket| bucket.into_refs().0) + } + + fn take(&mut self, key: &Q) -> Option { + if self.table.size() == 0 { + return None + } + + self.search_mut(key).map(|bucket| pop_internal(bucket).0) + } + + fn replace(&mut self, key: K) -> Option { + let hash = self.make_hash(&key); + self.reserve(1); + + let mut retkey = None; + self.insert_or_replace_with(hash, key, (), |key_ref, _, key, _| { + retkey = Some(replace(key_ref, key)); + }); + retkey + } +} + #[cfg(test)] mod test_map { use prelude::v1::*; diff --git a/src/libstd/collections/hash/mod.rs b/src/libstd/collections/hash/mod.rs index 47e300af269..4a6fcf44926 100644 --- a/src/libstd/collections/hash/mod.rs +++ b/src/libstd/collections/hash/mod.rs @@ -15,3 +15,11 @@ mod table; pub mod map; pub mod set; pub mod state; + +trait Recover { + type Key; + + fn get(&self, key: &Q) -> Option<&Self::Key>; + fn take(&mut self, key: &Q) -> Option; + fn replace(&mut self, key: Self::Key) -> Option; +} diff --git a/src/libstd/collections/hash/set.rs b/src/libstd/collections/hash/set.rs index ccad088a298..1f19a72371c 100644 --- a/src/libstd/collections/hash/set.rs +++ b/src/libstd/collections/hash/set.rs @@ -20,6 +20,7 @@ use iter::{Iterator, IntoIterator, ExactSizeIterator, FromIterator, Map, Chain, use ops::{BitOr, BitAnd, BitXor, Sub}; use option::Option::{Some, None, self}; +use super::Recover; use super::map::{self, HashMap, Keys, RandomState}; use super::state::HashState; @@ -459,6 +460,18 @@ impl HashSet self.map.contains_key(value) } + /// Returns a reference to the value in the set, if any, that is equal to the given value. + /// + /// The value may be any borrowed form of the set's value type, but + /// `Hash` and `Eq` on the borrowed form *must* match those for + /// the value type. + #[unstable(feature = "set_recovery", issue = "28050")] + pub fn get(&self, value: &Q) -> Option<&T> + where T: Borrow, Q: Hash + Eq + { + Recover::get(&self.map, value) + } + /// Returns `true` if the set has no elements in common with `other`. /// This is equivalent to checking for an empty intersection. /// @@ -544,6 +557,13 @@ impl HashSet #[stable(feature = "rust1", since = "1.0.0")] pub fn insert(&mut self, value: T) -> bool { self.map.insert(value, ()).is_none() } + /// Adds a value to the set, replacing the existing value, if any, that is equal to the given + /// one. Returns the replaced value. + #[unstable(feature = "set_recovery", issue = "28050")] + pub fn replace(&mut self, value: T) -> Option { + Recover::replace(&mut self.map, value) + } + /// Removes a value from the set. Returns `true` if the value was /// present in the set. /// @@ -568,6 +588,18 @@ impl HashSet { self.map.remove(value).is_some() } + + /// Removes and returns the value in the set, if any, that is equal to the given one. + /// + /// The value may be any borrowed form of the set's value type, but + /// `Hash` and `Eq` on the borrowed form *must* match those for + /// the value type. + #[unstable(feature = "set_recovery", issue = "28050")] + pub fn take(&mut self, value: &Q) -> Option + where T: Borrow, Q: Hash + Eq + { + Recover::take(&mut self.map, value) + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -1261,4 +1293,36 @@ mod test_set { s.extend(1..100); } } + + #[test] + fn test_replace() { + use hash; + + #[derive(Debug)] + struct Foo(&'static str, i32); + + impl PartialEq for Foo { + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } + } + + impl Eq for Foo {} + + impl hash::Hash for Foo { + fn hash(&self, h: &mut H) { + self.0.hash(h); + } + } + + let mut s = HashSet::new(); + assert_eq!(s.replace(Foo("a", 1)), None); + assert_eq!(s.len(), 1); + assert_eq!(s.replace(Foo("a", 2)), Some(Foo("a", 1))); + assert_eq!(s.len(), 1); + + let mut it = s.iter(); + assert_eq!(it.next(), Some(&Foo("a", 2))); + assert_eq!(it.next(), None); + } }