std.Io.Writer.Allocating: configurable bump amount

This commit is contained in:
Andrew Kelley 2025-08-08 16:43:24 -07:00
parent 5f7a0bbabf
commit df46ee61c4
2 changed files with 15 additions and 3 deletions

View file

@ -2497,6 +2497,10 @@ pub fn Hashing(comptime Hasher: type) type {
pub const Allocating = struct { pub const Allocating = struct {
allocator: Allocator, allocator: Allocator,
writer: Writer, writer: Writer,
/// Every call to `drain` ensures at least this amount of unused capacity
/// before it returns. This prevents an infinite loop in interface logic
/// that calls `drain`.
minimum_unused_capacity: usize = 1,
pub fn init(allocator: Allocator) Allocating { pub fn init(allocator: Allocator) Allocating {
return .{ return .{
@ -2607,14 +2611,13 @@ pub const Allocating = struct {
const gpa = a.allocator; const gpa = a.allocator;
const pattern = data[data.len - 1]; const pattern = data[data.len - 1];
const splat_len = pattern.len * splat; const splat_len = pattern.len * splat;
const bump = a.minimum_unused_capacity;
var list = a.toArrayList(); var list = a.toArrayList();
defer setArrayList(a, list); defer setArrayList(a, list);
const start_len = list.items.len; const start_len = list.items.len;
// Even if we append no data, this function needs to ensure there is more
// capacity in the buffer to avoid infinite loop, hence the +1 in this loop.
assert(data.len != 0); assert(data.len != 0);
for (data) |bytes| { for (data) |bytes| {
list.ensureUnusedCapacity(gpa, bytes.len + splat_len + 1) catch return error.WriteFailed; list.ensureUnusedCapacity(gpa, bytes.len + splat_len + bump) catch return error.WriteFailed;
list.appendSliceAssumeCapacity(bytes); list.appendSliceAssumeCapacity(bytes);
} }
if (splat == 0) { if (splat == 0) {

View file

@ -73,7 +73,12 @@ const indirect_vtable: Reader.VTable = .{
.readVec = readVec, .readVec = readVec,
}; };
/// `input` buffer is asserted to be at least 10 bytes, or EOF before then.
///
/// If `buffer` is provided then asserted to have `flate.max_window_len`
/// capacity, as well as `flate.history_len` unused capacity on every write.
pub fn init(input: *Reader, container: Container, buffer: []u8) Decompress { pub fn init(input: *Reader, container: Container, buffer: []u8) Decompress {
if (buffer.len != 0) assert(buffer.len >= flate.max_window_len);
return .{ return .{
.reader = .{ .reader = .{
.vtable = if (buffer.len == 0) &direct_vtable else &indirect_vtable, .vtable = if (buffer.len == 0) &direct_vtable else &indirect_vtable,
@ -234,6 +239,8 @@ fn decodeSymbol(self: *Decompress, decoder: anytype) !Symbol {
} }
fn streamDirect(r: *Reader, w: *Writer, limit: std.Io.Limit) Reader.StreamError!usize { fn streamDirect(r: *Reader, w: *Writer, limit: std.Io.Limit) Reader.StreamError!usize {
assert(w.buffer.len >= flate.max_window_len);
assert(w.unusedCapacityLen() >= flate.history_len);
const d: *Decompress = @alignCast(@fieldParentPtr("reader", r)); const d: *Decompress = @alignCast(@fieldParentPtr("reader", r));
return streamFallible(d, w, limit); return streamFallible(d, w, limit);
} }
@ -1246,6 +1253,7 @@ test "zlib should not overshoot" {
fn testFailure(container: Container, in: []const u8, expected_err: anyerror) !void { fn testFailure(container: Container, in: []const u8, expected_err: anyerror) !void {
var reader: Reader = .fixed(in); var reader: Reader = .fixed(in);
var aw: Writer.Allocating = .init(testing.allocator); var aw: Writer.Allocating = .init(testing.allocator);
aw.minimum_unused_capacity = flate.history_len;
try aw.ensureUnusedCapacity(flate.max_window_len); try aw.ensureUnusedCapacity(flate.max_window_len);
defer aw.deinit(); defer aw.deinit();
@ -1257,6 +1265,7 @@ fn testFailure(container: Container, in: []const u8, expected_err: anyerror) !vo
fn testDecompress(container: Container, compressed: []const u8, expected_plain: []const u8) !void { fn testDecompress(container: Container, compressed: []const u8, expected_plain: []const u8) !void {
var in: std.Io.Reader = .fixed(compressed); var in: std.Io.Reader = .fixed(compressed);
var aw: std.Io.Writer.Allocating = .init(testing.allocator); var aw: std.Io.Writer.Allocating = .init(testing.allocator);
aw.minimum_unused_capacity = flate.history_len;
try aw.ensureUnusedCapacity(flate.max_window_len); try aw.ensureUnusedCapacity(flate.max_window_len);
defer aw.deinit(); defer aw.deinit();