Add a get_mut method for TLS

Simulates borrow checks for '@mut' boxes, or at least it's the same idea.

Author: alexcrichton

src/libstd/local_data.rs

@@ -41,7 +41,7 @@ magic.
 
 use prelude::*;
 
-use task::local_data_priv::{local_get, local_pop, local_set, Handle};
+use task::local_data_priv::*;
 
 #[cfg(test)] use task;
 
@@ -95,6 +95,13 @@ pub fn get<T: 'static, U>(key: Key<@T>, f: &fn(Option<&@T>) -> U) -> U {
 pub fn get<T: 'static, U>(key: Key<T>, f: &fn(Option<&T>) -> U) -> U {
     unsafe { local_get(Handle::new(), key, f) }
 }
+/**
+ * Retrieve a mutable borrowed pointer to a task-local data value.
+ */
+#[cfg(not(stage0))]
+pub fn get_mut<T: 'static, U>(key: Key<T>, f: &fn(Option<&mut T>) -> U) -> U {
+    unsafe { local_get_mut(Handle::new(), key, f) }
+}
 /**
  * Store a value in task-local data. If this key already has a value,
  * that value is overwritten (and its destructor is run).
@@ -262,6 +269,34 @@ fn test_static_pointer() {
 fn test_owned() {
     static key: Key<~int> = &Key;
     set(key, ~1);
+
+    do get(key) |v| {
+        do get(key) |v| {
+            do get(key) |v| {
+                assert_eq!(**v.unwrap(), 1);
+            }
+            assert_eq!(**v.unwrap(), 1);
+        }
+        assert_eq!(**v.unwrap(), 1);
+    }
+    set(key, ~2);
+    do get(key) |v| {
+        assert_eq!(**v.unwrap(), 2);
+    }
+}
+
+#[test]
+fn test_get_mut() {
+    static key: Key<int> = &Key;
+    set(key, 1);
+
+    do get_mut(key) |v| {
+        *v.unwrap() = 2;
+    }
+
+    do get(key) |v| {
+        assert_eq!(*v.unwrap(), 2);
+    }
 }
 
 #[test]
@@ -283,3 +318,43 @@ fn test_same_key_type() {
     get(key4, |x| assert_eq!(*x.unwrap(), 4));
     get(key5, |x| assert_eq!(*x.unwrap(), 5));
 }
+
+#[test]
+#[should_fail]
+fn test_nested_get_set1() {
+    static key: Key<int> = &Key;
+    set(key, 4);
+    do get(key) |_| {
+        set(key, 4);
+    }
+}
+
+#[test]
+#[should_fail]
+fn test_nested_get_mut2() {
+    static key: Key<int> = &Key;
+    set(key, 4);
+    do get(key) |_| {
+        get_mut(key, |_| {})
+    }
+}
+
+#[test]
+#[should_fail]
+fn test_nested_get_mut3() {
+    static key: Key<int> = &Key;
+    set(key, 4);
+    do get_mut(key) |_| {
+        get(key, |_| {})
+    }
+}
+
+#[test]
+#[should_fail]
+fn test_nested_get_mut4() {
+    static key: Key<int> = &Key;
+    set(key, 4);
+    do get_mut(key) |_| {
+        get_mut(key, |_| {})
+    }
+}

src/libstd/task/local_data_priv.rs

@@ -44,6 +44,21 @@ impl Handle {
     }
 }
 
+#[deriving(Eq)]
+enum LoanState {
+    NoLoan, ImmLoan, MutLoan
+}
+
+impl LoanState {
+    fn describe(&self) -> &'static str {
+        match *self {
+            NoLoan => "no loan",
+            ImmLoan => "immutable",
+            MutLoan => "mutable"
+        }
+    }
+}
+
 trait LocalData {}
 impl<T: 'static> LocalData for T {}
 
@@ -77,7 +92,7 @@ impl<T: 'static> LocalData for T {}
 //
 // n.b. If TLS is used heavily in future, this could be made more efficient with
 //      a proper map.
-type TaskLocalMap = ~[Option<(*libc::c_void, TLSValue, uint)>];
+type TaskLocalMap = ~[Option<(*libc::c_void, TLSValue, LoanState)>];
 type TLSValue = ~LocalData:;
 
 fn cleanup_task_local_map(map_ptr: *libc::c_void) {
@@ -152,9 +167,10 @@ pub unsafe fn local_pop<T: 'static>(handle: Handle,
 
     for map.mut_iter().advance |entry| {
         match *entry {
-            Some((k, _, loans)) if k == key_value => {
-                if loans != 0 {
-                    fail!("TLS value has been loaned via get already");
+            Some((k, _, loan)) if k == key_value => {
+                if loan != NoLoan {
+                    fail!("TLS value cannot be removed because it is already \
+                          borrowed as %s", loan.describe());
                 }
                 // Move the data out of the `entry` slot via util::replace. This
                 // is guaranteed to succeed because we already matched on `Some`
@@ -192,6 +208,29 @@ pub unsafe fn local_pop<T: 'static>(handle: Handle,
 pub unsafe fn local_get<T: 'static, U>(handle: Handle,
                                        key: local_data::Key<T>,
                                        f: &fn(Option<&T>) -> U) -> U {
+    local_get_with(handle, key, ImmLoan, f)
+}
+
+pub unsafe fn local_get_mut<T: 'static, U>(handle: Handle,
+                                           key: local_data::Key<T>,
+                                           f: &fn(Option<&mut T>) -> U) -> U {
+    do local_get_with(handle, key, MutLoan) |x| {
+        match x {
+            None => f(None),
+            // We're violating a lot of compiler guarantees with this
+            // invocation of `transmute_mut`, but we're doing runtime checks to
+            // ensure that it's always valid (only one at a time).
+            //
+            // there is no need to be upset!
+            Some(x) => { f(Some(cast::transmute_mut(x))) }
+        }
+    }
+}
+
+unsafe fn local_get_with<T: 'static, U>(handle: Handle,
+                                        key: local_data::Key<T>,
+                                        state: LoanState,
+                                        f: &fn(Option<&T>) -> U) -> U {
     // This function must be extremely careful. Because TLS can store owned
     // values, and we must have some form of `get` function other than `pop`,
     // this function has to give a `&` reference back to the caller.
@@ -218,12 +257,24 @@ pub unsafe fn local_get<T: 'static, U>(handle: Handle,
         None => { return f(None); }
         Some(i) => {
             let ret;
+            let mut return_loan = false;
             match map[i] {
-                Some((_, ref data, ref mut loans)) => {
-                    *loans = *loans + 1;
+                Some((_, ref data, ref mut loan)) => {
+                    match (state, *loan) {
+                        (_, NoLoan) => {
+                            *loan = state;
+                            return_loan = true;
+                        }
+                        (ImmLoan, ImmLoan) => {}
+                        (want, cur) => {
+                            fail!("TLS slot cannot be borrowed as %s because \
+                                   it is already borrowed as %s",
+                                  want.describe(), cur.describe());
+                        }
+                    }
                     // data was created with `~~T as ~LocalData`, so we extract
                     // pointer part of the trait, (as ~~T), and then use
-                    // compiler coercions to achieve a '&' pointer
+                    // compiler coercions to achieve a '&' pointer.
                     match *cast::transmute::<&TLSValue, &(uint, ~~T)>(data) {
                         (_vtable, ref box) => {
                             let value: &T = **box;
@@ -238,9 +289,11 @@ pub unsafe fn local_get<T: 'static, U>(handle: Handle,
             // 'f' returned because `f` could have appended more TLS items which
             // in turn relocated the vector. Hence we do another lookup here to
             // fixup the loans.
-            match map[i] {
-                Some((_, _, ref mut loans)) => { *loans = *loans - 1; }
-                None => { libc::abort(); }
+            if return_loan {
+                match map[i] {
+                    Some((_, _, ref mut loan)) => { *loan = NoLoan; }
+                    None => { libc::abort(); }
+                }
             }
             return ret;
         }
@@ -269,9 +322,10 @@ pub unsafe fn local_set<T: 'static>(handle: Handle,
         // First see if the map contains this key already
         let curspot = map.iter().position(|entry| {
             match *entry {
-                Some((ekey, _, loans)) if key == ekey => {
-                    if loans != 0 {
-                        fail!("TLS value has been loaned via get already");
+                Some((ekey, _, loan)) if key == ekey => {
+                    if loan != NoLoan {
+                        fail!("TLS value cannot be overwritten because it is
+                               already borrowed as %s", loan.describe())
                     }
                     true
                 }
@@ -286,7 +340,7 @@ pub unsafe fn local_set<T: 'static>(handle: Handle,
     }
 
     match insertion_position(map, keyval) {
-        Some(i) => { map[i] = Some((keyval, data, 0)); }
-        None => { map.push(Some((keyval, data, 0))); }
+        Some(i) => { map[i] = Some((keyval, data, NoLoan)); }
+        None => { map.push(Some((keyval, data, NoLoan))); }
     }
 }