diff --git a/lib/std/buffer.zig b/lib/std/buffer.zig index 93281c35ec..6d361bdb4a 100644 --- 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 index 237d001450..93734d42aa 100644 --- 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 };