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()
     }
 }