Allow main to take istrs. Issue #855

7 files changed, 93 insertions, 10 deletions

diff --git a/src/comp/middle/trans.rs b/src/comp/middle/trans.rs
index ee31ae3e356..ec4e77a1cce 100644
--- a/src/comp/middle/trans.rs
+++ b/src/comp/middle/trans.rs
@@ -5579,21 +5579,50 @@ fn create_main_wrapper(ccx: &@crate_ctxt, sp: &span, main_llfn: ValueRef,
         ccx.sess.span_fatal(sp, ~"multiple 'main' functions");
     }
 
-    let main_takes_ivec =
+    let (main_takes_argv, main_takes_istr) =
         alt ty::struct(ccx.tcx, main_node_type) {
-          ty::ty_fn(_, args, _, _, _) { std::vec::len(args) != 0u }
+          ty::ty_fn(_, args, _, _, _) {
+            if std::vec::len(args) == 0u {
+                (false, false)
+            } else {
+                alt ty::struct(ccx.tcx, args[0].ty) {
+                  ty::ty_vec({ty: t, _}) {
+                    alt ty::struct(ccx.tcx, t) {
+                      ty::ty_str. { (true, false) }
+                      ty::ty_istr. { (true, true) }
+                    }
+                  }
+                }
+            }
+          }
         };
 
-    let llfn = create_main(ccx, sp, main_llfn, main_takes_ivec);
+    let llfn = create_main(ccx, sp, main_llfn,
+                           main_takes_argv, main_takes_istr);
     ccx.main_fn = some(llfn);
 
+    // FIXME: This is a transitional way to let the runtime know
+    // it needs to feed us istrs
+    let lltakesistr = istr::as_buf(~"_rust_main_takes_istr", { |buf|
+        llvm::LLVMAddGlobal(ccx.llmod, T_int(), buf)
+    });
+    llvm::LLVMSetInitializer(lltakesistr, C_uint(main_takes_istr as uint));
+    llvm::LLVMSetGlobalConstant(lltakesistr, True);
+    llvm::LLVMSetLinkage(lltakesistr,
+                         lib::llvm::LLVMExternalLinkage as llvm::Linkage);
+
     fn create_main(ccx: &@crate_ctxt, sp: &span, main_llfn: ValueRef,
-                   takes_ivec: bool) -> ValueRef {
+                   takes_argv: bool, takes_istr: bool) -> ValueRef {
+        let unit_ty = if takes_istr {
+            ty::mk_istr(ccx.tcx)
+        } else {
+            ty::mk_str(ccx.tcx)
+        };
         let ivecarg_ty: ty::arg =
             {mode: ty::mo_val,
              ty:
                  ty::mk_vec(ccx.tcx,
-                            {ty: ty::mk_str(ccx.tcx), mut: ast::imm})};
+                            {ty: unit_ty, mut: ast::imm})};
         let llfty =
             type_of_fn(ccx, sp, ast::proto_fn, [ivecarg_ty],
                        ty::mk_nil(ccx.tcx), 0u);
@@ -5609,7 +5638,7 @@ fn create_main_wrapper(ccx: &@crate_ctxt, sp: &span, main_llfn: ValueRef,
         let llenvarg = llvm::LLVMGetParam(llfdecl, 2u);
         let llargvarg = llvm::LLVMGetParam(llfdecl, 3u);
         let args = [lloutputarg, lltaskarg, llenvarg];
-        if takes_ivec { args += [llargvarg]; }
+        if takes_argv { args += [llargvarg]; }
         FastCall(bcx, main_llfn, args);
         build_return(bcx);
 
diff --git a/src/comp/middle/typeck.rs b/src/comp/middle/typeck.rs
index c6e11bf722d..50c99615c3f 100644
--- a/src/comp/middle/typeck.rs
+++ b/src/comp/middle/typeck.rs
@@ -2747,6 +2747,7 @@ fn arg_is_argv_ty(tcx: &ty::ctxt, a: &ty::arg) -> bool {
         if mt.mut != ast::imm { ret false; }
         alt ty::struct(tcx, mt.ty) {
           ty::ty_str. { ret true; }
+          ty::ty_istr. { ret true; }
           _ { ret false; }
         }
       }
diff --git a/src/etc/combine-tests.py b/src/etc/combine-tests.py
index c6656014525..26b20f1f8bc 100755
--- a/src/etc/combine-tests.py
+++ b/src/etc/combine-tests.py
@@ -20,6 +20,7 @@ run_pass = os.path.join(src_dir, "src", "test", "run-pass")
 run_pass = os.path.abspath(run_pass)
 stage2_tests = []
 take_args = {}
+take_iargs = {}
 
 for t in os.listdir(run_pass):
     if t.endswith(".rs") and not (
@@ -31,6 +32,9 @@ for t in os.listdir(run_pass):
             stage2_tests.append(t)
             if "main(args: [str])" in s:
                 take_args[t] = True
+            # FIXME: Transitional. Remove me
+            if "main(args: [istr])" in s:
+                take_iargs[t] = True
         f.close()
 
 stage2_tests.sort()
@@ -61,6 +65,8 @@ for t in stage2_tests:
     d.write("    out.write_str(~\"run-pass [stage2]: %s\\n\");\n" % p)
     if t in take_args:
         d.write("    t_%d::main([\"arg0\"]);\n" % i)
+    elif t in take_iargs:
+        d.write("    t_%d::main([~\"arg0\"]);\n" % i)
     else:
         d.write("    t_%d::main();\n" % i)
     i += 1
diff --git a/src/rt/main.ll.in b/src/rt/main.ll.in
index 6df51e67503..0a3cb38ef11 100644
--- a/src/rt/main.ll.in
+++ b/src/rt/main.ll.in
@@ -24,7 +24,13 @@ define void @_rust_main_wrap(i1* nocapture, %task *, %2* nocapture, %vec *)
   ret void
 }
 
+; FIXME: Transitional. Please remove
+@_rust_main_takes_istr = external global i32
+declare void @set_main_takes_istr(i32)
+
 define i32 @"MAIN"(i32, i32) {
+  %i = load i32* @_rust_main_takes_istr
+  call void @set_main_takes_istr(i32 %i)
   %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
 }
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;
 
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index a0eab5cd123..2f3f25eb64f 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -55,6 +55,7 @@ rust_ptr_eq
 rust_run_program
 rust_start
 rust_getcwd
+set_main_takes_istr
 set_min_stack
 sched_threads
 size_of
diff --git a/src/test/run-pass/main-ivec.rs b/src/test/run-pass/main-ivec.rs
index ec2403a9151..376374f68eb 100644
--- a/src/test/run-pass/main-ivec.rs
+++ b/src/test/run-pass/main-ivec.rs
@@ -1 +1,3 @@
-fn main(args: [str]) { for s in args { log s } }
+fn main(args: [istr]) {
+    for s in args { log s }
+}