feat(cli)!: Allow specifying WASI environment variables and CLI args (#1840)

Author: alex-snezhko

cli/bin/exec.js

@@ -122,20 +122,37 @@ function getGrainrun() {
 
 const grainrun = getGrainrun();
 
-function execGrainrun(file, options, program, execOpts = { stdio: "inherit" }) {
+function execGrainrun(
+  unprocessedArgs,
+  file,
+  options,
+  program,
+  execOpts = { stdio: "inherit" }
+) {
   const preopens = {};
   options.dir?.forEach((preopen) => {
     const [guestDir, hostDir = guestDir] = preopen.split("=");
     preopens[guestDir] = hostDir;
   });
 
+  const cliEnv = {};
+  options.env?.forEach((env) => {
+    const [name, ...rest] = env.split("=");
+    const val = rest.join("=");
+    cliEnv[name] = val;
+  });
+
   const env = {
+    ENV_VARS: JSON.stringify(cliEnv),
     PREOPENS: JSON.stringify(preopens),
     NODE_OPTIONS: `--experimental-wasi-unstable-preview1 --no-warnings`,
   };
 
   try {
-    execSync(`${grainrun} ${file}`, { ...execOpts, env });
+    execSync(`${grainrun} ${file} ${unprocessedArgs.join(" ")}`, {
+      ...execOpts,
+      env,
+    });
   } catch (e) {
     process.exit(e.status);
   }

cli/bin/grain.js

@@ -107,6 +107,7 @@ class GrainCommand extends commander.Command {
     );
     cmd.forwardOption("--import-memory", "import the memory from `env.memory`");
     cmd.option("--dir <dir...>", "directory to preopen");
+    cmd.option("--env <env...>", "WASI environment variables");
     cmd.forwardOption(
       "--compilation-mode <mode>",
       "compilation mode (advanced use only)"
@@ -164,6 +165,13 @@ class GrainCommand extends commander.Command {
   }
 }
 
+let endOptsI = process.argv.findIndex((x) => x === "--");
+if (endOptsI === -1) {
+  endOptsI = Infinity;
+}
+const argsToProcess = process.argv.slice(0, endOptsI);
+const unprocessedArgs = process.argv.slice(endOptsI + 1);
+
 const program = new GrainCommand();
 
 program
@@ -178,11 +186,8 @@ program
   .forwardOption("-o <filename>", "output filename")
   .action(function (file, options, program) {
     exec.grainc(file, options, program);
-    if (options.o) {
-      exec.grainrun(options.o, options, program);
-    } else {
-      exec.grainrun(file.replace(/\.gr$/, ".gr.wasm"), options, program);
-    }
+    const outFile = options.o ?? file.replace(/\.gr$/, ".gr.wasm");
+    exec.grainrun(unprocessedArgs, outFile, options, program);
   });
 
 program
@@ -195,7 +200,7 @@ program
 program
   .command("run <file>")
   .description("run a wasm file via grain's WASI runner")
-  .action(exec.grainrun);
+  .action((...args) => exec.grainrun(unprocessedArgs, ...args));
 
 program
   .command("lsp")
@@ -218,4 +223,4 @@ program
   .forwardOption("-o <file|dir>", "output file or directory")
   .action(exec.grainformat);
 
-program.parse(process.argv);
+program.parse(argsToProcess);

cli/bin/grainrun.js

@@ -17,13 +17,11 @@ const { WASI } = require("wasi");
 const { argv, env } = require("process");
 
 const preopens = JSON.parse(env.PREOPENS);
-
-delete env.PREOPENS;
-delete env.NODE_OPTIONS;
+const envVars = JSON.parse(env.ENV_VARS);
 
 const wasi = new WASI({
   args: argv.slice(2),
-  env,
+  env: envVars,
   preopens,
 });
 

compiler/test/runner.re

@@ -158,7 +158,7 @@ let open_process = args => {
   (code, out, err);
 };
 
-let run = (~num_pages=?, file) => {
+let run = (~num_pages=?, ~extra_args=[||], file) => {
   let mem_flags =
     switch (num_pages) {
     | Some(x) => [|
@@ -179,12 +179,22 @@ let run = (~num_pages=?, file) => {
       Filepath.to_string(test_data_dir),
     );
 
+  let extra_args =
+    Array.map(
+      arg =>
+        switch (Sys.win32, arg) {
+        | (true, "--") => "`--"
+        | _ => arg
+        },
+      extra_args,
+    );
   let cmd =
     Array.concat([
       [|"grain", "run"|],
       mem_flags,
       [|"-S", stdlib, "-I", Filepath.to_string(test_libs_dir), preopen|],
       [|file|],
+      extra_args,
     ]);
 
   let (code, out, err) = open_process(cmd);
@@ -289,11 +299,12 @@ let makeNoWarningRunner = (test, name, prog) => {
   });
 };
 
-let makeRunner = (test, ~num_pages=?, ~config_fn=?, name, prog, expected) => {
+let makeRunner =
+    (test, ~num_pages=?, ~config_fn=?, ~extra_args=?, name, prog, expected) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
       ignore @@ compile(~num_pages?, ~config_fn?, name, module_header ++ prog);
-      let (result, _) = run(~num_pages?, wasmfile(name));
+      let (result, _) = run(~num_pages?, ~extra_args?, wasmfile(name));
       expect.string(result).toEqual(expected);
     })
   });

compiler/test/stdlib/sys.process.test.gr

@@ -11,7 +11,7 @@ match (Process.argv()) {
 
 // Just a smoke test
 match (Process.env()) {
-  Ok(arr) => assert Array.length(arr) > 0,
+  Ok(arr) => assert Array.length(arr) == 0,
   Err(err) => throw err,
 }
 

compiler/test/suites/wasi_args.re

@@ -0,0 +1,61 @@
+open Grain_tests.TestFramework;
+open Grain_tests.Runner;
+
+describe("wasi args and env", ({test, testSkip}) => {
+  let test_or_skip =
+    Sys.backend_type == Other("js_of_ocaml") ? testSkip : test;
+
+  let assertRun = makeRunner(test_or_skip);
+
+  let print_wasi_info = {|
+  include "sys/process"
+  include "array"
+  include "string"
+
+  match (Process.argv()) {
+    Ok(args) => print(Array.slice(1, args)),
+    _ => print("Error reading args")
+  }
+  match (Process.env()) {
+    Ok(env) => print(env),
+    _ => print("Error reading env")
+  }
+  |};
+
+  assertRun(
+    ~extra_args=[|"--", "a", "b"|],
+    "print_args1",
+    print_wasi_info,
+    "[> \"a\", \"b\"]\n[> ]\n",
+  );
+  assertRun(
+    ~extra_args=[|"a", "b"|],
+    "print_args2",
+    print_wasi_info,
+    "[> ]\n[> ]\n",
+  );
+  assertRun(
+    ~extra_args=[|"--env=FOO=bar", "a", "b"|],
+    "print_args3",
+    print_wasi_info,
+    "[> ]\n[> \"FOO=bar\"]\n",
+  );
+  assertRun(
+    ~extra_args=[|"--env", "FOO=bar", "BAR=baz", "BAZ"|],
+    "print_args4",
+    print_wasi_info,
+    "[> ]\n[> \"FOO=bar\", \"BAR=baz\", \"BAZ=\"]\n",
+  );
+  assertRun(
+    ~extra_args=[|"--env", "FOO=bar", "--", "a", "b"|],
+    "print_args5",
+    print_wasi_info,
+    "[> \"a\", \"b\"]\n[> \"FOO=bar\"]\n",
+  );
+  assertRun(
+    ~extra_args=[|"--", "a", "b", "--env", "FOO=bar"|],
+    "print_args6",
+    print_wasi_info,
+    "[> \"a\", \"b\", \"--env\", \"FOO=bar\"]\n[> ]\n",
+  );
+});

stdlib/sys/process.gr

@@ -152,6 +152,11 @@ provide let env = () => {
   let envc = WasmI32.load(envcPtr, 0n)
   let envvBufSize = WasmI32.load(envvBufSizePtr, 0n)
 
+  if (WasmI32.eqz(envc)) {
+    Memory.free(envcPtr)
+    return Ok([>]: Array<String>)
+  }
+
   let envvPtr = Memory.malloc(envc * 4n)
   let envvBufPtr = Memory.malloc(envvBufSize)