From 01365be82fb413061d052c2ecdb761d24d73125d Mon Sep 17 00:00:00 2001
From: Shritesh Bhattarai <shritesh@shritesh.com>
Date: Mon, 29 Apr 2019 20:54:30 -0500
Subject: [PATCH] WASI: implement argsAlloc and argsFree (#2364)

* wasi: change URL to canon WASI-core.md

* wasi: import args_get and args_sizes_get

* wasi: Implement argsAlloc and argsFree

* test return value for wasi arg syscalls

* wasi: return unexpectedErrorPosix in argsAlloc

* wasi: Add TODO for ArgIterator
---
 std/os.zig           | 43 +++++++++++++++++++++++++++++++++++++++++++
 std/os/wasi.zig      |  2 +-
 std/os/wasi/core.zig |  3 +++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/std/os.zig b/std/os.zig
index ae2b04b759..139bc5b09d 100644
--- a/std/os.zig
+++ b/std/os.zig
@@ -2134,6 +2134,11 @@ pub const ArgIterator = struct {
     inner: InnerType,
 
     pub fn init() ArgIterator {
+        if (builtin.os == Os.wasi) {
+            // TODO: Figure out a compatible interface accomodating WASI
+            @compileError("ArgIterator is not yet supported in WASI. Use argsAlloc and argsFree instead.");
+        }
+
         return ArgIterator{ .inner = InnerType.init() };
     }
 
@@ -2166,6 +2171,34 @@ pub fn args() ArgIterator {
 
 /// Caller must call argsFree on result.
 pub fn argsAlloc(allocator: *mem.Allocator) ![]const []u8 {
+    if (builtin.os == Os.wasi) {
+        var count: usize = undefined;
+        var buf_size: usize = undefined;
+
+        const args_sizes_get_ret = os.wasi.args_sizes_get(&count, &buf_size);
+        if (args_sizes_get_ret != os.wasi.ESUCCESS) {
+            return unexpectedErrorPosix(args_sizes_get_ret);
+        }
+
+        var argv = try allocator.alloc([*]u8, count);
+        defer allocator.free(argv);
+
+        var argv_buf = try allocator.alloc(u8, buf_size);
+        const args_get_ret = os.wasi.args_get(argv.ptr, argv_buf.ptr);
+        if (args_get_ret != os.wasi.ESUCCESS) {
+            return unexpectedErrorPosix(args_get_ret);
+        }
+
+        var result_slice = try allocator.alloc([]u8, count);
+        
+        var i: usize = 0;
+        while (i < count) : (i += 1) {
+            result_slice[i] = mem.toSlice(u8, argv[i]);
+        }
+
+        return result_slice;
+    }
+
     // TODO refactor to only make 1 allocation.
     var it = args();
     var contents = try Buffer.initSize(allocator, 0);
@@ -2203,6 +2236,16 @@ pub fn argsAlloc(allocator: *mem.Allocator) ![]const []u8 {
 }
 
 pub fn argsFree(allocator: *mem.Allocator, args_alloc: []const []u8) void {
+    if (builtin.os == Os.wasi) {
+        const last_item = args_alloc[args_alloc.len - 1];
+        const last_byte_addr = @ptrToInt(last_item.ptr) + last_item.len + 1; // null terminated
+        const first_item_ptr = args_alloc[0].ptr;
+        const len = last_byte_addr - @ptrToInt(first_item_ptr);
+        allocator.free(first_item_ptr[0..len]);
+
+        return allocator.free(args_alloc);
+    }
+
     var total_bytes: usize = 0;
     for (args_alloc) |arg| {
         total_bytes += @sizeOf([]u8) + arg.len;
diff --git a/std/os/wasi.zig b/std/os/wasi.zig
index ea3dd8d740..c273fc6c9e 100644
--- a/std/os/wasi.zig
+++ b/std/os/wasi.zig
@@ -1,7 +1,7 @@
 pub use @import("wasi/core.zig");
 
 // Based on https://github.com/CraneStation/wasi-sysroot/blob/wasi/libc-bottom-half/headers/public/wasi/core.h
-// and https://github.com/CraneStation/wasmtime/blob/master/docs/WASI-api.md
+// and https://github.com/WebAssembly/WASI/blob/master/design/WASI-core.md
 
 pub const STDIN_FILENO = 0;
 pub const STDOUT_FILENO = 1;
diff --git a/std/os/wasi/core.zig b/std/os/wasi/core.zig
index 903442296d..dc24218dfc 100644
--- a/std/os/wasi/core.zig
+++ b/std/os/wasi/core.zig
@@ -10,6 +10,9 @@ pub const ciovec_t = extern struct {
 
 pub const SIGABRT: signal_t = 6;
 
+pub extern "wasi_unstable" fn args_get(argv: [*][*]u8, argv_buf: [*]u8) errno_t;
+pub extern "wasi_unstable" fn args_sizes_get(argc: *usize, argv_buf_size: *usize) errno_t;
+
 pub extern "wasi_unstable" fn proc_raise(sig: signal_t) errno_t;
 
 pub extern "wasi_unstable" fn proc_exit(rval: exitcode_t) noreturn;