http fixes

Author: andrewrk

lib/std/Io/Writer.zig

@@ -191,29 +191,87 @@ pub fn writeSplatHeader(
     data: []const []const u8,
     splat: usize,
 ) Error!usize {
-    const new_end = w.end + header.len;
-    if (new_end <= w.buffer.len) {
-        @memcpy(w.buffer[w.end..][0..header.len], header);
-        w.end = new_end;
-        return header.len + try writeSplat(w, data, splat);
-    }
-    var vecs: [8][]const u8 = undefined; // Arbitrarily chosen size.
-    var i: usize = 1;
-    vecs[0] = header;
-    for (data[0 .. data.len - 1]) |buf| {
-        if (buf.len == 0) continue;
-        vecs[i] = buf;
-        i += 1;
-        if (vecs.len - i == 0) break;
+    return writeSplatHeaderLimit(w, header, data, splat, .unlimited);
+}
+
+/// Equivalent to `writeSplatHeader` but writes at most `limit` bytes.
+pub fn writeSplatHeaderLimit(
+    w: *Writer,
+    header: []const u8,
+    data: []const []const u8,
+    splat: usize,
+    limit: Limit,
+) Error!usize {
+    var remaining = @intFromEnum(limit);
+    {
+        const copy_len = @min(header.len, w.buffer.len - w.end, remaining);
+        if (header.len - copy_len != 0) return writeSplatHeaderLimitFinish(w, header, data, splat, remaining);
+        @memcpy(w.buffer[w.end..][0..copy_len], header[0..copy_len]);
+        w.end += copy_len;
+        remaining -= copy_len;
+    }
+    for (data[0 .. data.len - 1], 0..) |buf, i| {
+        const copy_len = @min(buf.len, w.buffer.len - w.end, remaining);
+        if (buf.len - copy_len != 0) return @intFromEnum(limit) - remaining +
+            try writeSplatHeaderLimitFinish(w, &.{}, data[i..], splat, remaining);
+        @memcpy(w.buffer[w.end..][0..copy_len], buf[0..copy_len]);
+        w.end += copy_len;
+        remaining -= copy_len;
     }
     const pattern = data[data.len - 1];
-    const new_splat = s: {
-        if (pattern.len == 0 or vecs.len - i == 0) break :s 1;
+    const splat_n = pattern.len * splat;
+    if (splat_n > @min(w.buffer.len - w.end, remaining)) {
+        const buffered_n = @intFromEnum(limit) - remaining;
+        const written = try writeSplatHeaderLimitFinish(w, &.{}, data[data.len - 1 ..][0..1], splat, remaining);
+        return buffered_n + written;
+    }
+
+    for (0..splat) |_| {
+        @memcpy(w.buffer[w.end..][0..pattern.len], pattern);
+        w.end += pattern.len;
+    }
+
+    remaining -= splat_n;
+    return @intFromEnum(limit) - remaining;
+}
+
+fn writeSplatHeaderLimitFinish(
+    w: *Writer,
+    header: []const u8,
+    data: []const []const u8,
+    splat: usize,
+    limit: usize,
+) Error!usize {
+    var remaining = limit;
+    var vecs: [8][]const u8 = undefined;
+    var i: usize = 0;
+    v: {
+        if (header.len != 0) {
+            const copy_len = @min(header.len, remaining);
+            vecs[i] = header[0..copy_len];
+            i += 1;
+            remaining -= copy_len;
+            if (remaining == 0) break :v;
+        }
+        for (data[0 .. data.len - 1]) |buf| if (buf.len != 0) {
+            const copy_len = @min(header.len, remaining);
+            vecs[i] = buf;
+            i += 1;
+            remaining -= copy_len;
+            if (remaining == 0) break :v;
+            if (vecs.len - i == 0) break :v;
+        };
+        const pattern = data[data.len - 1];
+        if (splat == 1) {
+            vecs[i] = pattern[0..@min(remaining, pattern.len)];
+            i += 1;
+            break :v;
+        }
         vecs[i] = pattern;
         i += 1;
-        break :s splat;
-    };
-    return w.vtable.drain(w, vecs[0..i], new_splat);
+        return w.vtable.drain(w, (&vecs)[0..i], @min(remaining / pattern.len, splat));
+    }
+    return w.vtable.drain(w, (&vecs)[0..i], 1);
 }
 
 test "writeSplatHeader splatting avoids buffer aliasing temptation" {