rt: Move some stack manipulation functions into rust_task

2 files changed, 136 insertions, 137 deletions

diff --git a/src/rt/rust_task.cpp b/src/rt/rust_task.cpp
index fcf0da29e45..dc8b911778a 100644
--- a/src/rt/rust_task.cpp
+++ b/src/rt/rust_task.cpp
@@ -57,133 +57,8 @@
 #endif
 #endif
 
-static size_t
-get_next_stk_size(rust_task_thread *thread, rust_task *task,
-                  size_t min, size_t current, size_t requested) {
-    LOG(task, mem, "calculating new stack size for 0x%" PRIxPTR, task);
-    LOG(task, mem,
-        "min: %" PRIdPTR " current: %" PRIdPTR " requested: %" PRIdPTR,
-        min, current, requested);
-
-    // Allocate at least enough to accomodate the next frame
-    size_t sz = std::max(min, requested);
-
-    // And double the stack size each allocation
-    const size_t max = 1024 * 1024;
-    size_t next = std::min(max, current * 2);
-
-    sz = std::max(sz, next);
-
-    LOG(task, mem, "next stack size: %" PRIdPTR, sz);
-    I(thread, requested <= sz);
-    return sz;
-}
-
-// Task stack segments. Heap allocated and chained together.
-
-// The amount of stack in a segment available to Rust code
-static size_t
-user_stack_size(stk_seg *stk) {
-    return (size_t)(stk->end
-                    - (uintptr_t)&stk->data[0]
-                    - RED_ZONE_SIZE);
-}
-
-static void
-free_stk(rust_task *task, stk_seg *stk) {
-    LOGPTR(task->thread, "freeing stk segment", (uintptr_t)stk);
-    task->total_stack_sz -= user_stack_size(stk);
-    task->free(stk);
-}
-
-static stk_seg*
-new_stk(rust_task_thread *thread, rust_task *task, size_t requested_sz)
-{
-    LOG(task, mem, "creating new stack for task %" PRIxPTR, task);
-    if (task->stk) {
-        check_stack_canary(task->stk);
-    }
-
-    // The minimum stack size, in bytes, of a Rust stack, excluding red zone
-    size_t min_sz = thread->min_stack_size;
-
-    // Try to reuse an existing stack segment
-    if (task->stk != NULL && task->stk->prev != NULL) {
-        size_t prev_sz = user_stack_size(task->stk->prev);
-        if (min_sz <= prev_sz && requested_sz <= prev_sz) {
-            LOG(task, mem, "reusing existing stack");
-            task->stk = task->stk->prev;
-            A(thread, task->stk->prev == NULL, "Bogus stack ptr");
-            config_valgrind_stack(task->stk);
-            return task->stk;
-        } else {
-            LOG(task, mem, "existing stack is not big enough");
-            free_stk(task, task->stk->prev);
-            task->stk->prev = NULL;
-        }
-    }
-
-    // The size of the current stack segment, excluding red zone
-    size_t current_sz = 0;
-    if (task->stk != NULL) {
-        current_sz = user_stack_size(task->stk);
-    }
-    // The calculated size of the new stack, excluding red zone
-    size_t rust_stk_sz = get_next_stk_size(thread, task, min_sz,
-                                           current_sz, requested_sz);
-
-    if (task->total_stack_sz + rust_stk_sz > thread->env->max_stack_size) {
-        LOG_ERR(task, task, "task %" PRIxPTR " ran out of stack", task);
-        task->fail();
-    }
-
-    size_t sz = sizeof(stk_seg) + rust_stk_sz + RED_ZONE_SIZE;
-    stk_seg *stk = (stk_seg *)task->malloc(sz, "stack");
-    LOGPTR(task->thread, "new stk", (uintptr_t)stk);
-    memset(stk, 0, sizeof(stk_seg));
-    add_stack_canary(stk);
-    stk->prev = NULL;
-    stk->next = task->stk;
-    stk->end = (uintptr_t) &stk->data[rust_stk_sz + RED_ZONE_SIZE];
-    LOGPTR(task->thread, "stk end", stk->end);
-
-    task->stk = stk;
-    config_valgrind_stack(task->stk);
-    task->total_stack_sz += user_stack_size(stk);
-    return stk;
-}
-
-static void
-del_stk(rust_task *task, stk_seg *stk)
-{
-    assert(stk == task->stk && "Freeing stack segments out of order!");
-    check_stack_canary(stk);
-
-    task->stk = stk->next;
-
-    bool delete_stack = false;
-    if (task->stk != NULL) {
-        // Don't actually delete this stack. Save it to reuse later,
-        // preventing the pathological case where we repeatedly reallocate
-        // the stack for the next frame.
-        task->stk->prev = stk;
-    } else {
-        // This is the last stack, delete it.
-        delete_stack = true;
-    }
-
-    // Delete the previous previous stack
-    if (stk->prev != NULL) {
-        free_stk(task, stk->prev);
-        stk->prev = NULL;
-    }
-
-    unconfig_valgrind_stack(stk);
-    if (delete_stack) {
-        free_stk(task, stk);
-        A(task->thread, task->total_stack_sz == 0, "Stack size should be 0");
-    }
-}
+extern "C" CDECL void
+record_sp(void *limit);
 
 // Tasks
 rust_task::rust_task(rust_task_thread *thread, rust_task_list *state,
@@ -218,7 +93,7 @@ rust_task::rust_task(rust_task_thread *thread, rust_task_list *state,
     LOGPTR(thread, "new task", (uintptr_t)this);
     DLOG(thread, task, "sizeof(task) = %d (0x%x)", sizeof *this, sizeof *this);
 
-    stk = new_stk(thread, this, init_stack_sz);
+    new_stack(init_stack_sz);
     if (supervisor) {
         supervisor->ref();
     }
@@ -246,7 +121,7 @@ rust_task::delete_this()
     // and no landing pads stopped to clean up.
     // FIXME: We should do this when the task exits, not in the destructor
     while (stk != NULL) {
-        del_stk(this, stk);
+        del_stack();
     }
 
     thread->release_task(this);
@@ -630,16 +505,134 @@ rust_task::notify(bool success) {
     }
 }
 
-extern "C" CDECL void
-record_sp(void *limit);
+size_t
+rust_task::get_next_stack_size(size_t min, size_t current, size_t requested) {
+    LOG(this, mem, "calculating new stack size for 0x%" PRIxPTR, this);
+    LOG(this, mem,
+        "min: %" PRIdPTR " current: %" PRIdPTR " requested: %" PRIdPTR,
+        min, current, requested);
+
+    // Allocate at least enough to accomodate the next frame
+    size_t sz = std::max(min, requested);
+
+    // And double the stack size each allocation
+    const size_t max = 1024 * 1024;
+    size_t next = std::min(max, current * 2);
+
+    sz = std::max(sz, next);
+
+    LOG(this, mem, "next stack size: %" PRIdPTR, sz);
+    I(thread, requested <= sz);
+    return sz;
+}
+
+// The amount of stack in a segment available to Rust code
+static size_t
+user_stack_size(stk_seg *stk) {
+    return (size_t)(stk->end
+                    - (uintptr_t)&stk->data[0]
+                    - RED_ZONE_SIZE);
+}
+
+void
+rust_task::free_stack(stk_seg *stk) {
+    LOGPTR(thread, "freeing stk segment", (uintptr_t)stk);
+    total_stack_sz -= user_stack_size(stk);
+    free(stk);
+}
+
+void
+rust_task::new_stack(size_t requested_sz) {
+    LOG(this, mem, "creating new stack for task %" PRIxPTR, this);
+    if (stk) {
+        ::check_stack_canary(stk);
+    }
+
+    // The minimum stack size, in bytes, of a Rust stack, excluding red zone
+    size_t min_sz = thread->min_stack_size;
+
+    // Try to reuse an existing stack segment
+    if (stk != NULL && stk->prev != NULL) {
+        size_t prev_sz = user_stack_size(stk->prev);
+        if (min_sz <= prev_sz && requested_sz <= prev_sz) {
+            LOG(this, mem, "reusing existing stack");
+            stk = stk->prev;
+            A(thread, stk->prev == NULL, "Bogus stack ptr");
+            config_valgrind_stack(stk);
+            return;
+        } else {
+            LOG(this, mem, "existing stack is not big enough");
+            free_stack(stk->prev);
+            stk->prev = NULL;
+        }
+    }
+
+    // The size of the current stack segment, excluding red zone
+    size_t current_sz = 0;
+    if (stk != NULL) {
+        current_sz = user_stack_size(stk);
+    }
+    // The calculated size of the new stack, excluding red zone
+    size_t rust_stk_sz = get_next_stack_size(min_sz,
+                                             current_sz, requested_sz);
+
+    if (total_stack_sz + rust_stk_sz > thread->env->max_stack_size) {
+        LOG_ERR(this, task, "task %" PRIxPTR " ran out of stack", this);
+        fail();
+    }
+
+    size_t sz = sizeof(stk_seg) + rust_stk_sz + RED_ZONE_SIZE;
+    stk_seg *new_stk = (stk_seg *)malloc(sz, "stack");
+    LOGPTR(thread, "new stk", (uintptr_t)new_stk);
+    memset(new_stk, 0, sizeof(stk_seg));
+    add_stack_canary(new_stk);
+    new_stk->prev = NULL;
+    new_stk->next = stk;
+    new_stk->end = (uintptr_t) &new_stk->data[rust_stk_sz + RED_ZONE_SIZE];
+    LOGPTR(thread, "stk end", new_stk->end);
+
+    stk = new_stk;
+    config_valgrind_stack(stk);
+    total_stack_sz += user_stack_size(new_stk);
+}
+
+void
+rust_task::del_stack() {
+    stk_seg *old_stk = stk;
+    ::check_stack_canary(old_stk);
+
+    stk = old_stk->next;
+
+    bool delete_stack = false;
+    if (stk != NULL) {
+        // Don't actually delete this stack. Save it to reuse later,
+        // preventing the pathological case where we repeatedly reallocate
+        // the stack for the next frame.
+        stk->prev = old_stk;
+    } else {
+        // This is the last stack, delete it.
+        delete_stack = true;
+    }
+
+    // Delete the previous previous stack
+    if (old_stk->prev != NULL) {
+        free_stack(old_stk->prev);
+        old_stk->prev = NULL;
+    }
+
+    unconfig_valgrind_stack(old_stk);
+    if (delete_stack) {
+        free_stack(old_stk);
+        A(thread, total_stack_sz == 0, "Stack size should be 0");
+    }
+}
 
 void *
 rust_task::next_stack(size_t stk_sz, void *args_addr, size_t args_sz) {
-
-    stk_seg *stk_seg = new_stk(thread, this, stk_sz + args_sz);
-    A(thread, stk_seg->end - (uintptr_t)stk_seg->data >= stk_sz + args_sz,
+    new_stack(stk_sz + args_sz);
+    A(thread, stk->end - (uintptr_t)stk->data >= stk_sz + args_sz,
       "Did not receive enough stack");
-    uint8_t *new_sp = (uint8_t*)stk_seg->end;
+    uint8_t *new_sp = (uint8_t*)stk->end;
     // Push the function arguments to the new stack
     new_sp = align_down(new_sp - args_sz);
     memcpy(new_sp, args_addr, args_sz);
@@ -651,7 +644,7 @@ rust_task::next_stack(size_t stk_sz, void *args_addr, size_t args_sz) {
 
 void
 rust_task::prev_stack() {
-    del_stk(this, stk);
+    del_stack();
     A(thread, rust_task_thread::get_task() == this,
       "Recording the stack limit for the wrong thread");
     record_stack_limit();
@@ -695,7 +688,7 @@ void
 rust_task::reset_stack_limit() {
     uintptr_t sp = get_sp();
     while (!sp_in_stk_seg(sp, stk)) {
-        del_stk(this, stk);
+        del_stack();
         A(thread, stk != NULL, "Failed to find the current stack");
     }
     record_stack_limit();
diff --git a/src/rt/rust_task.h b/src/rt/rust_task.h
index f78a0ca5805..7f270f88826 100644
--- a/src/rt/rust_task.h
+++ b/src/rt/rust_task.h
@@ -106,6 +106,12 @@ rust_task : public kernel_owned<rust_task>, rust_cond
 private:
     // Called when the atomic refcount reaches zero
     void delete_this();
+
+    void new_stack(size_t sz);
+    void del_stack();
+    void free_stack(stk_seg *stk);
+    size_t get_next_stack_size(size_t min, size_t current, size_t requested);
+
 public:
 
     // Only a pointer to 'name' is kept, so it must live as long as this task.