Accept main(args: [str]) as main signature

1 files changed, 44 insertions, 4 deletions

diff --git a/src/rt/rust.cpp b/src/rt/rust.cpp
index cc2f86da4b2..c211408d4d4 100644
--- a/src/rt/rust.cpp
+++ b/src/rt/rust.cpp
@@ -10,10 +10,12 @@ command_line_args : public kernel_owned<command_line_args>
 
     // vec[str] passed to rust_task::start.
     rust_vec *args;
+    rust_ivec *args_ivec;
 
     command_line_args(rust_task *task,
                       int sys_argc,
-                      char **sys_argv)
+                      char **sys_argv,
+                      bool main_is_ivec)
         : kernel(task->kernel),
           task(task),
           argc(sys_argc),
@@ -49,15 +51,39 @@ command_line_args : public kernel_owned<command_line_args>
             mem = kernel->malloc(str_alloc, "command line arg");
             strs[i] = new (mem) rust_str(str_alloc, str_fill,
                                          (uint8_t const *)argv[i]);
+            strs[i]->ref_count++;
         }
         args->fill = vec_fill;
         // If the caller has a declared args array, they may drop; but
         // we don't know if they have such an array. So we pin the args
         // array here to ensure it survives to program-shutdown.
         args->ref();
+
+        if (main_is_ivec) {
+            size_t ivec_interior_sz =
+                sizeof(size_t) * 2 + sizeof(rust_str *) * 4;
+            args_ivec = (rust_ivec *)
+                kernel->malloc(ivec_interior_sz,
+                               "command line arg interior");
+            args_ivec->fill = 0;
+            size_t ivec_exterior_sz = sizeof(rust_str *) * argc;
+            args_ivec->alloc = ivec_exterior_sz;
+            // NB: This is freed by some ivec machinery, probably the drop
+            // glue in main, so we don't free it ourselves
+            args_ivec->payload.ptr = (rust_ivec_heap *)
+                kernel->malloc(ivec_exterior_sz + sizeof(size_t),
+                               "command line arg exterior");
+            args_ivec->payload.ptr->fill = ivec_exterior_sz;
+            memcpy(&args_ivec->payload.ptr->data, strs, ivec_exterior_sz);
+        } else {
+            args_ivec = NULL;
+        }
     }
 
     ~command_line_args() {
+        if (args_ivec) {
+            kernel->free(args_ivec);
+        }
         if (args) {
             // Drop the args we've had pinned here.
             rust_str **strs = (rust_str**) &args->data[0];
@@ -84,7 +110,8 @@ command_line_args : public kernel_owned<command_line_args>
 int check_claims = 0;
 
 extern "C" CDECL int
-rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
+rust_start_ivec(uintptr_t main_fn, int argc, char **argv,
+                void* crate_map, int main_is_ivec) {
 
     rust_env *env = load_env();
 
@@ -99,7 +126,7 @@ rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
     rust_scheduler *sched = root_task->sched;
     command_line_args *args
         = new (kernel, "main command line args")
-        command_line_args(root_task, argc, argv);
+        command_line_args(root_task, argc, argv, main_is_ivec);
 
     DLOG(sched, dom, "startup: %d args in 0x%" PRIxPTR,
              args->argc, (uintptr_t)args->args);
@@ -107,7 +134,13 @@ rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
         DLOG(sched, dom, "startup: arg[%d] = '%s'", i, args->argv[i]);
     }
 
-    root_task->start(main_fn, (uintptr_t)args->args);
+    if (main_is_ivec) {
+        DLOG(sched, dom, "main takes ivec");
+        root_task->start(main_fn, (uintptr_t)args->args_ivec);
+    } else {
+        DLOG(sched, dom, "main takes vec");
+        root_task->start(main_fn, (uintptr_t)args->args);
+    }
     root_task->deref();
     root_task = NULL;
 
@@ -127,6 +160,13 @@ rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
     return ret;
 }
 
+extern "C" CDECL int
+rust_start(uintptr_t main_fn, int argc, char **argv,
+           void* crate_map) {
+    return rust_start_ivec(main_fn, argc, argv, crate_map, 0);
+}
+
+
 //
 // Local Variables:
 // mode: C++