rt: Don't hit TLS on upcall_vec_push unless necessary

2 files changed, 15 insertions, 5 deletions

diff --git a/src/rt/rust_upcall.cpp b/src/rt/rust_upcall.cpp
index 15b149404c1..905b546b0b7 100644
--- a/src/rt/rust_upcall.cpp
+++ b/src/rt/rust_upcall.cpp
@@ -434,7 +434,7 @@ upcall_vec_grow(rust_vec** vp, size_t new_sz) {
 // Copy elements from one vector to another,
 // dealing with reference counts
 static inline void
-copy_elements(rust_task *task, type_desc *elem_t,
+copy_elements(type_desc *elem_t,
               void *pdst, void *psrc, size_t n) {
     char *dst = (char *)pdst, *src = (char *)psrc;
     memmove(dst, src, n);
@@ -454,12 +454,10 @@ extern "C" CDECL void
 upcall_vec_push(rust_vec** vp, type_desc* elt_ty, void* elt) {
     // NB: This runs entirely on the Rust stack because it invokes take glue
 
-    rust_task *task = rust_task_thread::get_task();
-
     size_t new_sz = (*vp)->fill + elt_ty->size;
-    reserve_vec(task, vp, new_sz);
+    reserve_vec_fast(vp, new_sz);
     rust_vec* v = *vp;
-    copy_elements(task, elt_ty, &v->data[0] + v->fill,
+    copy_elements(elt_ty, &v->data[0] + v->fill,
                   elt, elt_ty->size);
     v->fill += elt_ty->size;
 }
diff --git a/src/rt/rust_util.h b/src/rt/rust_util.h
index 1b611dd0be6..a0530b80dab 100644
--- a/src/rt/rust_util.h
+++ b/src/rt/rust_util.h
@@ -189,6 +189,18 @@ inline void reserve_vec(rust_task* task, rust_vec** vpp, size_t size) {
     reserve_vec_exact(task, vpp, next_power_of_two(size));
 }
 
+// Call this when you don't already have a task pointer and it will
+// avoid hitting the TLS if it doesn't have to
+inline void reserve_vec_fast(rust_vec **vpp, size_t size) {
+    size_t new_size = next_power_of_two(size);
+    if (new_size > (*vpp)->alloc) {
+        rust_task *task = rust_task_thread::get_task();
+        size_t alloc_size = new_size + sizeof(rust_vec);
+        *vpp = (rust_vec*)task->kernel->realloc(*vpp, alloc_size);
+        (*vpp)->alloc = new_size;
+    }
+}
+
 typedef rust_vec rust_str;
 
 inline rust_str *