motiejus/zig
1
Fork 0
Code Releases Activity

Merge pull request #5388 from kubkon/wasi-args-iter

Browse Source 
Add ArgIteratorWasi and integrate it with ArgIterator
This commit is contained in:
Andrew Kelley
2020-05-29 15:55:14 -04:00
committed by GitHub
parent e9f344dcd4 6e347e6180
commit 65c3833ec2
1 changed files with 133 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

175
lib/std/process.zig
View File

@@ -185,6 +185,79 @@ pub const ArgIteratorPosix = struct {
}
};
pub const ArgIteratorWasi = struct {
allocator: *mem.Allocator,
index: usize,
args: [][]u8,
pub const InitError = error{OutOfMemory} || os.UnexpectedError;
/// You must call deinit to free the internal buffer of the
/// iterator after you are done.
pub fn init(allocator: *mem.Allocator) InitError!ArgIteratorWasi {
const fetched_args = try ArgIteratorWasi.internalInit(allocator);
return ArgIteratorWasi{
.allocator = allocator,
.index = 0,
.args = fetched_args,
};
}
fn internalInit(allocator: *mem.Allocator) InitError![][]u8 {
const w = os.wasi;
var count: usize = undefined;
var buf_size: usize = undefined;
switch (w.args_sizes_get(&count, &buf_size)) {
w.ESUCCESS => {},
else => |err| return os.unexpectedErrno(err),
}
var argv = try allocator.alloc([*:0]u8, count);
defer allocator.free(argv);
var argv_buf = try allocator.alloc(u8, buf_size);
switch (w.args_get(argv.ptr, argv_buf.ptr)) {
w.ESUCCESS => {},
else => |err| return os.unexpectedErrno(err),
}
var result_args = try allocator.alloc([]u8, count);
var i: usize = 0;
while (i < count) : (i += 1) {
result_args[i] = mem.spanZ(argv[i]);
}
return result_args;
}
pub fn next(self: *ArgIteratorWasi) ?[]const u8 {
if (self.index == self.args.len) return null;
const arg = self.args[self.index];
self.index += 1;
return arg;
}
pub fn skip(self: *ArgIteratorWasi) bool {
if (self.index == self.args.len) return false;
self.index += 1;
return true;
}
/// Call to free the internal buffer of the iterator.
pub fn deinit(self: *ArgIteratorWasi) void {
const last_item = self.args[self.args.len - 1];
const last_byte_addr = @ptrToInt(last_item.ptr) + last_item.len + 1; // null terminated
const first_item_ptr = self.args[0].ptr;
const len = last_byte_addr - @ptrToInt(first_item_ptr);
self.allocator.free(first_item_ptr[0..len]);
self.allocator.free(self.args);
}
};
pub const ArgIteratorWindows = struct {
index: usize,
cmd_line: [*]const u8,
@@ -335,14 +408,29 @@ pub const ArgIteratorWindows = struct {
};
pub const ArgIterator = struct {
const InnerType = if (builtin.os.tag == .windows) ArgIteratorWindows else ArgIteratorPosix;
const InnerType = switch (builtin.os.tag) {
.windows => ArgIteratorWindows,
.wasi => ArgIteratorWasi,
else => ArgIteratorPosix,
};
inner: InnerType,
/// Initialize the args iterator.
pub fn init() ArgIterator {
if (builtin.os.tag == .wasi) {
// TODO: Figure out a compatible interface accomodating WASI
@compileError("ArgIterator is not yet supported in WASI. Use argsAlloc and argsFree instead.");
@compileError("In WASI, use initWithAllocator instead.");
}
return ArgIterator{ .inner = InnerType.init() };
}
pub const InitError = ArgIteratorWasi.InitError;
/// You must deinitialize iterator's internal buffers by calling `deinit` when done.
pub fn initWithAllocator(allocator: *mem.Allocator) InitError!ArgIterator {
if (builtin.os.tag == .wasi) {
return ArgIterator{ .inner = try InnerType.init(allocator) };
}
return ArgIterator{ .inner = InnerType.init() };
@@ -364,49 +452,62 @@ pub const ArgIterator = struct {
return self.inner.next();
}
/// If you only are targeting WASI, you can call this and not need an allocator.
pub fn nextWasi(self: *ArgIterator) ?[]const u8 {
return self.inner.next();
}
/// Parse past 1 argument without capturing it.
/// Returns `true` if skipped an arg, `false` if we are at the end.
pub fn skip(self: *ArgIterator) bool {
return self.inner.skip();
}
/// Call this to free the iterator's internal buffer if the iterator
/// was created with `initWithAllocator` function.
pub fn deinit(self: *ArgIterator) void {
// Unless we're targeting WASI, this is a no-op.
if (builtin.os.tag == .wasi) {
self.inner.deinit();
}
}
};
pub fn args() ArgIterator {
return ArgIterator.init();
}
/// You must deinitialize iterator's internal buffers by calling `deinit` when done.
pub fn argsWithAllocator(allocator: *mem.Allocator) ArgIterator.InitError!ArgIterator {
return ArgIterator.initWithAllocator(allocator);
}
test "args iterator" {
var ga = std.testing.allocator;
var it = if (builtin.os.tag == .wasi) try argsWithAllocator(ga) else args();
defer it.deinit(); // no-op unless WASI
const prog_name = try it.next(ga) orelse unreachable;
defer ga.free(prog_name);
const expected_suffix = switch (builtin.os.tag) {
.wasi => "test.wasm",
.windows => "test.exe",
else => "test",
};
const given_suffix = std.fs.path.basename(prog_name);
testing.expect(mem.eql(u8, expected_suffix, given_suffix));
testing.expectEqual(it.next(ga), null);
testing.expect(!it.skip());
}
/// Caller must call argsFree on result.
pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {
if (builtin.os.tag == .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 os.unexpectedErrno(args_sizes_get_ret);
}
var argv = try allocator.alloc([*:0]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 os.unexpectedErrno(args_get_ret);
}
var result_slice = try allocator.alloc([]u8, count);
var i: usize = 0;
while (i < count) : (i += 1) {
result_slice[i] = mem.spanZ(argv[i]);
}
return result_slice;
}
// TODO refactor to only make 1 allocation.
var it = args();
var it = if (builtin.os.tag == .wasi) try argsWithAllocator(allocator) else args();
defer it.deinit();
var contents = std.ArrayList(u8).init(allocator);
defer contents.deinit();
@@ -442,16 +543,6 @@ pub fn argsAlloc(allocator: *mem.Allocator) ![][]u8 {
}
pub fn argsFree(allocator: *mem.Allocator, args_alloc: []const []u8) void {
if (builtin.os.tag == .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;