zig

commit 6a53fe7c93ddc6216e7cd41514bebe701531c9c3 (tree)
parent 0059d9ee3e7298df03723d97adbea72ff142cacd
Author: Benjamin Feng <benjamin.feng@glassdoor.com>
Date:   Fri,  6 Mar 2020 12:03:15 -0600

Handle potential downcast when translating stream size

Diffstat:
M
lib/std/buffer.zig
 | 
4
+++-
M
lib/std/fmtstream.zig
 | 
10
+++++++---

2 files changed, 10 insertions(+), 4 deletions(-)
diff --git a/lib/std/buffer.zig b/lib/std/buffer.zig
@@ -65,7 +65,9 @@ pub const Buffer = struct {
     }
 
     pub fn allocPrint(allocator: *Allocator, comptime format: []const u8, args: var) !Buffer {
-        const size = std.fmtstream.count(format, args);
+        const size = std.fmtstream.count(format, args) catch |err| switch (err) {
+            error.Overflow => return error.OutOfMemory,
+        };
         var self = try Buffer.initSize(allocator, size);
         assert((std.fmtstream.bufPrint(self.list.items, format, args) catch unreachable).len == size);
         return self;
diff --git a/lib/std/fmtstream.zig b/lib/std/fmtstream.zig
@@ -1081,16 +1081,20 @@ pub fn bufPrint(buf: []u8, comptime fmt: []const u8, args: var) BufPrintError![]
 }
 
 // Count the characters needed for format. Useful for preallocating memory
-pub fn count(comptime fmt: []const u8, args: var) usize {
+pub fn count(comptime fmt: []const u8, args: var) !usize {
     var counting_stream = std.io.countingOutStream(std.io.null_out_stream);
     format(counting_stream.outStream(), fmt, args) catch |err| switch (err) {};
-    return counting_stream.bytes_written;
+    return std.math.cast(usize, counting_stream.bytes_written);
 }
 
 pub const AllocPrintError = error{OutOfMemory};
 
 pub fn allocPrint(allocator: *mem.Allocator, comptime fmt: []const u8, args: var) AllocPrintError![]u8 {
-    const buf = try allocator.alloc(u8, count(fmt, args));
+    const size = count(fmt, args) catch |err| switch (err) {
+        // Output too long. Can't possibly allocate enough memory to display it.
+        error.Overflow => return error.OutOfMemory,
+    };
+    const buf = try allocator.alloc(u8, size);
     return bufPrint(buf, fmt, args) catch |err| switch (err) {
         error.BufferTooSmall => unreachable, // we just counted the size above
     };