1const MultiReader = @This();
2
3const std = @import("../../std.zig");
4const Io = std.Io;
5const File = Io.File;
6const Allocator = std.mem.Allocator;
7const assert = std.debug.assert;
8
9gpa: Allocator,
10streams: *Streams,
11batch: Io.Batch,
12
13pub const Context = struct {
14 mr: *MultiReader,
15 fr: File.Reader,
16 vec: [1][]u8,
17 err: ?Error,
18};
19
20pub const Error = UnendingError || error{EndOfStream};
21pub const UnendingError = Allocator.Error || File.Reader.Error || Io.ConcurrentError;
22
23/// Trailing:
24/// * `contexts: [len]Context`
25/// * `storage: [len]Io.Operation.Storage`
26pub const Streams = extern struct {
27 len: u32,
28
29 pub fn contexts(s: *Streams) []Context {
30 const base: usize = @intFromPtr(s);
31 const ptr: [*]Context = @ptrFromInt(std.mem.alignForward(usize, base + @sizeOf(Streams), @alignOf(Context)));
32 return ptr[0..s.len];
33 }
34
35 pub fn storage(s: *Streams) []Io.Operation.Storage {
36 const prev = contexts(s);
37 const end = prev.ptr + prev.len;
38 const ptr: [*]Io.Operation.Storage = @ptrFromInt(std.mem.alignForward(usize, @intFromPtr(end), @alignOf(Io.Operation.Storage)));
39 return ptr[0..s.len];
40 }
41};
42
43pub fn Buffer(comptime n: usize) type {
44 return extern struct {
45 len: u32,
46 contexts: [n][@sizeOf(Context)]u8 align(@alignOf(Context)),
47 storage: [n][@sizeOf(Io.Operation.Storage)]u8 align(@alignOf(Io.Operation.Storage)),
48
49 pub fn toStreams(b: *@This()) *Streams {
50 b.len = n;
51 return @ptrCast(b);
52 }
53 };
54}
55
56/// See `Streams.Buffer` for convenience API to obtain the `streams` parameter.
57pub fn init(mr: *MultiReader, gpa: Allocator, io: Io, streams: *Streams, files: []const File) void {
58 const contexts = streams.contexts();
59 for (contexts, files) |*context, file| context.* = .{
60 .mr = mr,
61 .fr = .{
62 .io = io,
63 .file = file,
64 .mode = .streaming,
65 .interface = .{
66 .vtable = &.{
67 .stream = stream,
68 .discard = discard,
69 .readVec = readVec,
70 .rebase = rebase,
71 },
72 .buffer = &.{},
73 .seek = 0,
74 .end = 0,
75 },
76 },
77 .vec = .{&.{}},
78 .err = null,
79 };
80 mr.* = .{
81 .gpa = gpa,
82 .streams = streams,
83 .batch = .init(streams.storage()),
84 };
85 for (contexts, 0..) |*context, i| {
86 const r = &context.fr.interface;
87 rebaseGrowing(mr, context, 1) catch |err| {
88 context.err = err;
89 continue;
90 };
91 context.vec[0] = r.buffer;
92 mr.batch.addAt(@intCast(i), .{ .file_read_streaming = .{
93 .file = context.fr.file,
94 .data = &context.vec,
95 } });
96 }
97}
98
99pub fn deinit(mr: *MultiReader) void {
100 const gpa = mr.gpa;
101 const contexts = mr.streams.contexts();
102 const io = contexts[0].fr.io;
103 mr.batch.cancel(io);
104 for (contexts) |*context| {
105 gpa.free(context.fr.interface.buffer);
106 }
107}
108
109pub fn fileReader(mr: *MultiReader, index: usize) *File.Reader {
110 return &mr.streams.contexts()[index].fr;
111}
112
113pub fn reader(mr: *MultiReader, index: usize) *Io.Reader {
114 return &mr.streams.contexts()[index].fr.interface;
115}
116
117/// Checks for errors in all streams, prioritizing `error.Canceled` if it
118/// occurred anywhere, and ignoring `error.EndOfStream`.
119pub fn checkAnyError(mr: *const MultiReader) UnendingError!void {
120 const contexts = mr.streams.contexts();
121 var other: UnendingError!void = {};
122 for (contexts) |*context| {
123 if (context.err) |err| switch (err) {
124 error.Canceled => |e| return e,
125 error.EndOfStream => continue,
126 else => |e| other = e,
127 };
128 }
129 return other;
130}
131
132pub fn toOwnedSlice(mr: *MultiReader, index: usize) Allocator.Error![]u8 {
133 const gpa = mr.gpa;
134 const r: *Io.Reader = reader(mr, index);
135 if (r.seek == 0) {
136 const new = try gpa.realloc(r.buffer, r.end);
137 r.buffer = &.{};
138 r.end = 0;
139 return new;
140 }
141 const new = try gpa.dupe(u8, r.buffered());
142 gpa.free(r.buffer);
143 r.buffer = &.{};
144 r.seek = 0;
145 r.end = 0;
146 return new;
147}
148
149fn stream(r: *Io.Reader, w: *Io.Writer, limit: Io.Limit) Io.Reader.StreamError!usize {
150 _ = limit;
151 _ = w;
152 const fr: *File.Reader = @alignCast(@fieldParentPtr("interface", r));
153 const context: *Context = @fieldParentPtr("fr", fr);
154 try fillUntimed(context, 1);
155 return 0;
156}
157
158fn discard(r: *Io.Reader, limit: Io.Limit) Io.Reader.Error!usize {
159 _ = limit;
160 const fr: *File.Reader = @alignCast(@fieldParentPtr("interface", r));
161 const context: *Context = @fieldParentPtr("fr", fr);
162 try fillUntimed(context, 1);
163 return 0;
164}
165
166fn readVec(r: *Io.Reader, data: [][]u8) Io.Reader.Error!usize {
167 _ = data;
168 const fr: *File.Reader = @alignCast(@fieldParentPtr("interface", r));
169 const context: *Context = @fieldParentPtr("fr", fr);
170 try fillUntimed(context, 1);
171 return 0;
172}
173
174fn rebase(r: *Io.Reader, capacity: usize) Io.Reader.RebaseError!void {
175 const fr: *File.Reader = @alignCast(@fieldParentPtr("interface", r));
176 const context: *Context = @fieldParentPtr("fr", fr);
177 try fillUntimed(context, capacity);
178}
179
180fn fillUntimed(context: *Context, capacity: usize) Io.Reader.Error!void {
181 fill(context.mr, capacity, .none) catch |err| switch (err) {
182 error.Timeout => unreachable,
183 error.Canceled, error.ConcurrencyUnavailable => |e| {
184 context.err = e;
185 return error.ReadFailed;
186 },
187 error.EndOfStream => |e| return e,
188 };
189 if (context.err) |err| switch (err) {
190 error.EndOfStream => |e| return e,
191 else => return error.ReadFailed,
192 };
193}
194
195pub const FillError = Io.Batch.AwaitConcurrentError || error{
196 /// `fill` was called when all streams already have failed or reached the
197 /// end.
198 EndOfStream,
199};
200
201/// Wait until at least one stream receives more data.
202pub fn fill(mr: *MultiReader, unused_capacity: usize, timeout: Io.Timeout) FillError!void {
203 const contexts = mr.streams.contexts();
204 const io = contexts[0].fr.io;
205 var any_completed = false;
206
207 try mr.batch.awaitConcurrent(io, timeout);
208
209 while (mr.batch.next()) |operation| {
210 any_completed = true;
211 const context = &contexts[operation.index];
212 const n = operation.result.file_read_streaming catch |err| {
213 context.err = err;
214 continue;
215 };
216 const r = &context.fr.interface;
217 r.end += n;
218 if (r.buffer.len - r.end < unused_capacity) {
219 rebaseGrowing(mr, context, r.bufferedLen() + unused_capacity) catch |err| {
220 context.err = err;
221 continue;
222 };
223 assert(r.seek == 0);
224 }
225 context.vec[0] = r.buffer[r.end..];
226 mr.batch.addAt(operation.index, .{ .file_read_streaming = .{
227 .file = context.fr.file,
228 .data = &context.vec,
229 } });
230 }
231
232 if (!any_completed) return error.EndOfStream;
233}
234
235/// Wait until all streams fail or reach the end.
236pub fn fillRemaining(mr: *MultiReader, timeout: Io.Timeout) Io.Batch.AwaitConcurrentError!void {
237 while (fill(mr, 1, timeout)) |_| {} else |err| switch (err) {
238 error.EndOfStream => return,
239 else => |e| return e,
240 }
241}
242
243fn rebaseGrowing(mr: *MultiReader, context: *Context, capacity: usize) Allocator.Error!void {
244 const gpa = mr.gpa;
245 const r = &context.fr.interface;
246 if (r.buffer.len >= capacity) {
247 const data = r.buffer[r.seek..r.end];
248 @memmove(r.buffer[0..data.len], data);
249 r.seek = 0;
250 r.end = data.len;
251 } else {
252 const adjusted_capacity = std.ArrayList(u8).growCapacity(capacity);
253
254 if (r.seek == 0) {
255 if (gpa.remap(r.buffer, adjusted_capacity)) |new_memory| {
256 r.buffer = new_memory;
257 return;
258 }
259 }
260
261 const data = r.buffer[r.seek..r.end];
262 const new = try gpa.alloc(u8, adjusted_capacity);
263 @memcpy(new[0..data.len], data);
264 gpa.free(r.buffer);
265 r.buffer = new;
266 r.seek = 0;
267 r.end = data.len;
268 }
269}