Work around stage0 to remove '@' requirements from TLS

5 files changed, 334 insertions, 75 deletions

diff --git a/src/libstd/local_data.rs b/src/libstd/local_data.rs
index d7ab7236b5c..a117d461cfd 100644
--- a/src/libstd/local_data.rs
+++ b/src/libstd/local_data.rs
@@ -66,6 +66,15 @@ pub type Key<'self,T> = &'self fn:Copy(v: T);
  * Remove a task-local data value from the table, returning the
  * reference that was originally created to insert it.
  */
+#[cfg(stage0)]
+pub unsafe fn pop<T: 'static>(key: Key<@T>) -> Option<@T> {
+    local_pop(Handle::new(), key)
+}
+/**
+ * Remove a task-local data value from the table, returning the
+ * reference that was originally created to insert it.
+ */
+#[cfg(not(stage0))]
 pub unsafe fn pop<T: 'static>(key: Key<T>) -> Option<T> {
     local_pop(Handle::new(), key)
 }
@@ -73,6 +82,15 @@ pub unsafe fn pop<T: 'static>(key: Key<T>) -> Option<T> {
  * Retrieve a task-local data value. It will also be kept alive in the
  * table until explicitly removed.
  */
+#[cfg(stage0)]
+pub unsafe fn get<T: 'static>(key: Key<@T>) -> Option<@T> {
+    local_get(Handle::new(), key, |loc| loc.map(|&x| *x))
+}
+/**
+ * Retrieve a task-local data value. It will also be kept alive in the
+ * table until explicitly removed.
+ */
+#[cfg(not(stage0))]
 pub unsafe fn get<T: 'static>(key: Key<@T>) -> Option<@T> {
     local_get(Handle::new(), key, |loc| loc.map(|&x| *x))
 }
@@ -80,13 +98,23 @@ pub unsafe fn get<T: 'static>(key: Key<@T>) -> Option<@T> {
  * Store a value in task-local data. If this key already has a value,
  * that value is overwritten (and its destructor is run).
  */
+#[cfg(stage0)]
 pub unsafe fn set<T: 'static>(key: Key<@T>, data: @T) {
     local_set(Handle::new(), key, data)
 }
 /**
+ * Store a value in task-local data. If this key already has a value,
+ * that value is overwritten (and its destructor is run).
+ */
+#[cfg(not(stage0))]
+pub unsafe fn set<T: 'static>(key: Key<T>, data: T) {
+    local_set(Handle::new(), key, data)
+}
+/**
  * Modify a task-local data value. If the function returns 'None', the
  * data is removed (and its reference dropped).
  */
+#[cfg(stage0)]
 pub unsafe fn modify<T: 'static>(key: Key<@T>,
                                  f: &fn(Option<@T>) -> Option<@T>) {
     match f(pop(key)) {
@@ -94,6 +122,18 @@ pub unsafe fn modify<T: 'static>(key: Key<@T>,
         None => {}
     }
 }
+/**
+ * Modify a task-local data value. If the function returns 'None', the
+ * data is removed (and its reference dropped).
+ */
+#[cfg(not(stage0))]
+pub unsafe fn modify<T: 'static>(key: Key<T>,
+                                 f: &fn(Option<T>) -> Option<T>) {
+    match f(pop(key)) {
+        Some(next) => { set(key, next); }
+        None => {}
+    }
+}
 
 #[test]
 fn test_tls_multitask() {
diff --git a/src/libstd/rt/task.rs b/src/libstd/rt/task.rs
index 04c1f972a82..55e633e8496 100644
--- a/src/libstd/rt/task.rs
+++ b/src/libstd/rt/task.rs
@@ -167,15 +167,15 @@ mod test {
 
     #[test]
     fn tls() {
-        use local_data::*;
+        use local_data;
         do run_in_newsched_task() {
             unsafe {
                 fn key(_x: @~str) { }
-                local_data_set(key, @~"data");
-                assert!(*local_data_get(key).get() == ~"data");
+                local_data::set(key, @~"data");
+                assert!(*local_data::get(key).get() == ~"data");
                 fn key2(_x: @~str) { }
-                local_data_set(key2, @~"data");
-                assert!(*local_data_get(key2).get() == ~"data");
+                local_data::set(key2, @~"data");
+                assert!(*local_data::get(key2).get() == ~"data");
             }
         }
     }
diff --git a/src/libstd/task/local_data_priv.rs b/src/libstd/task/local_data_priv.rs
index 17d534cfd03..66a459c23e6 100644
--- a/src/libstd/task/local_data_priv.rs
+++ b/src/libstd/task/local_data_priv.rs
@@ -13,12 +13,10 @@
 use cast;
 use libc;
 use local_data;
-use managed::raw::BoxRepr;
 use prelude::*;
 use ptr;
 use sys;
 use task::rt;
-use unstable::intrinsics;
 use util;
 
 use super::rt::rust_task;
@@ -50,7 +48,7 @@ impl Handle {
 trait LocalData {}
 impl<T: 'static> LocalData for T {}
 
-// The task-local-map actuall stores all TLS information. Right now it's a list
+// The task-local-map actually stores all TLS information. Right now it's a list
 // of triples of (key, value, loans). The key is a code pointer (right now at
 // least), the value is a trait so destruction can work, and the loans value
 // is a count of the number of times the value is currently on loan via
@@ -58,7 +56,7 @@ impl<T: 'static> LocalData for T {}
 //
 // TLS is designed to be able to store owned data, so `local_data_get` must
 // return a borrowed pointer to this data. In order to have a proper lifetime, a
-// borrowed pointer is insted yielded to a closure specified to the `get`
+// borrowed pointer is instead yielded to a closure specified to the `get`
 // function. As a result, it would be unsound to perform `local_data_set` on the
 // same key inside of a `local_data_get`, so we ensure at runtime that this does
 // not happen.
@@ -68,7 +66,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 TLSValue = @LocalData;
+type TLSValue = ~LocalData:;
 
 fn cleanup_task_local_map(map_ptr: *libc::c_void) {
     unsafe {
@@ -136,28 +134,8 @@ unsafe fn key_to_key_value<T: 'static>(key: local_data::Key<T>) -> *libc::c_void
     return pair.code as *libc::c_void;
 }
 
-unsafe fn transmute_back<'a, T>(data: &'a TLSValue) -> (*BoxRepr, &'a T) {
-    // Currently, a TLSValue is an '@Trait' instance which means that its actual
-    // representation is a pair of (vtable, box). Also, because of issue #7673
-    // the box actually points to another box which has the data. Hence, to get
-    // a pointer to the actual value that we're interested in, we decode the
-    // trait pointer and pass through one layer of boxes to get to the actual
-    // data we're interested in.
-    //
-    // The reference count of the containing @Trait box is already taken care of
-    // because the TLSValue is owned by the containing TLS map which means that
-    // the reference count is at least one. Extra protections are then added at
-    // runtime to ensure that once a loan on a value in TLS has been given out,
-    // the value isn't modified by another user.
-    let (_vt, box) = *cast::transmute::<&TLSValue, &(uint, *BoxRepr)>(data);
-
-    return (box, cast::transmute(&(*box).data));
-}
-
 pub unsafe fn local_pop<T: 'static>(handle: Handle,
                                     key: local_data::Key<T>) -> Option<T> {
-    // If you've never seen horrendously unsafe code written in rust before,
-    // just feel free to look a bit farther...
     let map = get_local_map(handle);
     let key_value = key_to_key_value(key);
 
@@ -175,25 +153,23 @@ pub unsafe fn local_pop<T: 'static>(handle: Handle,
                     None => libc::abort(),
                 };
 
-                // First, via some various cheats/hacks, we extract the value
-                // contained within the TLS box. This leaves a big chunk of
-                // memory which needs to be deallocated now.
-                let (chunk, inside) = transmute_back(&data);
-                let inside = cast::transmute_mut(inside);
-                let ret = ptr::read_ptr(inside);
+                // 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);
+
+                // 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));
 
-                // Forget the trait box because we're about to manually
-                // deallocate the other box. And for my next trick (kids don't
-                // try this at home), transmute the chunk of @ memory from the
-                // @-trait box to a pointer to a zero-sized '@' block which will
-                // then cause it to get properly deallocated, but it won't touch
-                // any of the uninitialized memory beyond the end.
-                cast::forget(data);
-                let chunk: *mut BoxRepr = cast::transmute(chunk);
-                (*chunk).header.type_desc =
-                    cast::transmute(intrinsics::get_tydesc::<()>());
-                let _: @() = cast::transmute(chunk);
+                // 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);
             }
             _ => {}
@@ -213,9 +189,17 @@ pub unsafe fn local_get<T: 'static, U>(handle: Handle,
     for map.mut_iter().advance |entry| {
         match *entry {
             Some((k, ref data, ref mut loans)) if k == key_value => {
+                let ret;
                 *loans = *loans + 1;
-                let (_, val) = transmute_back(data);
-                let ret = f(Some(val));
+                // 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));
+                    }
+                }
                 *loans = *loans - 1;
                 return ret;
             }
@@ -225,44 +209,46 @@ pub unsafe fn local_get<T: 'static, U>(handle: Handle,
     return f(None);
 }
 
-// FIXME(#7673): This shouldn't require '@', it should use '~'
 pub unsafe fn local_set<T: 'static>(handle: Handle,
-                                    key: local_data::Key<@T>,
-                                    data: @T) {
+                                    key: local_data::Key<T>,
+                                    data: T) {
     let map = get_local_map(handle);
     let keyval = key_to_key_value(key);
 
     // 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
+    // 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.
-    let data = @data as @LocalData;
-
-    // First, try to insert it if we already have it.
-    for map.mut_iter().advance |entry| {
-        match *entry {
-            Some((key, ref mut value, loans)) if key == keyval => {
-                if loans != 0 {
-                    fail!("TLS value has been loaned via get already");
+    //
+    // 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
                 }
-                util::replace(value, data);
-                return;
+                _ => 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())
         }
     }
-    // Next, search for an open spot
-    for map.mut_iter().advance |entry| {
-        match *entry {
-            Some(*) => {}
-            None => {
-                *entry = Some((keyval, data, 0));
-                return;
-            }
-        }
+
+    match insertion_position(map, keyval) {
+        Some(i) => { map[i] = Some((keyval, data, 0)); }
+        None => { map.push(Some((keyval, data, 0))); }
     }
-    // Finally push it on the end of the list
-    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 a8e8cfd163a..b012a834ed0 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;