From 8299f3a447ca6416b23c20871172d784a28762bd Mon Sep 17 00:00:00 2001
From: Erick Tryzelaar <erick.tryzelaar@gmail.com>
Date: Fri, 7 Sep 2012 10:41:07 -0700
Subject: [PATCH] libcore: add send_map::with_find_ref and with_get_ref.

---
 src/libcore/send_map.rs | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/src/libcore/send_map.rs b/src/libcore/send_map.rs
index e9eff6db31a2..49a376bb2731 100644
--- a/src/libcore/send_map.rs
+++ b/src/libcore/send_map.rs
@@ -22,6 +22,8 @@ trait SendMap<K:Eq Hash, V: Copy> {
     fn each_value_ref(&self, blk: fn(v: &V) -> bool);
     fn find(&const self, k: &K) -> Option<V>;
     fn get(&const self, k: &K) -> V;
+    fn with_find_ref<T>(&const self, k: &K, blk: fn(Option<&V>) -> T) -> T;
+    fn with_get_ref<T>(&const self, k: &K, blk: fn(v: &V) -> T) -> T;
 }
 
 /// Open addressing with linear probing.
@@ -290,6 +292,27 @@ mod linear {
         }
         */
 
+        fn with_find_ref<T>(&self, k: &K, blk: fn(Option<&V>) -> T) -> T {
+            match self.bucket_for_key(self.buckets, k) {
+                FoundEntry(idx) => {
+                    match self.buckets[idx] {
+                        Some(bkt) => blk(Some(&bkt.value)),
+                        None => fail ~"LinearMap::find: internal logic error"
+                    }
+                }
+                TableFull | FoundHole(_) => blk(None),
+            }
+        }
+
+        fn with_get_ref<T>(&self, k: &K, blk: fn(v: &V) -> T) -> T {
+            do self.with_find_ref(k) |v| {
+                match v {
+                    Some(v) => blk(v),
+                    None => fail fmt!("No entry found for key: %?", k),
+                }
+            }
+        }
+
         fn each_ref(&self, blk: fn(k: &K, v: &V) -> bool) {
             for vec::each(self.buckets) |slot| {
                 let mut broke = false;
@@ -426,4 +449,12 @@ mod test {
         }
         assert observed == 0xFFFF_FFFF;
     }
+
+    #[test]
+    fn with_find_ref() {
+        let mut m = ~LinearMap();
+        m.with_find_ref(&1, |v| assert v.is_none());
+        m.insert(1, 2);
+        m.with_find_ref(&1, |v| assert *v.get() == 2);
+    }
 }