| ... | ... | @@ -0,0 +1,318 @@ |
| 1 | // FIFO of fixed size items |
| 2 | // Usually used for e.g. byte buffers |
| 3 | |
| 4 | const std = @import("std"); |
| 5 | const math = std.math; |
| 6 | const mem = std.mem; |
| 7 | const Allocator = mem.Allocator; |
| 8 | const debug = std.debug; |
| 9 | const assert = debug.assert; |
| 10 | const testing = std.testing; |
| 11 | |
| 12 | pub fn FixedSizeFifo(comptime T: type) type { |
| 13 | return struct { |
| 14 | allocator: *Allocator, |
| 15 | buf: []u8, |
| 16 | head: usize, |
| 17 | count: usize, |
| 18 | |
| 19 | const Self = @This(); |
| 20 | |
| 21 | pub fn init(allocator: *Allocator) Self { |
| 22 | return Self{ |
| 23 | .allocator = allocator, |
| 24 | .buf = [_]T{}, |
| 25 | .head = 0, |
| 26 | .count = 0, |
| 27 | }; |
| 28 | } |
| 29 | |
| 30 | pub fn deinit(self: *Self) void { |
| 31 | self.allocator.free(self.buf); |
| 32 | self.* = undefined; |
| 33 | } |
| 34 | |
| 35 | pub fn realign(self: *Self) void { |
| 36 | if (self.buf.len - self.head >= self.count) { |
| 37 | // this copy overlaps |
| 38 | mem.copy(T, self.buf[0..self.count], self.buf[self.head..][0..self.count]); |
| 39 | self.head = 0; |
| 40 | } else { |
| 41 | var tmp: [mem.page_size / 2 / @sizeOf(T)]T = undefined; |
| 42 | |
| 43 | while (self.head != 0) { |
| 44 | const n = math.min(self.head, tmp.len); |
| 45 | const m = self.buf.len - n; |
| 46 | mem.copy(T, tmp[0..n], self.buf[0..n]); |
| 47 | // this middle copy overlaps; the others here don't |
| 48 | mem.copy(T, self.buf[0..m], self.buf[n..][0..m]); |
| 49 | mem.copy(T, self.buf[m..], tmp[0..n]); |
| 50 | self.head -= n; |
| 51 | } |
| 52 | } |
| 53 | { // set unused area to undefined |
| 54 | const unused = @sliceToBytes(self.buf[self.count..]); |
| 55 | @memset(unused.ptr, undefined, unused.len); |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | /// Reduce allocated capacity to `size`. |
| 60 | pub fn shrink(self: *Self, size: usize) void { |
| 61 | assert(size >= self.count); |
| 62 | self.realign(); |
| 63 | self.buf = self.allocator.realloc(self.buf, size) catch |e| switch (e) { |
| 64 | error.OutOfMemory => return, // no problem, capacity is still correct then. |
| 65 | }; |
| 66 | } |
| 67 | |
| 68 | /// Ensure that the buffer can fit at least `size` items |
| 69 | pub fn ensureCapacity(self: *Self, size: usize) error{OutOfMemory}!void { |
| 70 | if (self.buf.len >= size) return; |
| 71 | self.realign(); |
| 72 | const new_size = math.ceilPowerOfTwo(usize, size) catch return error.OutOfMemory; |
| 73 | self.buf = try self.allocator.realloc(self.buf, new_size); |
| 74 | } |
| 75 | |
| 76 | /// Makes sure at least `size` items are unused |
| 77 | pub fn ensureUnusedCapacity(self: *Self, size: usize) error{OutOfMemory}!void { |
| 78 | if (self.writableLength() >= size) return; |
| 79 | |
| 80 | return try self.ensureCapacity(math.add(usize, self.count, size) catch return error.OutOfMemory); |
| 81 | } |
| 82 | |
| 83 | /// Returns number of items currently in fifo |
| 84 | pub fn readableLength(self: Self) usize { |
| 85 | return self.count; |
| 86 | } |
| 87 | |
| 88 | /// Returns a writable slice from the 'read' end of the fifo |
| 89 | fn readableSliceMut(self: Self, offset: usize) []T { |
| 90 | if (offset > self.count) return [_]T{}; |
| 91 | |
| 92 | const start = self.head + offset; |
| 93 | if (start >= self.buf.len) { |
| 94 | return self.buf[start - self.buf.len ..][0 .. self.count - offset]; |
| 95 | } else { |
| 96 | const end: usize = self.head + self.count; |
| 97 | if (end >= self.buf.len) { |
| 98 | return self.buf[start..self.buf.len]; |
| 99 | } else { |
| 100 | return self.buf[start..end]; |
| 101 | } |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | /// Returns a readable slice from `offset` |
| 106 | pub fn readableSlice(self: Self, offset: usize) []const T { |
| 107 | return self.readableSliceMut(offset); |
| 108 | } |
| 109 | |
| 110 | const autoalign = false; |
| 111 | |
| 112 | /// Discard first `count` bytes of readable data |
| 113 | pub fn discard(self: *Self, count: usize) void { |
| 114 | assert(count <= self.count); |
| 115 | { // set old range to undefined. Note: may be wrapped around |
| 116 | const slice = self.readableSliceMut(0); |
| 117 | if (slice.len >= count) { |
| 118 | const unused = @sliceToBytes(slice[0..count]); |
| 119 | @memset(unused.ptr, undefined, unused.len); |
| 120 | } else { |
| 121 | const unused = @sliceToBytes(slice[0..]); |
| 122 | @memset(unused.ptr, undefined, unused.len); |
| 123 | const unused2 = @sliceToBytes(self.readableSliceMut(slice.len)[0 .. count - slice.len]); |
| 124 | @memset(unused2.ptr, undefined, unused2.len); |
| 125 | } |
| 126 | } |
| 127 | self.head = (self.head + count) % self.buf.len; |
| 128 | self.count -= count; |
| 129 | if (autoalign and self.count == 0) |
| 130 | self.head = 0; |
| 131 | } |
| 132 | |
| 133 | /// Read the next item from the fifo |
| 134 | pub fn readItem(self: *Self) !T { |
| 135 | if (self.count == 0) return error.EndOfStream; |
| 136 | |
| 137 | const c = self.buf[self.head]; |
| 138 | self.discard(1); |
| 139 | return c; |
| 140 | } |
| 141 | |
| 142 | /// Read data from the fifo into `dst`, returns slice of bytes copied (subslice of `dst`) |
| 143 | pub fn read(self: *Self, dst: []T) []T { |
| 144 | var dst_left = dst; |
| 145 | |
| 146 | while (dst_left.len > 0) { |
| 147 | const slice = self.readableSlice(0); |
| 148 | if (slice.len == 0) break; |
| 149 | const n = math.min(slice.len, dst_left.len); |
| 150 | mem.copy(T, dst_left, slice[0..n]); |
| 151 | self.discard(n); |
| 152 | dst_left = dst_left[n..]; |
| 153 | } |
| 154 | |
| 155 | return dst[0 .. dst.len - dst_left.len]; |
| 156 | } |
| 157 | |
| 158 | /// Returns number of bytes available in fifo |
| 159 | pub fn writableLength(self: Self) usize { |
| 160 | return self.buf.len - self.count; |
| 161 | } |
| 162 | |
| 163 | /// Returns the first section of writable buffer |
| 164 | /// Note that this may be of length 0 |
| 165 | pub fn writableSlice(self: Self, offset: usize) []T { |
| 166 | if (offset > self.buf.len) return [_]T{}; |
| 167 | |
| 168 | const tail = self.head + offset + self.count; |
| 169 | if (tail < self.buf.len) { |
| 170 | return self.buf[tail..]; |
| 171 | } else { |
| 172 | return self.buf[tail - self.buf.len ..][0 .. self.writableLength() - offset]; |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | /// Returns a writable buffer of at least `size` bytes, allocating memory as needed. |
| 177 | /// Use `fifo.update` once you've written data to it. |
| 178 | pub fn writeableWithSize(self: *Self, size: usize) ![]T { |
| 179 | try self.ensureUnusedCapacity(size); |
| 180 | |
| 181 | // try to avoid realigning buffer |
| 182 | var slice = self.writableSlice(0); |
| 183 | if (slice.len < size) { |
| 184 | self.realign(); |
| 185 | slice = self.writableSlice(0); |
| 186 | } |
| 187 | return slice; |
| 188 | } |
| 189 | |
| 190 | /// Update the tail location of the buffer (usually follows use of writable/writeableWithSize) |
| 191 | pub fn update(self: *Self, count: usize) void { |
| 192 | assert(self.count + count <= self.buf.len); |
| 193 | self.count += count; |
| 194 | } |
| 195 | |
| 196 | /// Appends the data in `src` to the fifo. You must |
| 197 | pub fn writeAssumeCapacity(self: *Self, src: []const T) void { |
| 198 | assert(self.writableLength() >= src.len); |
| 199 | |
| 200 | var src_left = src; |
| 201 | while (src_left.len > 0) { |
| 202 | const writable_slice = self.writableSlice(0); |
| 203 | assert(writable_slice.len != 0); |
| 204 | const n = math.min(writable_slice.len, src_left.len); |
| 205 | mem.copy(T, writable_slice, src_left[0..n]); |
| 206 | self.update(n); |
| 207 | src_left = src_left[n..]; |
| 208 | } |
| 209 | } |
| 210 | |
| 211 | /// Appends the data in `src` to the fifo. |
| 212 | /// Allocates more memory as necessary |
| 213 | pub fn write(self: *Self, src: []const T) !void { |
| 214 | try self.ensureUnusedCapacity(src.len); |
| 215 | |
| 216 | return self.writeAssumeCapacity(src); |
| 217 | } |
| 218 | |
| 219 | pub fn print(self: *Self, comptime format: []const u8, args: ...) !void { |
| 220 | return std.fmt.format(self, error{OutOfMemory}, Self.write, format, args); |
| 221 | } |
| 222 | |
| 223 | /// Make `count` bytes available before the current read location |
| 224 | fn rewind(self: *Self, size: usize) void { |
| 225 | assert(self.writableLength() >= size); |
| 226 | |
| 227 | self.head = (self.head + (self.buf.len - size)) % self.buf.len; |
| 228 | self.count += size; |
| 229 | } |
| 230 | |
| 231 | /// Place data back into the read stream |
| 232 | pub fn unget(self: *Self, src: []const T) !void { |
| 233 | try self.ensureUnusedCapacity(src.len); |
| 234 | |
| 235 | self.rewind(src.len); |
| 236 | |
| 237 | const slice = self.readableSliceMut(0); |
| 238 | mem.copy(T, slice, src[0..slice.len]); |
| 239 | const slice2 = self.readableSliceMut(slice.len); |
| 240 | mem.copy(T, slice2, src[slice.len..]); |
| 241 | } |
| 242 | |
| 243 | /// Peek at the item at `offset` |
| 244 | pub fn peekItem(self: Self, offset: usize) error{EndOfStream}!T { |
| 245 | if (offset >= self.count) |
| 246 | return error.EndOfStream; |
| 247 | |
| 248 | return self.buf[(self.head + offset) % self.buf.len]; |
| 249 | } |
| 250 | }; |
| 251 | } |
| 252 | |
| 253 | const ByteFifo = FixedSizeFifo(u8); |
| 254 | |
| 255 | test "ByteFifo" { |
| 256 | var fifo = ByteFifo.init(debug.global_allocator); |
| 257 | defer fifo.deinit(); |
| 258 | |
| 259 | try fifo.write("HELLO"); |
| 260 | testing.expectEqual(usize(5), fifo.readableLength()); |
| 261 | testing.expectEqualSlices(u8, "HELLO", fifo.readableSlice(0)); |
| 262 | |
| 263 | { |
| 264 | var i: usize = 0; |
| 265 | while (i < 5) : (i += 1) { |
| 266 | try fifo.write([_]u8{try fifo.peekItem(i)}); |
| 267 | } |
| 268 | testing.expectEqual(usize(10), fifo.readableLength()); |
| 269 | testing.expectEqualSlices(u8, "HELLOHELLO", fifo.readableSlice(0)); |
| 270 | } |
| 271 | |
| 272 | { |
| 273 | testing.expectEqual(u8('H'), try fifo.readItem()); |
| 274 | testing.expectEqual(u8('E'), try fifo.readItem()); |
| 275 | testing.expectEqual(u8('L'), try fifo.readItem()); |
| 276 | testing.expectEqual(u8('L'), try fifo.readItem()); |
| 277 | testing.expectEqual(u8('O'), try fifo.readItem()); |
| 278 | } |
| 279 | testing.expectEqual(usize(5), fifo.readableLength()); |
| 280 | |
| 281 | { // Writes that wrap around |
| 282 | testing.expectEqual(usize(11), fifo.writableLength()); |
| 283 | testing.expectEqual(usize(6), fifo.writableSlice(0).len); |
| 284 | fifo.writeAssumeCapacity("6<chars<11"); |
| 285 | testing.expectEqualSlices(u8, "HELLO6<char", fifo.readableSlice(0)); |
| 286 | testing.expectEqualSlices(u8, "s<11", fifo.readableSlice(11)); |
| 287 | fifo.discard(11); |
| 288 | testing.expectEqualSlices(u8, "s<11", fifo.readableSlice(0)); |
| 289 | fifo.discard(4); |
| 290 | testing.expectEqual(usize(0), fifo.readableLength()); |
| 291 | } |
| 292 | |
| 293 | { |
| 294 | const buf = try fifo.writeableWithSize(12); |
| 295 | testing.expectEqual(usize(12), buf.len); |
| 296 | var i: u8 = 0; |
| 297 | while (i < 10) : (i += 1) { |
| 298 | buf[i] = i + 'a'; |
| 299 | } |
| 300 | fifo.update(10); |
| 301 | testing.expectEqualSlices(u8, "abcdefghij", fifo.readableSlice(0)); |
| 302 | } |
| 303 | |
| 304 | { |
| 305 | try fifo.unget("prependedstring"); |
| 306 | var result: [30]u8 = undefined; |
| 307 | testing.expectEqualSlices(u8, "prependedstringabcdefghij", fifo.read(&result)); |
| 308 | } |
| 309 | |
| 310 | fifo.shrink(0); |
| 311 | |
| 312 | { |
| 313 | try fifo.print("{}, {}!", "Hello", "World"); |
| 314 | var result: [30]u8 = undefined; |
| 315 | testing.expectEqualSlices(u8, "Hello, World!", fifo.read(&result)); |
| 316 | testing.expectEqual(usize(0), fifo.readableLength()); |
| 317 | } |
| 318 | } |
daurnimator <quae@daurnimator.com> 2019-11-05 07:05:29+11:00
Andrew Kelley <andrew@ziglang.org> 2019-11-04 17:44:01-05:00