Allow main to take istrs. Issue #855

1 files changed, 41 insertions, 3 deletions

diff --git a/src/rt/rust.cpp b/src/rt/rust.cpp
index b5a462154ed..a7411b98cd0 100644
--- a/src/rt/rust.cpp
+++ b/src/rt/rust.cpp
@@ -11,6 +11,7 @@ command_line_args : public kernel_owned<command_line_args>
 
     // [str] passed to rust_task::start.
     rust_vec *args;
+    rust_vec *args_istr;
 
     command_line_args(rust_task *task,
                       int sys_argc,
@@ -38,6 +39,8 @@ command_line_args : public kernel_owned<command_line_args>
         }
         LocalFree(wargv);
 #endif
+
+        // Allocate a vector of estrs
         size_t vec_fill = sizeof(rust_str *) * argc;
         size_t vec_alloc = next_power_of_two(vec_fill);
         void *mem = kernel->malloc(vec_alloc, "command line");
@@ -55,17 +58,40 @@ command_line_args : public kernel_owned<command_line_args>
             kernel->malloc(vec_size<rust_str*>(argc),
                            "command line arg interior");
         args->fill = args->alloc = sizeof(rust_str *) * argc;
-        // NB: _rust_main owns the vec and will be responsible for
-        // freeing it
         memcpy(&args->data[0], strs, args->fill);
+
+        // Allocate a vector of istrs
+        args_istr = (rust_vec *)
+            kernel->malloc(vec_size<rust_vec*>(argc),
+                           "command line arg interior");
+        args_istr->fill = args_istr->alloc = sizeof(rust_vec*) * argc;
+        for (int i = 0; i < argc; ++i) {
+            size_t str_fill = strlen(argv[i]) + 1;
+            size_t str_alloc = str_fill;
+            rust_vec *str = (rust_vec *)
+                kernel->malloc(vec_size<char>(str_fill),
+                               "command line arg");
+            str->fill = str_fill;
+            str->alloc = str_alloc;
+            memcpy(&str->data, argv[i], str_fill);
+            ((rust_vec**)&args_istr->data)[i] = str;
+        }
     }
 
     ~command_line_args() {
+        // Free the estr args
         kernel->free(args);
         for (int i = 0; i < argc; ++i)
             kernel->free(strs[i]);
         kernel->free(strs);
 
+        // Free the istr args
+        for (int i = 0; i < argc; ++i) {
+            rust_vec *s = ((rust_vec**)&args_istr->data)[i];
+            kernel->free(s);
+        }
+        kernel->free(args_istr);
+
 #ifdef __WIN32__
         for (int i = 0; i < argc; ++i) {
             kernel->free(argv[i]);
@@ -76,6 +102,14 @@ command_line_args : public kernel_owned<command_line_args>
 };
 
 
+// FIXME: Transitional. Please remove.
+bool main_takes_istr = false;
+
+extern "C" CDECL void
+set_main_takes_istr(uintptr_t flag) {
+    main_takes_istr = flag != 0;
+}
+
 /**
  * Main entry point into the Rust runtime. Here we create a Rust service,
  * initialize the kernel, create the root domain and run it.
@@ -108,7 +142,11 @@ rust_start(uintptr_t main_fn, int argc, char **argv,
         DLOG(sched, dom, "startup: arg[%d] = '%s'", i, args->argv[i]);
     }
 
-    root_task->start(main_fn, (uintptr_t)args->args);
+    if (main_takes_istr) {
+        root_task->start(main_fn, (uintptr_t)args->args_istr);
+    } else {
+        root_task->start(main_fn, (uintptr_t)args->args);
+    }
     root_task->deref();
     root_task = NULL;