Rewrite spawn yet again

The motivation here is that the bottom of each stack needs to contain a C++
try/catch block so that we can unwind. This is already the case for main, but
not spawned tasks.

Issue #236

Author: brson

src/lib/task.rs

@@ -39,20 +39,17 @@ native "rust" mod rustrt {
     fn new_task() -> task_id;
     fn drop_task(task: *rust_task);
     fn get_task_pointer(id: task_id) -> *rust_task;
-    fn start_task(id: task_id);
     fn get_task_trampoline() -> u32;
 
     fn migrate_alloc(alloc: *u8, target: task_id);
+    fn start_task(id: task_id, closure: *u8);
 }
 
 type rust_task =
     {id: task,
      mutable notify_enabled: u32,
      mutable notify_chan: comm::chan<task_notification>,
-     ctx: task_context,
-     stack_ptr: *u8};
-
-type task_context = {regs: x86_registers, next: *u8};
+     mutable stack_ptr: *u8};
 
 resource rust_task_ptr(task: *rust_task) { rustrt::drop_task(task); }
 
@@ -115,20 +112,22 @@ fn spawn_inner(thunk: -fn(), notify: option<comm::chan<task_notification>>) ->
    task_id {
     let id = rustrt::new_task();
 
-    // the order of arguments are outptr, taskptr, envptr.
-    // LLVM fastcall puts the first two in ecx, edx, and the rest on the
-    // stack.
+    let raw_thunk: {code: u32, env: u32} = cast(thunk);
 
     // set up the task pointer
     let task_ptr = rust_task_ptr(rustrt::get_task_pointer(id));
-    let regs = ptr::addr_of((**task_ptr).ctx.regs);
-    (*regs).edx = cast(*task_ptr);;
-    (*regs).esp = cast((**task_ptr).stack_ptr);
 
     assert (ptr::null() != (**task_ptr).stack_ptr);
 
-    let raw_thunk: {code: u32, env: u32} = cast(thunk);
-    (*regs).eip = raw_thunk.code;
+    // copy the thunk from our stack to the new stack
+    let sp: uint = cast((**task_ptr).stack_ptr);
+    let ptrsize = sys::size_of::<*u8>();
+    let thunkfn: *mutable uint = cast(sp - ptrsize * 2u);
+    let thunkenv: *mutable uint = cast(sp - ptrsize);
+    *thunkfn = cast(raw_thunk.code);
+    *thunkenv = cast(raw_thunk.env);
+    // align the stack to 16 bytes
+    (**task_ptr).stack_ptr = cast(sp - ptrsize * 4u);
 
     // set up notifications if they are enabled.
     alt notify {
@@ -139,60 +138,14 @@ fn spawn_inner(thunk: -fn(), notify: option<comm::chan<task_notification>>) ->
       none { }
     };
 
-    // okay, now we align the stack and add the environment pointer and a fake
-    // return address.
-
-    // -12 for the taskm output location, the env pointer
-    // -4 for the return address.
-    (*regs).esp = align_down((*regs).esp - 12u32) - 4u32;
-
-    let ra: *mutable u32 = cast((*regs).esp);
-    let env: *mutable u32 = cast((*regs).esp + 4u32);
-    let tptr: *mutable u32 = cast((*regs).esp + 12u32);
-
-    // put the return pointer in ecx.
-    (*regs).ecx = (*regs).esp + 8u32;;
-
-    *tptr = cast(*task_ptr);;
-    *env = raw_thunk.env;;
-    *ra = rustrt::get_task_trampoline();
-
+    // give the thunk environment's allocation to the new task
     rustrt::migrate_alloc(cast(raw_thunk.env), id);
-    rustrt::start_task(id);
-
+    rustrt::start_task(id, cast(thunkfn));
+    // don't cleanup the thunk in this task
     unsafe::leak(thunk);
-
     ret id;
 }
 
-// Who says we can't write an operating system in Rust?
-type x86_registers =
-    // This needs to match the structure in context.h
-
-
-    {mutable eax: u32,
-     mutable ebx: u32,
-     mutable ecx: u32,
-     mutable edx: u32,
-     mutable ebp: u32,
-     mutable esi: u32,
-     mutable edi: u32,
-     mutable esp: u32,
-     mutable cs: u16,
-     mutable ds: u16,
-     mutable ss: u16,
-     mutable es: u16,
-     mutable fs: u16,
-     mutable gs: u16,
-     mutable eflags: u32,
-     mutable eip: u32};
-
-fn align_down(x: u32) -> u32 {
-
-    // Aligns x down to 16 bytes
-    x & !15u32
-}
-
 // Local Variables:
 // mode: rust;
 // fill-column: 78;

src/rt/arch/i386/_context.S

@@ -75,11 +75,3 @@ swap_registers:
     jmp *48(%eax)
 
 
-.globl task_trampoline
-task_trampoline:
-    // This gets set up by std::task::_spawn.
-#if defined(__APPLE__) || defined(__WIN32__)
-    call _task_exit
-#else
-    call task_exit
-#endif

src/rt/main.ll.in

@@ -24,7 +24,20 @@ define void @_rust_main_wrap(i1* nocapture, %task *, %2* nocapture, %vec *)
   ret void
 }
 
+%nullary_fn = type void (i1*, %task*, %2*)
+
+define void @_rust_spawn_wrap(
+       i1* nocapture, %task*, %2* nocapture, %nullary_fn* %f)
+{
+  call fastcc void %f(i1* %0, %task *%1, %2* nocapture %2)
+  ret void
+}
+
+declare external void @set_spawn_wrapper(void (i1*, %task*, %2*, %nullary_fn*)*);
+
 define i32 @"MAIN"(i32, i32) {
-  %3 = tail call i32 @rust_start(i32 ptrtoint (void (i1*, %task*, %2*, %vec*)* @_rust_main_wrap to i32), i32 %0, i32 %1, i32 ptrtoint (%0* @_rust_crate_map_toplevel to i32))
-  ret i32 %3
+  call void @set_spawn_wrapper(void (i1*, %task*, %2*, %nullary_fn*)* @_rust_spawn_wrap)
+
+  %result = tail call i32 @rust_start(i32 ptrtoint (void (i1*, %task*, %2*, %vec*)* @_rust_main_wrap to i32), i32 %0, i32 %1, i32 ptrtoint (%0* @_rust_crate_map_toplevel to i32))
+  ret i32 %result
 }

src/rt/rust_builtin.cpp

@@ -440,18 +440,24 @@ get_task_pointer(rust_task *task, rust_task_id id) {
     return task->kernel->get_task_by_id(id);
 }
 
-extern "C" CDECL void
-start_task(rust_task *task, rust_task_id id) {
-    rust_task * target = task->kernel->get_task_by_id(id);
-    target->start();
-    target->deref();
+// FIXME: Transitional. Remove
+extern "C" CDECL void **
+get_task_trampoline(rust_task *task) {
+    return NULL;
 }
 
-extern "C" void *task_trampoline asm("task_trampoline");
+struct fn_env_pair {
+    intptr_t f;
+    intptr_t env;
+};
+
+extern "C" CDECL uintptr_t get_spawn_wrapper();
 
-extern "C" CDECL void **
-get_task_trampoline(rust_task *task) {
-    return &task_trampoline;
+extern "C" CDECL void
+start_task(rust_task *task, rust_task_id id, fn_env_pair *f) {
+    rust_task *target = task->kernel->get_task_by_id(id);
+    target->start(get_spawn_wrapper(), f->f, f->env);
+    target->deref();
 }
 
 extern "C" CDECL void

src/rt/rust_scheduler.cpp

@@ -48,10 +48,10 @@ void
 rust_scheduler::activate(rust_task *task) {
     context ctx;
 
-    task->user.ctx.next = &ctx;
+    task->ctx.next = &ctx;
     DLOG(this, task, "descheduling...");
     lock.unlock();
-    task->user.ctx.swap(ctx);
+    task->ctx.swap(ctx);
     lock.lock();
     DLOG(this, task, "task has returned");
 }

src/rt/rust_task.cpp

@@ -125,13 +125,13 @@ struct spawn_args {
                        uintptr_t, uintptr_t);
 };
 
-struct rust_closure {
+struct rust_closure_env {
     intptr_t ref_count;
     type_desc *td;
 };
 
 extern "C" CDECL
-void task_exit(rust_closure *env, int rval, rust_task *task) {
+void task_exit(rust_closure_env *env, int rval, rust_task *task) {
     LOG(task, task, "task exited with value %d", rval);
     if(env) {
         // free the environment.
@@ -155,14 +155,32 @@ void task_start_wrapper(spawn_args *a)
     int rval = 42;
 
     a->f(&rval, task, a->a3, a->a4);
-    task_exit(NULL, rval, task);
+    task_exit((rust_closure_env*)a->a3, rval, task);
+}
+
+/* We spawn a rust (fastcc) function through a CDECL function
+   defined in main.ll, which is built as part of each crate. These accessors
+   allow each rust program to install that function at startup */
+
+uintptr_t spawn_wrapper;
+
+extern "C" CDECL void
+set_spawn_wrapper(uintptr_t f) {
+    spawn_wrapper = f;
+}
+
+extern "C" CDECL uintptr_t
+get_spawn_wrapper() {
+    return spawn_wrapper;
 }
 
 void
 rust_task::start(uintptr_t spawnee_fn,
-                 uintptr_t args)
+                 uintptr_t args,
+                 uintptr_t env)
 {
-    LOGPTR(sched, "from spawnee", spawnee_fn);
+    LOG(this, task, "starting task from fn 0x%" PRIxPTR
+        " with args 0x%" PRIxPTR, spawnee_fn, args);
 
     I(sched, stk->data != NULL);
 
@@ -173,16 +191,23 @@ rust_task::start(uintptr_t spawnee_fn,
     spawn_args *a = (spawn_args *)sp;
 
     a->task = this;
-    a->a3 = 0;
+    a->a3 = env;
     a->a4 = args;
     void **f = (void **)&a->f;
     *f = (void *)spawnee_fn;
 
-    user.ctx.call((void *)task_start_wrapper, a, sp);
+    ctx.call((void *)task_start_wrapper, a, sp);
 
     this->start();
 }
 
+void
+rust_task::start(uintptr_t spawnee_fn,
+                 uintptr_t args)
+{
+    start(spawnee_fn, args, 0);
+}
+
 void rust_task::start()
 {
     yield_timer.reset_us(0);
@@ -213,7 +238,7 @@ rust_task::yield(size_t time_in_us) {
     yield_timer.reset_us(time_in_us);
 
     // Return to the scheduler.
-    user.ctx.next->swap(user.ctx);
+    ctx.next->swap(ctx);
 }
 
 void

src/rt/rust_task.h

@@ -29,7 +29,6 @@ struct rust_task_user {
     uint32_t notify_enabled;   // this is way more bits than necessary, but it
                                // simplifies the alignment.
     chan_handle notify_chan;
-    context ctx;
     uintptr_t rust_sp;         // Saved sp when not running.
 };
 
@@ -55,6 +54,7 @@ rust_task : public kernel_owned<rust_task>, rust_cond
     RUST_ATOMIC_REFCOUNT();
 
     // Fields known to the compiler.
+    context ctx;
     stk_seg *stk;
     uintptr_t runtime_sp;      // Runtime sp while task running.
     void *gc_alloc_chain;      // Linked list of GC allocations.
@@ -118,6 +118,9 @@ rust_task : public kernel_owned<rust_task>, rust_cond
 
     ~rust_task();
 
+    void start(uintptr_t spawnee_fn,
+               uintptr_t args,
+               uintptr_t env);
     void start(uintptr_t spawnee_fn,
                uintptr_t args);
     void start();

src/rt/rustrt.def.in

@@ -55,6 +55,7 @@ rust_run_program
 rust_start
 rust_getcwd
 set_min_stack
+set_spawn_wrapper
 sched_threads
 size_of
 squareroot