authorgravatar for squeek502@hotmail.comRyan Liptak <squeek502@hotmail.com> 2024-07-12 00:38:10-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-07-13 14:48:17-07:00
log959d227d1370c9cf9198ed111e3a8bca15a7e47b
tree4141025bd114fa54fedf1e76adf19202561bc018
parent11534aa34d6e9c66081eb2918bfaacbb21db0e56

ArgIteratorWindows: Reduce allocated memory by parsing the WTF-16 string directly

Before this commit, the WTF-16 command line string would be converted to WTF-8 in `init`, and then a second buffer of the WTF-8 size + 1 would be allocated to store the parsed arguments. The converted WTF-8 command line would then be parsed and the relevant bytes would be copied into the argument buffer before being returned. After this commit, only the WTF-8 size of the WTF-16 string is calculated (without conversion) which is then used to allocate the buffer for the parsed arguments. Parsing is then done on the WTF-16 slice directly, with the arguments being converted to WTF-8 on-the-fly. This has a few (minor) benefits: - Cuts the amount of memory allocated by ArgIteratorWindows in half (or better) - Makes the total amount of memory allocated by ArgIteratorWindows predictable, since, before, the upfront `wtf16LeToWtf8Alloc` call could end up allocating more-memory-than-necessary temporarily due to its internal use of an ArrayList. Now, the amount of memory allocated is always exactly `calcWtf8Len(cmd_line) + 1`.

2 files changed, 67 insertions(+), 23 deletions(-)

lib/std/process.zig+34-23
......@@ -663,11 +663,11 @@ pub const ArgIteratorWasi = struct {
663663/// - https://daviddeley.com/autohotkey/parameters/parameters.htm#WINCRULES
664664pub const ArgIteratorWindows = struct {
665665 allocator: Allocator,
666 /// Owned by the iterator.
667 /// Encoded as WTF-8.
668 cmd_line: []const u8,
666 /// Encoded as WTF-16 LE.
667 cmd_line: [:0]const u16,
669668 index: usize = 0,
670 /// Owned by the iterator. Long enough to hold the entire `cmd_line` plus a null terminator.
669 /// Owned by the iterator. Long enough to hold contiguous NUL-terminated slices
670 /// of each argument encoded as WTF-8.
671671 buffer: []u8,
672672 start: usize = 0,
673673 end: usize = 0,
......@@ -676,13 +676,18 @@ pub const ArgIteratorWindows = struct {
676676
677677 /// `cmd_line_w` *must* be a WTF16-LE-encoded string.
678678 ///
679 /// The iterator makes a copy of `cmd_line_w` converted WTF-8 and keeps it; it does *not* take
680 /// ownership of `cmd_line_w`.
679 /// The iterator stores and uses `cmd_line_w`, so its memory must be valid for
680 /// at least as long as the returned ArgIteratorWindows.
681681 pub fn init(allocator: Allocator, cmd_line_w: [*:0]const u16) InitError!ArgIteratorWindows {
682 const cmd_line = try unicode.wtf16LeToWtf8Alloc(allocator, mem.sliceTo(cmd_line_w, 0));
683 errdefer allocator.free(cmd_line);
684
685 const buffer = try allocator.alloc(u8, cmd_line.len + 1);
682 const cmd_line = mem.sliceTo(cmd_line_w, 0);
683 const wtf8_len = unicode.calcWtf8Len(cmd_line);
684
685 // This buffer must be large enough to contain contiguous NUL-terminated slices
686 // of each argument. For arguments past the first one, space for the NUL-terminator
687 // is guaranteed due to the necessary whitespace between arugments. However, we need
688 // one extra byte to guarantee enough room for the NUL terminator if the command line
689 // ends up being exactly 1 argument long with no quotes, etc.
690 const buffer = try allocator.alloc(u8, wtf8_len + 1);
686691 errdefer allocator.free(buffer);
687692
688693 return .{
......@@ -714,11 +719,11 @@ pub const ArgIteratorWindows = struct {
714719 for (0..count) |_| emitCharacter(self, '\\');
715720 }
716721
717 fn emitCharacter(self: *ArgIteratorWindows, char: u8) void {
718 self.buffer[self.end] = char;
719 self.end += 1;
722 fn emitCharacter(self: *ArgIteratorWindows, code_unit: u16) void {
723 const wtf8_len = std.unicode.wtf8Encode(code_unit, self.buffer[self.end..]) catch unreachable;
724 self.end += wtf8_len;
720725
721 // Because we are emitting WTF-8 byte-by-byte, we need to
726 // Because we are emitting WTF-8, we need to
722727 // check to see if we've emitted two consecutive surrogate
723728 // codepoints that form a valid surrogate pair in order
724729 // to ensure that we're always emitting well-formed WTF-8
......@@ -732,9 +737,7 @@ pub const ArgIteratorWindows = struct {
732737 // This is relevant when dealing with a WTF-16 encoded
733738 // command line like this:
734739 // "<0xD801>"<0xDC37>
735 // which would get converted to WTF-8 in `cmd_line` as:
736 // "<0xED><0xA0><0x81>"<0xED><0xB0><0xB7>
737 // and then after parsing it'd naively get emitted as:
740 // which would get parsed and converted to WTF-8 as:
738741 // <0xED><0xA0><0x81><0xED><0xB0><0xB7>
739742 // but instead, we need to recognize the surrogate pair
740743 // and emit the codepoint it encodes, which in this
......@@ -780,7 +783,7 @@ pub const ArgIteratorWindows = struct {
780783
781784 fn emitBackslashes(_: *ArgIteratorWindows, _: usize) void {}
782785
783 fn emitCharacter(_: *ArgIteratorWindows, _: u8) void {}
786 fn emitCharacter(_: *ArgIteratorWindows, _: u16) void {}
784787
785788 fn yieldArg(_: *ArgIteratorWindows) bool {
786789 return true;
......@@ -798,7 +801,10 @@ pub const ArgIteratorWindows = struct {
798801
799802 var inside_quotes = false;
800803 while (true) : (self.index += 1) {
801 const char = if (self.index != self.cmd_line.len) self.cmd_line[self.index] else 0;
804 const char = if (self.index != self.cmd_line.len)
805 mem.littleToNative(u16, self.cmd_line[self.index])
806 else
807 0;
802808 switch (char) {
803809 0 => {
804810 return strategy.yieldArg(self);
......@@ -823,7 +829,10 @@ pub const ArgIteratorWindows = struct {
823829
824830 // Skip spaces and tabs. The iterator completes if we reach the end of the string here.
825831 while (true) : (self.index += 1) {
826 const char = if (self.index != self.cmd_line.len) self.cmd_line[self.index] else 0;
832 const char = if (self.index != self.cmd_line.len)
833 mem.littleToNative(u16, self.cmd_line[self.index])
834 else
835 0;
827836 switch (char) {
828837 0 => return strategy.eof,
829838 ' ', '\t' => continue,
......@@ -844,7 +853,10 @@ pub const ArgIteratorWindows = struct {
844853 var backslash_count: usize = 0;
845854 var inside_quotes = false;
846855 while (true) : (self.index += 1) {
847 const char = if (self.index != self.cmd_line.len) self.cmd_line[self.index] else 0;
856 const char = if (self.index != self.cmd_line.len)
857 mem.littleToNative(u16, self.cmd_line[self.index])
858 else
859 0;
848860 switch (char) {
849861 0 => {
850862 strategy.emitBackslashes(self, backslash_count);
......@@ -867,7 +879,7 @@ pub const ArgIteratorWindows = struct {
867879 } else {
868880 if (inside_quotes and
869881 self.index + 1 != self.cmd_line.len and
870 self.cmd_line[self.index + 1] == '"')
882 mem.littleToNative(u16, self.cmd_line[self.index + 1]) == '"')
871883 {
872884 strategy.emitCharacter(self, '"');
873885 self.index += 1;
......@@ -892,7 +904,6 @@ pub const ArgIteratorWindows = struct {
892904 /// argument slices.
893905 pub fn deinit(self: *ArgIteratorWindows) void {
894906 self.allocator.free(self.buffer);
895 self.allocator.free(self.cmd_line);
896907 }
897908};
898909
lib/std/unicode.zig+33
......@@ -2107,3 +2107,36 @@ test "well-formed WTF-16 roundtrips" {
21072107 mem.nativeToLittle(u16, 0xDC00), // low surrogate
21082108 });
21092109}
2110
2111/// Returns the length, in bytes, that would be necessary to encode the
2112/// given WTF-16 LE slice as WTF-8.
2113pub fn calcWtf8Len(wtf16le: []const u16) usize {
2114 var it = Wtf16LeIterator.init(wtf16le);
2115 var num_wtf8_bytes: usize = 0;
2116 while (it.nextCodepoint()) |codepoint| {
2117 // Note: If utf8CodepointSequenceLength is ever changed to error on surrogate
2118 // codepoints, then it would no longer be eligible to be used in this context.
2119 num_wtf8_bytes += utf8CodepointSequenceLength(codepoint) catch |err| switch (err) {
2120 error.CodepointTooLarge => unreachable,
2121 };
2122 }
2123 return num_wtf8_bytes;
2124}
2125
2126fn testCalcWtf8Len() !void {
2127 const L = utf8ToUtf16LeStringLiteral;
2128 try testing.expectEqual(@as(usize, 1), calcWtf8Len(L("a")));
2129 try testing.expectEqual(@as(usize, 10), calcWtf8Len(L("abcdefghij")));
2130 // unpaired surrogate
2131 try testing.expectEqual(@as(usize, 3), calcWtf8Len(&[_]u16{
2132 mem.nativeToLittle(u16, 0xD800),
2133 }));
2134 try testing.expectEqual(@as(usize, 15), calcWtf8Len(L("こんにちは")));
2135 // First codepoints that are encoded as 1, 2, 3, and 4 bytes
2136 try testing.expectEqual(@as(usize, 1 + 2 + 3 + 4), calcWtf8Len(L("\u{0}\u{80}\u{800}\u{10000}")));
2137}
2138
2139test "calculate wtf8 string length of given wtf16 string" {
2140 try testCalcWtf8Len();
2141 try comptime testCalcWtf8Len();
2142}