auto merge of #7677 : alexcrichton/rust/tls-gc, r=pcwalton

cc #6004 and #3273

This is a rewrite of TLS to get towards not requiring `@` when using task local storage. Most of the rewrite is straightforward, although there are two caveats:

1. Changing `local_set` to not require `@` is blocked on #7673
2. The code in `local_pop` is some of the most unsafe code I've written. A second set of eyes should definitely scrutinize it...

The public-facing interface currently hasn't changed, although it will have to change because `local_data::get` cannot return `Option<T>`, nor can it return `Option<&T>` (the lifetime isn't known). This will have to be changed to be given a closure which yield `&T` (or as an Option). I didn't do this part of the api rewrite in this pull request as I figured that it could wait until when `@` is fully removed.

This also doesn't deal with the issue of using something other than functions as keys, but I'm looking into using static slices (as mentioned in the issues).

5 files changed, 462 insertions, 170 deletions

diff --git a/src/libstd/task/local_data_priv.rs b/src/libstd/task/local_data_priv.rs
index 8dd96df4545..42cfcbc16db 100644
--- a/src/libstd/task/local_data_priv.rs
+++ b/src/libstd/task/local_data_priv.rs
@@ -11,11 +11,13 @@
 #[allow(missing_doc)];
 
 use cast;
-use cmp::Eq;
 use libc;
+use local_data;
 use prelude::*;
+use ptr;
+use sys;
 use task::rt;
-use local_data::LocalDataKey;
+use util;
 
 use super::rt::rust_task;
 use rt::task::{Task, LocalStorage};
@@ -43,25 +45,41 @@ impl Handle {
     }
 }
 
-pub trait LocalData { }
-impl<T: 'static> LocalData for @T { }
+trait LocalData {}
+impl<T: 'static> LocalData for T {}
 
-impl Eq for @LocalData {
-    fn eq(&self, other: &@LocalData) -> bool {
-        unsafe {
-            let ptr_a: &(uint, uint) = cast::transmute(self);
-            let ptr_b: &(uint, uint) = cast::transmute(other);
-            return ptr_a == ptr_b;
-        }
-    }
-    fn ne(&self, other: &@LocalData) -> bool { !(*self).eq(other) }
-}
-
-// If TLS is used heavily in future, this could be made more efficient with a
-// proper map.
-type TaskLocalElement = (*libc::c_void, *libc::c_void, @LocalData);
-// Has to be a pointer at outermost layer; the foreign call returns void *.
-type TaskLocalMap = @mut ~[Option<TaskLocalElement>];
+// The task-local-map stores all TLS information for the currently running task.
+// It is stored as an owned pointer into the runtime, and it's only allocated
+// when TLS is used for the first time. This map must be very carefully
+// constructed because it has many mutable loans unsoundly handed out on it to
+// the various invocations of TLS requests.
+//
+// One of the most important operations is loaning a value via `get` to a
+// caller. In doing so, the slot that the TLS entry is occupying cannot be
+// invalidated because upon returning it's loan state must be updated. Currently
+// the TLS map is a vector, but this is possibly dangerous because the vector
+// can be reallocated/moved when new values are pushed onto it.
+//
+// This problem currently isn't solved in a very elegant way. Inside the `get`
+// function, it internally "invalidates" all references after the loan is
+// finished and looks up into the vector again. In theory this will prevent
+// pointers from being moved under our feet so long as LLVM doesn't go too crazy
+// with the optimizations.
+//
+// n.b. Other structures are not sufficient right now:
+//          * HashMap uses ~[T] internally (push reallocates/moves)
+//          * TreeMap is plausible, but it's in extra
+//          * dlist plausible, but not in std
+//          * a custom owned linked list was attempted, but difficult to write
+//            and involved a lot of extra code bloat
+//
+// n.b. Has to be stored with a pointer at outermost layer; the foreign call
+//      returns void *.
+//
+// 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 TLSValue = ~LocalData:;
 
 fn cleanup_task_local_map(map_ptr: *libc::c_void) {
     unsafe {
@@ -74,162 +92,203 @@ fn cleanup_task_local_map(map_ptr: *libc::c_void) {
 }
 
 // Gets the map from the runtime. Lazily initialises if not done so already.
-unsafe fn get_local_map(handle: Handle) -> TaskLocalMap {
-    match handle {
-        OldHandle(task) => get_task_local_map(task),
-        NewHandle(local_storage) => get_newsched_local_map(local_storage)
-    }
-}
+unsafe fn get_local_map(handle: Handle) -> &mut TaskLocalMap {
 
-unsafe fn get_task_local_map(task: *rust_task) -> TaskLocalMap {
+    unsafe fn oldsched_map(task: *rust_task) -> &mut TaskLocalMap {
+        extern fn cleanup_extern_cb(map_ptr: *libc::c_void) {
+            cleanup_task_local_map(map_ptr);
+        }
 
-    extern fn cleanup_task_local_map_extern_cb(map_ptr: *libc::c_void) {
-        cleanup_task_local_map(map_ptr);
+        // Relies on the runtime initialising the pointer to null.
+        // Note: the map is an owned pointer and is "owned" by TLS. It is moved
+        // into the tls slot for this task, and then mutable loans are taken
+        // from this slot to modify the map.
+        let map_ptr = rt::rust_get_task_local_data(task);
+        if (*map_ptr).is_null() {
+            // First time TLS is used, create a new map and set up the necessary
+            // TLS information for its safe destruction
+            let map: TaskLocalMap = ~[];
+            *map_ptr = cast::transmute(map);
+            rt::rust_task_local_data_atexit(task, cleanup_extern_cb);
+        }
+        return cast::transmute(map_ptr);
     }
 
-    // Relies on the runtime initialising the pointer to null.
-    // Note: The map's box lives in TLS invisibly referenced once. Each time
-    // we retrieve it for get/set, we make another reference, which get/set
-    // drop when they finish. No "re-storing after modifying" is needed.
-    let map_ptr = rt::rust_get_task_local_data(task);
-    if map_ptr.is_null() {
-        let map: TaskLocalMap = @mut ~[];
-        // NB: This bumps the ref count before converting to an unsafe pointer,
-        // keeping the map alive until TLS is destroyed
-        rt::rust_set_task_local_data(task, cast::transmute(map));
-        rt::rust_task_local_data_atexit(task, cleanup_task_local_map_extern_cb);
-        map
-    } else {
-        let map = cast::transmute(map_ptr);
-        let nonmut = cast::transmute::<TaskLocalMap,
-                                       @~[Option<TaskLocalElement>]>(map);
-        cast::bump_box_refcount(nonmut);
-        map
+    unsafe fn newsched_map(local: *mut LocalStorage) -> &mut TaskLocalMap {
+        // This is based on the same idea as the oldsched code above.
+        match &mut *local {
+            // If the at_exit function is already set, then we just need to take
+            // a loan out on the TLS map stored inside
+            &LocalStorage(ref mut map_ptr, Some(_)) => {
+                assert!(map_ptr.is_not_null());
+                return cast::transmute(map_ptr);
+            }
+            // If this is the first time we've accessed TLS, perform similar
+            // actions to the oldsched way of doing things.
+            &LocalStorage(ref mut map_ptr, ref mut at_exit) => {
+                assert!(map_ptr.is_null());
+                assert!(at_exit.is_none());
+                let map: TaskLocalMap = ~[];
+                *map_ptr = cast::transmute(map);
+                *at_exit = Some(cleanup_task_local_map);
+                return cast::transmute(map_ptr);
+            }
+        }
     }
-}
 
-unsafe fn get_newsched_local_map(local: *mut LocalStorage) -> TaskLocalMap {
-    match &mut *local {
-        &LocalStorage(map_ptr, Some(_)) => {
-            assert!(map_ptr.is_not_null());
-            let map = cast::transmute(map_ptr);
-            let nonmut = cast::transmute::<TaskLocalMap,
-            @~[Option<TaskLocalElement>]>(map);
-            cast::bump_box_refcount(nonmut);
-            return map;
-        }
-        &LocalStorage(ref mut map_ptr, ref mut at_exit) => {
-            assert!((*map_ptr).is_null());
-            let map: TaskLocalMap = @mut ~[];
-            *map_ptr = cast::transmute(map);
-            let at_exit_fn: ~fn(*libc::c_void) = |p|cleanup_task_local_map(p);
-            *at_exit = Some(at_exit_fn);
-            return map;
-        }
+    match handle {
+        OldHandle(task) => oldsched_map(task),
+        NewHandle(local_storage) => newsched_map(local_storage)
     }
 }
 
-unsafe fn key_to_key_value<T: 'static>(key: LocalDataKey<T>) -> *libc::c_void {
-    // Keys are closures, which are (fnptr,envptr) pairs. Use fnptr.
-    // Use reinterpret_cast -- transmute would leak (forget) the closure.
-    let pair: (*libc::c_void, *libc::c_void) = cast::transmute_copy(&key);
-    pair.first()
+unsafe fn key_to_key_value<T: 'static>(key: local_data::Key<T>) -> *libc::c_void {
+    let pair: sys::Closure = cast::transmute(key);
+    return pair.code as *libc::c_void;
 }
 
-// If returning Some(..), returns with @T with the map's reference. Careful!
-unsafe fn local_data_lookup<T: 'static>(
-    map: TaskLocalMap, key: LocalDataKey<T>)
-    -> Option<(uint, *libc::c_void)> {
-
+pub unsafe fn local_pop<T: 'static>(handle: Handle,
+                                    key: local_data::Key<T>) -> Option<T> {
+    let map = get_local_map(handle);
     let key_value = key_to_key_value(key);
-    let map_pos = (*map).iter().position(|entry|
+
+    for map.mut_iter().advance |entry| {
         match *entry {
-            Some((k,_,_)) => k == key_value,
-            None => false
-        }
-    );
-    do map_pos.map |index| {
-        // .get() is guaranteed because of "None { false }" above.
-        let (_, data_ptr, _) = (*map)[*index].get();
-        (*index, data_ptr)
-    }
-}
+            Some((k, _, loans)) if k == key_value => {
+                if loans != 0 {
+                    fail!("TLS value has been loaned via get already");
+                }
+                // Move the data out of the `entry` slot via util::replace. This
+                // is guaranteed to succeed because we already matched on `Some`
+                // above.
+                let data = match util::replace(entry, None) {
+                    Some((_, data, _)) => data,
+                    None => libc::abort(),
+                };
 
-unsafe fn local_get_helper<T: 'static>(
-    handle: Handle, key: LocalDataKey<T>,
-    do_pop: bool) -> Option<@T> {
+                // Move `data` into transmute to get out the memory that it
+                // owns, we must free it manually later.
+                let (_vtable, box): (uint, ~~T) = cast::transmute(data);
 
-    let map = get_local_map(handle);
-    // Interpreturn our findings from the map
-    do local_data_lookup(map, key).map |result| {
-        // A reference count magically appears on 'data' out of thin air. It
-        // was referenced in the local_data box, though, not here, so before
-        // overwriting the local_data_box we need to give an extra reference.
-        // We must also give an extra reference when not removing.
-        let (index, data_ptr) = *result;
-        let data: @T = cast::transmute(data_ptr);
-        cast::bump_box_refcount(data);
-        if do_pop {
-            map[index] = None;
+                // Read the box's value (using the compiler's built-in
+                // auto-deref functionality to obtain a pointer to the base)
+                let ret = ptr::read_ptr(cast::transmute::<&T, *mut T>(*box));
+
+                // Finally free the allocated memory. we don't want this to
+                // actually touch the memory inside because it's all duplicated
+                // now, so the box is transmuted to a 0-sized type. We also use
+                // a type which references `T` because currently the layout
+                // could depend on whether T contains managed pointers or not.
+                let _: ~~[T, ..0] = cast::transmute(box);
+
+                // Everything is now deallocated, and we own the value that was
+                // located inside TLS, so we now return it.
+                return Some(ret);
+            }
+            _ => {}
         }
-        data
     }
+    return None;
 }
 
+pub unsafe fn local_get<T: 'static, U>(handle: Handle,
+                                       key: local_data::Key<T>,
+                                       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.
+    //
+    // One option is to return the reference, but this cannot be sound because
+    // the actual lifetime of the object is not known. The slot in TLS could not
+    // be modified until the object goes out of scope, but the TLS code cannot
+    // know when this happens.
+    //
+    // For this reason, the reference is yielded to a specified closure. This
+    // way the TLS code knows exactly what the lifetime of the yielded pointer
+    // is, allowing callers to acquire references to owned data. This is also
+    // sound so long as measures are taken to ensure that while a TLS slot is
+    // loaned out to a caller, it's not modified recursively.
+    let map = get_local_map(handle);
+    let key_value = key_to_key_value(key);
 
-pub unsafe fn local_pop<T: 'static>(
-    handle: Handle,
-    key: LocalDataKey<T>) -> Option<@T> {
-
-    local_get_helper(handle, key, true)
-}
-
-pub unsafe fn local_get<T: 'static>(
-    handle: Handle,
-    key: LocalDataKey<T>) -> Option<@T> {
+    let pos = map.iter().position(|entry| {
+        match *entry {
+            Some((k, _, _)) if k == key_value => true, _ => false
+        }
+    });
+    match pos {
+        None => { return f(None); }
+        Some(i) => {
+            let ret;
+            match map[i] {
+                Some((_, ref data, ref mut loans)) => {
+                    *loans = *loans + 1;
+                    // 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
+                    match *cast::transmute::<&TLSValue, &(uint, ~~T)>(data) {
+                        (_vtable, ref box) => {
+                            let value: &T = **box;
+                            ret = f(Some(value));
+                        }
+                    }
+                }
+                _ => libc::abort()
+            }
 
-    local_get_helper(handle, key, false)
+            // n.b. 'data' and 'loans' are both invalid pointers at the point
+            // '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(); }
+            }
+            return ret;
+        }
+    }
 }
 
-pub unsafe fn local_set<T: 'static>(
-    handle: Handle, key: LocalDataKey<T>, data: @T) {
-
+pub unsafe fn local_set<T: 'static>(handle: Handle,
+                                    key: local_data::Key<T>,
+                                    data: T) {
     let map = get_local_map(handle);
-    // Store key+data as *voids. Data is invisibly referenced once; key isn't.
     let keyval = key_to_key_value(key);
-    // We keep the data in two forms: one as an unsafe pointer, so we can get
-    // it back by casting; another in an existential box, so the reference we
-    // own on it can be dropped when the box is destroyed. The unsafe pointer
-    // does not have a reference associated with it, so it may become invalid
-    // when the box is destroyed.
-    let data_ptr = *cast::transmute::<&@T, &*libc::c_void>(&data);
-    let data_box = @data as @LocalData;
-    // Construct new entry to store in the map.
-    let new_entry = Some((keyval, data_ptr, data_box));
-    // Find a place to put it.
-    match local_data_lookup(map, key) {
-        Some((index, _old_data_ptr)) => {
-            // Key already had a value set, _old_data_ptr, whose reference
-            // will get dropped when the local_data box is overwritten.
-            map[index] = new_entry;
-        }
-        None => {
-            // Find an empty slot. If not, grow the vector.
-            match (*map).iter().position(|x| x.is_none()) {
-                Some(empty_index) => { map[empty_index] = new_entry; }
-                None => { map.push(new_entry); }
+
+    // When the task-local map is destroyed, all the data needs to be cleaned
+    // up. For this reason we can't do some clever tricks to store '~T' as a
+    // '*c_void' or something like that. To solve the problem, we cast
+    // everything to a trait (LocalData) which is then stored inside the map.
+    // Upon destruction of the map, all the objects will be destroyed and the
+    // traits have enough information about them to destroy themselves.
+    //
+    // FIXME(#7673): This should be "~data as ~LocalData" (without the colon at
+    //               the end, and only one sigil)
+    let data = ~~data as ~LocalData:;
+
+    fn insertion_position(map: &mut TaskLocalMap,
+                          key: *libc::c_void) -> Option<uint> {
+        // 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");
+                    }
+                    true
+                }
+                _ => false,
             }
+        });
+        // If it doesn't contain the key, just find a slot that's None
+        match curspot {
+            Some(i) => Some(i),
+            None => map.iter().position(|entry| entry.is_none())
         }
     }
-}
-
-pub unsafe fn local_modify<T: 'static>(
-    handle: Handle, key: LocalDataKey<T>,
-    modify_fn: &fn(Option<@T>) -> Option<@T>) {
 
-    // Could be more efficient by doing the lookup work, but this is easy.
-    let newdata = modify_fn(local_pop(handle, key));
-    if newdata.is_some() {
-        local_set(handle, key, newdata.unwrap());
+    match insertion_position(map, keyval) {
+        Some(i) => { map[i] = Some((keyval, data, 0)); }
+        None => { map.push(Some((keyval, data, 0))); }
     }
 }
diff --git a/src/libstd/task/local_data_priv_stage0.rs b/src/libstd/task/local_data_priv_stage0.rs
new file mode 100644
index 00000000000..fe80ec06c69
--- /dev/null
+++ b/src/libstd/task/local_data_priv_stage0.rs
@@ -0,0 +1,229 @@
+// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#[allow(missing_doc)];
+
+use cast;
+use cmp::Eq;
+use libc;
+use local_data;
+use prelude::*;
+use sys;
+use task::rt;
+
+use super::rt::rust_task;
+use rt::task::{Task, LocalStorage};
+
+pub enum Handle {
+    OldHandle(*rust_task),
+    NewHandle(*mut LocalStorage)
+}
+
+impl Handle {
+    pub fn new() -> Handle {
+        use rt::{context, OldTaskContext};
+        use rt::local::Local;
+        unsafe {
+            match context() {
+                OldTaskContext => {
+                    OldHandle(rt::rust_get_task())
+                }
+                _ => {
+                    let task = Local::unsafe_borrow::<Task>();
+                    NewHandle(&mut (*task).storage)
+                }
+            }
+        }
+    }
+}
+
+pub trait LocalData { }
+impl<T: 'static> LocalData for @T { }
+
+impl Eq for @LocalData {
+    fn eq(&self, other: &@LocalData) -> bool {
+        unsafe {
+            let ptr_a: &(uint, uint) = cast::transmute(self);
+            let ptr_b: &(uint, uint) = cast::transmute(other);
+            return ptr_a == ptr_b;
+        }
+    }
+    fn ne(&self, other: &@LocalData) -> bool { !(*self).eq(other) }
+}
+
+// If TLS is used heavily in future, this could be made more efficient with a
+// proper map.
+type TaskLocalElement = (*libc::c_void, *libc::c_void, @LocalData);
+// Has to be a pointer at outermost layer; the foreign call returns void *.
+type TaskLocalMap = ~[Option<TaskLocalElement>];
+
+fn cleanup_task_local_map(map_ptr: *libc::c_void) {
+    unsafe {
+        assert!(!map_ptr.is_null());
+        // Get and keep the single reference that was created at the
+        // beginning.
+        let _map: TaskLocalMap = cast::transmute(map_ptr);
+        // All local_data will be destroyed along with the map.
+    }
+}
+
+// Gets the map from the runtime. Lazily initialises if not done so already.
+unsafe fn get_local_map(handle: Handle) -> &mut TaskLocalMap {
+    match handle {
+        OldHandle(task) => get_task_local_map(task),
+        NewHandle(local_storage) => get_newsched_local_map(local_storage)
+    }
+}
+
+unsafe fn get_task_local_map(task: *rust_task) -> &mut TaskLocalMap {
+
+    extern fn cleanup_task_local_map_extern_cb(map_ptr: *libc::c_void) {
+        cleanup_task_local_map(map_ptr);
+    }
+
+    // Relies on the runtime initialising the pointer to null.
+    // Note: the map is an owned pointer and is "owned" by TLS. It is moved
+    // into the tls slot for this task, and then mutable loans are taken from
+    // this slot to modify the map.
+    let map_ptr = rt::rust_get_task_local_data(task);
+    if (*map_ptr).is_null() {
+        // First time TLS is used, create a new map and set up the necessary
+        // TLS information for its safe destruction
+        let map: TaskLocalMap = ~[];
+        *map_ptr = cast::transmute(map);
+        rt::rust_task_local_data_atexit(task, cleanup_task_local_map_extern_cb);
+    }
+    return cast::transmute(map_ptr);
+}
+
+unsafe fn get_newsched_local_map(local: *mut LocalStorage) -> &mut TaskLocalMap {
+    // This is based on the same idea as the oldsched code above.
+    match &mut *local {
+        // If the at_exit function is already set, then we just need to take a
+        // loan out on the TLS map stored inside
+        &LocalStorage(ref mut map_ptr, Some(_)) => {
+            assert!(map_ptr.is_not_null());
+            return cast::transmute(map_ptr);
+        }
+        // If this is the first time we've accessed TLS, perform similar
+        // actions to the oldsched way of doing things.
+        &LocalStorage(ref mut map_ptr, ref mut at_exit) => {
+            assert!(map_ptr.is_null());
+            assert!(at_exit.is_none());
+            let map: TaskLocalMap = ~[];
+            *map_ptr = cast::transmute(map);
+            *at_exit = Some(cleanup_task_local_map);
+            return cast::transmute(map_ptr);
+        }
+    }
+}
+
+unsafe fn key_to_key_value<T: 'static>(key: local_data::Key<@T>) -> *libc::c_void {
+    let pair: sys::Closure = cast::transmute(key);
+    return pair.code as *libc::c_void;
+}
+
+// If returning Some(..), returns with @T with the map's reference. Careful!
+unsafe fn local_data_lookup<T: 'static>(
+    map: &mut TaskLocalMap, key: local_data::Key<@T>)
+    -> Option<(uint, *libc::c_void)> {
+
+    let key_value = key_to_key_value(key);
+    for map.iter().enumerate().advance |(i, entry)| {
+        match *entry {
+            Some((k, data, _)) if k == key_value => { return Some((i, data)); }
+            _ => {}
+        }
+    }
+    return None;
+}
+
+unsafe fn local_get_helper<T: 'static>(
+    handle: Handle, key: local_data::Key<@T>,
+    do_pop: bool) -> Option<@T> {
+
+    let map = get_local_map(handle);
+    // Interpreturn our findings from the map
+    do local_data_lookup(map, key).map |result| {
+        // A reference count magically appears on 'data' out of thin air. It
+        // was referenced in the local_data box, though, not here, so before
+        // overwriting the local_data_box we need to give an extra reference.
+        // We must also give an extra reference when not removing.
+        let (index, data_ptr) = *result;
+        let data: @T = cast::transmute(data_ptr);
+        cast::bump_box_refcount(data);
+        if do_pop {
+            map[index] = None;
+        }
+        data
+    }
+}
+
+
+pub unsafe fn local_pop<T: 'static>(
+    handle: Handle,
+    key: local_data::Key<@T>) -> Option<@T> {
+
+    local_get_helper(handle, key, true)
+}
+
+pub unsafe fn local_get<T: 'static, U>(
+    handle: Handle,
+    key: local_data::Key<@T>,
+    f: &fn(Option<&@T>) -> U) -> U {
+
+    match local_get_helper(handle, key, false) {
+        Some(ref x) => f(Some(x)),
+        None => f(None)
+    }
+}
+
+pub unsafe fn local_set<T: 'static>(
+    handle: Handle, key: local_data::Key<@T>, data: @T) {
+
+    let map = get_local_map(handle);
+    // Store key+data as *voids. Data is invisibly referenced once; key isn't.
+    let keyval = key_to_key_value(key);
+    // We keep the data in two forms: one as an unsafe pointer, so we can get
+    // it back by casting; another in an existential box, so the reference we
+    // own on it can be dropped when the box is destroyed. The unsafe pointer
+    // does not have a reference associated with it, so it may become invalid
+    // when the box is destroyed.
+    let data_ptr = *cast::transmute::<&@T, &*libc::c_void>(&data);
+    let data_box = @data as @LocalData;
+    // Construct new entry to store in the map.
+    let new_entry = Some((keyval, data_ptr, data_box));
+    // Find a place to put it.
+    match local_data_lookup(map, key) {
+        Some((index, _old_data_ptr)) => {
+            // Key already had a value set, _old_data_ptr, whose reference
+            // will get dropped when the local_data box is overwritten.
+            map[index] = new_entry;
+        }
+        None => {
+            // Find an empty slot. If not, grow the vector.
+            match map.iter().position(|x| x.is_none()) {
+                Some(empty_index) => { map[empty_index] = new_entry; }
+                None => { map.push(new_entry); }
+            }
+        }
+    }
+}
+
+pub unsafe fn local_modify<T: 'static>(
+    handle: Handle, key: local_data::Key<@T>,
+    modify_fn: &fn(Option<@T>) -> Option<@T>) {
+
+    // Could be more efficient by doing the lookup work, but this is easy.
+    let newdata = modify_fn(local_pop(handle, key));
+    if newdata.is_some() {
+        local_set(handle, key, newdata.unwrap());
+    }
+}
diff --git a/src/libstd/task/mod.rs b/src/libstd/task/mod.rs
index 5a3ff10ae83..b14100991df 100644
--- a/src/libstd/task/mod.rs
+++ b/src/libstd/task/mod.rs
@@ -54,6 +54,10 @@ use util;
 #[cfg(test)] use ptr;
 #[cfg(test)] use task;
 
+#[cfg(stage0)]
+#[path="local_data_priv_stage0.rs"]
+mod local_data_priv;
+#[cfg(not(stage0))]
 mod local_data_priv;
 pub mod rt;
 pub mod spawn;
diff --git a/src/libstd/task/rt.rs b/src/libstd/task/rt.rs
index 4860ab36f77..76fcad0759a 100644
--- a/src/libstd/task/rt.rs
+++ b/src/libstd/task/rt.rs
@@ -63,9 +63,7 @@ pub extern {
     fn rust_task_kill_all(task: *rust_task);
 
     #[rust_stack]
-    fn rust_get_task_local_data(task: *rust_task) -> *libc::c_void;
-    #[rust_stack]
-    fn rust_set_task_local_data(task: *rust_task, map: *libc::c_void);
+    fn rust_get_task_local_data(task: *rust_task) -> *mut *libc::c_void;
     #[rust_stack]
     fn rust_task_local_data_atexit(task: *rust_task, cleanup_fn: *u8);
 }
diff --git a/src/libstd/task/spawn.rs b/src/libstd/task/spawn.rs
index bcb7e06bf1f..f45d470a9f6 100644
--- a/src/libstd/task/spawn.rs
+++ b/src/libstd/task/spawn.rs
@@ -478,26 +478,28 @@ fn gen_child_taskgroup(linked: bool, supervised: bool)
          * Step 1. Get spawner's taskgroup info.
          *##################################################################*/
         let spawner_group: @@mut TCB =
-            match local_get(OldHandle(spawner), taskgroup_key!()) {
-                None => {
-                    // Main task, doing first spawn ever. Lazily initialise
-                    // here.
-                    let mut members = new_taskset();
-                    taskset_insert(&mut members, spawner);
-                    let tasks = exclusive(Some(TaskGroupData {
-                        members: members,
-                        descendants: new_taskset(),
-                    }));
-                    // Main task/group has no ancestors, no notifier, etc.
-                    let group = @@mut TCB(spawner,
-                                          tasks,
-                                          AncestorList(None),
-                                          true,
-                                          None);
-                    local_set(OldHandle(spawner), taskgroup_key!(), group);
-                    group
+            do local_get(OldHandle(spawner), taskgroup_key!()) |group| {
+                match group {
+                    None => {
+                        // Main task, doing first spawn ever. Lazily initialise
+                        // here.
+                        let mut members = new_taskset();
+                        taskset_insert(&mut members, spawner);
+                        let tasks = exclusive(Some(TaskGroupData {
+                            members: members,
+                            descendants: new_taskset(),
+                        }));
+                        // Main task/group has no ancestors, no notifier, etc.
+                        let group = @@mut TCB(spawner,
+                                              tasks,
+                                              AncestorList(None),
+                                              true,
+                                              None);
+                        local_set(OldHandle(spawner), taskgroup_key!(), group);
+                        group
+                    }
+                    Some(&group) => group
                 }
-                Some(group) => group
             };
         let spawner_group: &mut TCB = *spawner_group;