std: Fix perf of local allocations in newsched

src/libstd/os.rs

@@ -1142,9 +1142,9 @@ pub fn real_args() -> ~[~str] {
 #[cfg(target_os = "freebsd")]
 pub fn real_args() -> ~[~str] {
     use rt;
-    use rt::TaskContext;
+    use rt::NewRtContext;
 
-    if rt::context() == TaskContext {
+    if rt::context() == NewRtContext {
         match rt::args::clone() {
             Some(args) => args,
             None => fail!("process arguments not initialized")

src/libstd/rt/comm.rs

@@ -15,6 +15,7 @@ use cast;
 use ops::Drop;
 use rt::kill::BlockedTask;
 use kinds::Send;
+use rt;
 use rt::sched::Scheduler;
 use rt::local::Local;
 use rt::select::{Select, SelectPort};
@@ -24,7 +25,6 @@ use util::Void;
 use comm::{GenericChan, GenericSmartChan, GenericPort, Peekable};
 use cell::Cell;
 use clone::Clone;
-use rt::{context, SchedulerContext};
 use tuple::ImmutableTuple;
 
 /// A combined refcount / BlockedTask-as-uint pointer.
@@ -113,7 +113,7 @@ impl<T> ChanOne<T> {
     // 'do_resched' configures whether the scheduler immediately switches to
     // the receiving task, or leaves the sending task still running.
     fn try_send_inner(self, val: T, do_resched: bool) -> bool {
-        rtassert!(context() != SchedulerContext);
+        rtassert!(!rt::in_sched_context());
 
         let mut this = self;
         let mut recvr_active = true;

src/libstd/rt/local_heap.rs

@@ -13,6 +13,7 @@
 use libc;
 use libc::{c_void, uintptr_t, size_t};
 use ops::Drop;
+use option::{Some, None};
 use rt;
 use rt::OldTaskContext;
 use rt::local::Local;
@@ -86,20 +87,19 @@ impl Drop for LocalHeap {
 
 // A little compatibility function
 pub unsafe fn local_free(ptr: *libc::c_char) {
-    match rt::context() {
-        OldTaskContext => {
+    // XXX: Unsafe borrow for speed. Lame.
+    match Local::try_unsafe_borrow::<Task>() {
+        Some(task) => {
+            (*task).heap.free(ptr as *libc::c_void);
+        }
+        None => {
             rust_upcall_free_noswitch(ptr);
 
             extern {
                 #[fast_ffi]
                 fn rust_upcall_free_noswitch(ptr: *libc::c_char);
             }
         }
-        _ => {
-            do Local::borrow::<Task,()> |task| {
-                task.heap.free(ptr as *libc::c_void);
-            }
-        }
     }
 }
 
@@ -119,20 +119,28 @@ pub fn live_allocs() -> *raw::Box<()> {
 }
 
 extern {
+    #[fast_ffi]
     fn rust_new_memory_region(synchronized: uintptr_t,
                                detailed_leaks: uintptr_t,
                                poison_on_free: uintptr_t) -> *MemoryRegion;
+    #[fast_ffi]
     fn rust_delete_memory_region(region: *MemoryRegion);
+    #[fast_ffi]
     fn rust_new_boxed_region(region: *MemoryRegion,
                              poison_on_free: uintptr_t) -> *BoxedRegion;
+    #[fast_ffi]
     fn rust_delete_boxed_region(region: *BoxedRegion);
+    #[fast_ffi]
     fn rust_boxed_region_malloc(region: *BoxedRegion,
                                 td: *TypeDesc,
                                 size: size_t) -> *OpaqueBox;
+    #[fast_ffi]
     fn rust_boxed_region_realloc(region: *BoxedRegion,
                                  ptr: *OpaqueBox,
                                  size: size_t) -> *OpaqueBox;
+    #[fast_ffi]
     fn rust_boxed_region_free(region: *BoxedRegion, box: *OpaqueBox);
+    #[fast_ffi]
     fn rust_current_boxed_region() -> *BoxedRegion;
 }
 

src/libstd/rt/mod.rs

@@ -407,14 +407,10 @@ fn run_(main: ~fn(), use_main_sched: bool) -> int {
 /// or the old scheduler.
 #[deriving(Eq)]
 pub enum RuntimeContext {
-    // Only the exchange heap is available
-    GlobalContext,
-    // The scheduler may be accessed
-    SchedulerContext,
-    // Full task services, e.g. local heap, unwinding
-    TaskContext,
     // Running in an old-style task
-    OldTaskContext
+    OldTaskContext,
+    // Not old task context
+    NewRtContext
 }
 
 /// Determine the current RuntimeContext
@@ -424,23 +420,40 @@ pub fn context() -> RuntimeContext {
 
     if unsafe { rust_try_get_task().is_not_null() } {
         return OldTaskContext;
-    } else if Local::exists::<Task>() {
-        // In this case we know it is a new runtime context, but we
-        // need to check which one. Going to try borrowing task to
-        // check. Task should always be in TLS, so hopefully this
-        // doesn't conflict with other ops that borrow.
-        return do Local::borrow::<Task,RuntimeContext> |task| {
-            match task.task_type {
-                SchedTask => SchedulerContext,
-                GreenTask(_) => TaskContext
-            }
-        };
     } else {
-        return GlobalContext;
+        return NewRtContext;
     }
 
     extern {
         #[rust_stack]
         pub fn rust_try_get_task() -> *rust_task;
     }
 }
+
+pub fn in_sched_context() -> bool {
+    unsafe {
+        match Local::try_unsafe_borrow::<Task>() {
+            Some(task) => {
+                match (*task).task_type {
+                    SchedTask => true,
+                    _ => false
+                }
+            }
+            None => false
+        }
+    }
+}
+
+pub fn in_green_task_context() -> bool {
+    unsafe {
+        match Local::try_unsafe_borrow::<Task>() {
+            Some(task) => {
+                match (*task).task_type {
+                    GreenTask(_) => true,
+                    _ => false
+                }
+            }
+            None => false
+        }
+    }
+}

src/libstd/sys.rs

@@ -136,7 +136,7 @@ impl FailWithCause for &'static str {
 pub fn begin_unwind_(msg: *c_char, file: *c_char, line: size_t) -> ! {
     use either::Left;
     use option::{Some, None};
-    use rt::{context, OldTaskContext, TaskContext};
+    use rt::{context, OldTaskContext, in_green_task_context};
     use rt::task::Task;
     use rt::local::Local;
     use rt::logging::Logger;
@@ -158,7 +158,7 @@ pub fn begin_unwind_(msg: *c_char, file: *c_char, line: size_t) -> ! {
 
                 // XXX: Logging doesn't work correctly in non-task context because it
                 // invokes the local heap
-                if context == TaskContext {
+                if in_green_task_context() {
                     // XXX: Logging doesn't work here - the check to call the log
                     // function never passes - so calling the log function directly.
                     do Local::borrow::<Task, ()> |task| {

src/libstd/task/mod.rs

@@ -42,7 +42,7 @@ use cmp::Eq;
 use comm::{stream, Chan, GenericChan, GenericPort, Port};
 use result::Result;
 use result;
-use rt::{context, OldTaskContext, TaskContext};
+use rt::{context, OldTaskContext, in_green_task_context};
 use rt::local::Local;
 use unstable::finally::Finally;
 use util;
@@ -527,14 +527,15 @@ pub fn try<T:Send>(f: ~fn() -> T) -> Result<T,()> {
 pub fn with_task_name<U>(blk: &fn(Option<&str>) -> U) -> U {
     use rt::task::Task;
 
-    match context() {
-        TaskContext => do Local::borrow::<Task, U> |task| {
+    if in_green_task_context() {
+        do Local::borrow::<Task, U> |task| {
             match task.name {
                 Some(ref name) => blk(Some(name.as_slice())),
                 None => blk(None)
             }
-        },
-        _ => fail!("no task name exists in %?", context()),
+        }
+    } else {
+        fail!("no task name exists in %?", context())
     }
 }
 
@@ -614,7 +615,7 @@ pub fn unkillable<U>(f: &fn() -> U) -> U {
                     rt::rust_task_allow_kill(t);
                 }
             }
-            TaskContext => {
+            _ if in_green_task_context() => {
                 // The inhibits/allows might fail and need to borrow the task.
                 let t = Local::unsafe_borrow::<Task>();
                 do (|| {
@@ -645,7 +646,7 @@ pub unsafe fn rekillable<U>(f: &fn() -> U) -> U {
                 rt::rust_task_inhibit_kill(t);
             }
         }
-        TaskContext => {
+        _ if in_green_task_context() => {
             let t = Local::unsafe_borrow::<Task>();
             do (|| {
                 (*t).death.allow_kill((*t).unwinder.unwinding);

src/libstd/task/spawn.rs

@@ -91,7 +91,7 @@ use to_bytes::IterBytes;
 use uint;
 use util;
 use unstable::sync::Exclusive;
-use rt::{OldTaskContext, TaskContext, SchedulerContext, GlobalContext, context};
+use rt::{OldTaskContext, NewRtContext, context, in_green_task_context};
 use rt::local::Local;
 use rt::task::{Task, Sched};
 use rt::kill::KillHandle;
@@ -526,7 +526,7 @@ impl RuntimeGlue {
                 let me = rt::rust_get_task();
                 blk(OldTask(me), rt::rust_task_is_unwinding(me))
             },
-            TaskContext => unsafe {
+            NewRtContext if in_green_task_context() => unsafe {
                 // Can't use safe borrow, because the taskgroup destructor needs to
                 // access the scheduler again to send kill signals to other tasks.
                 let me = Local::unsafe_borrow::<Task>();
@@ -535,7 +535,7 @@ impl RuntimeGlue {
                 blk(NewTask((*me).death.kill_handle.get_ref().clone()),
                     (*me).unwinder.unwinding)
             },
-            SchedulerContext | GlobalContext => rtabort!("task dying in bad context"),
+            NewRtContext => rtabort!("task dying in bad context"),
         }
     }
 
@@ -563,7 +563,7 @@ impl RuntimeGlue {
                     }
                 }
             },
-            TaskContext => unsafe {
+            NewRtContext if in_green_task_context() => unsafe {
                 // Can't use safe borrow, because creating new hashmaps for the
                 // tasksets requires an rng, which needs to borrow the sched.
                 let me = Local::unsafe_borrow::<Task>();
@@ -588,7 +588,7 @@ impl RuntimeGlue {
                     Some(ref group) => group,
                 })
             },
-            SchedulerContext | GlobalContext => rtabort!("spawning in bad context"),
+            NewRtContext => rtabort!("spawning in bad context"),
         }
     }
 }
@@ -666,10 +666,9 @@ fn enlist_many(child: TaskHandle, child_arc: &TaskGroupArc,
 
 pub fn spawn_raw(opts: TaskOpts, f: ~fn()) {
     match context() {
-        OldTaskContext   => spawn_raw_oldsched(opts, f),
-        TaskContext      => spawn_raw_newsched(opts, f),
-        SchedulerContext => fail!("can't spawn from scheduler context"),
-        GlobalContext    => fail!("can't spawn from global context"),
+        OldTaskContext => spawn_raw_oldsched(opts, f),
+        _ if in_green_task_context() => spawn_raw_newsched(opts, f),
+        _ => fail!("can't spawn from this context")
     }
 }
 

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

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