std: Fix perf of local allocations in newsched

Mostly optimizing TLS accesses to bring local heap allocation performance
closer to that of oldsched. It's not completely at parity but removing the
branches involved in supporting oldsched and optimizing pthread_get/setspecific
to instead use our dedicated TCB slot will probably make up for it.

2 files changed, 23 insertions, 20 deletions

diff --git a/src/libstd/unstable/lang.rs b/src/libstd/unstable/lang.rs
index 98c0fe254b6..c5112529aed 100644
--- a/src/libstd/unstable/lang.rs
+++ b/src/libstd/unstable/lang.rs
@@ -12,9 +12,9 @@
 
 use cast::transmute;
 use libc::{c_char, c_uchar, c_void, size_t, uintptr_t, c_int};
+use option::{Some, None};
 use str;
 use sys;
-use rt::{context, OldTaskContext};
 use rt::task::Task;
 use rt::local::Local;
 use rt::borrowck;
@@ -56,16 +56,13 @@ pub fn fail_bounds_check(file: *c_char, line: size_t,
 
 #[lang="malloc"]
 pub unsafe fn local_malloc(td: *c_char, size: uintptr_t) -> *c_char {
-    match context() {
-        OldTaskContext => {
-            return rustrt::rust_upcall_malloc_noswitch(td, size);
+    // XXX: Unsafe borrow for speed. Lame.
+    match Local::try_unsafe_borrow::<Task>() {
+        Some(task) => {
+            (*task).heap.alloc(td as *c_void, size as uint) as *c_char
         }
-        _ => {
-            let mut alloc = ::ptr::null();
-            do Local::borrow::<Task,()> |task| {
-                alloc = task.heap.alloc(td as *c_void, size as uint) as *c_char;
-            }
-            return alloc;
+        None => {
+            rustrt::rust_upcall_malloc_noswitch(td, size)
         }
     }
 }
diff --git a/src/libstd/unstable/sync.rs b/src/libstd/unstable/sync.rs
index 225ac5c92ad..a8d942a46b3 100644
--- a/src/libstd/unstable/sync.rs
+++ b/src/libstd/unstable/sync.rs
@@ -282,7 +282,7 @@ pub unsafe fn atomically<U>(f: &fn() -> U) -> U {
     use rt::task::Task;
     use task::rt;
     use rt::local::Local;
-    use rt::{context, OldTaskContext, TaskContext};
+    use rt::{context, OldTaskContext};
 
     match context() {
         OldTaskContext => {
@@ -296,17 +296,23 @@ pub unsafe fn atomically<U>(f: &fn() -> U) -> U {
                 rt::rust_task_allow_kill(t);
             }
         }
-        TaskContext => {
-            let t = Local::unsafe_borrow::<Task>();
-            do (|| {
-                (*t).death.inhibit_yield();
-                f()
-            }).finally {
-                (*t).death.allow_yield();
+        _ => {
+            let t = Local::try_unsafe_borrow::<Task>();
+            match t {
+                Some(t) => {
+                    do (|| {
+                        (*t).death.inhibit_yield();
+                        f()
+                    }).finally {
+                        (*t).death.allow_yield();
+                    }
+                }
+                None => {
+                    // FIXME(#3095): As in unkillable().
+                    f()
+                }
             }
         }
-        // FIXME(#3095): As in unkillable().
-        _ => f()
     }
 }