memstream.ha (8861B)
1 // License: MPL-2.0 2 // (c) 2022 Alexey Yerin <yyp@disroot.org> 3 // (c) 2021 Bor Grošelj Simić <bor.groseljsimic@telemach.net> 4 // (c) 2021 Drew DeVault <sir@cmpwn.com> 5 // (c) 2021 Ember Sawady <ecs@d2evs.net> 6 use bytes; 7 use io; 8 use strings; 9 use errors; 10 11 export type memstream = struct { 12 stream: io::stream, 13 buf: []u8, 14 pos: size, 15 }; 16 17 const memstream_vt_r: io::vtable = io::vtable { 18 seeker = &seek, 19 copier = ©, 20 reader = &read, 21 ... 22 }; 23 24 const fixed_vt_w: io::vtable = io::vtable { 25 seeker = &seek, 26 copier = ©, 27 writer = &fixed_write, 28 ... 29 }; 30 31 const fixed_vt_rw: io::vtable = io::vtable { 32 seeker = &seek, 33 copier = ©, 34 reader = &read, 35 writer = &fixed_write, 36 ... 37 }; 38 39 // Creates a stream for a fixed, caller-supplied buffer. All fixed streams are 40 // seekable; seeking a write stream will cause subsequent writes to overwrite 41 // existing contents of the buffer. The program aborts if writes would exceed 42 // the buffer's capacity. The stream doesn't have to be closed. 43 export fn fixed(in: []u8, mode: io::mode) memstream = { 44 let s = memstream { 45 stream = &memstream_vt_r, 46 buf = in, 47 pos = 0, 48 }; 49 if (mode & io::mode::RDWR == io::mode::RDWR) { 50 s.stream = &fixed_vt_rw; 51 } else if (mode & io::mode::WRITE == io::mode::WRITE) { 52 s.stream = &fixed_vt_w; 53 } else if (mode & io::mode::READ == io::mode::READ) { 54 s.stream = &memstream_vt_r; 55 }; 56 return s; 57 }; 58 59 fn fixed_write(s: *io::stream, buf: const []u8) (size | io::error) = { 60 if (len(buf) == 0) { 61 return 0z; 62 }; 63 let s = s: *memstream; 64 if (s.pos >= len(s.buf)) { 65 abort("bufio::fixed buffer exceeded"); 66 }; 67 const n = if (len(buf) > len(s.buf[s.pos..])) { 68 yield len(s.buf[s.pos..]); 69 } else { 70 yield len(buf); 71 }; 72 s.buf[s.pos..s.pos+n] = buf[..n]; 73 s.pos += n; 74 return n; 75 }; 76 77 const dynamic_vt_w: io::vtable = io::vtable { 78 seeker = &seek, 79 copier = ©, 80 writer = &dynamic_write, 81 closer = &dynamic_close, 82 ... 83 }; 84 85 const dynamic_vt_rw: io::vtable = io::vtable { 86 seeker = &seek, 87 copier = ©, 88 reader = &read, 89 writer = &dynamic_write, 90 closer = &dynamic_close, 91 ... 92 }; 93 94 // Creates an [[io::stream]] which dynamically allocates a buffer to store 95 // writes into. Subsequent reads will consume the buffered data. Upon failure to 96 // allocate sufficient memory to store writes, the program aborts. 97 // 98 // Calling [[io::close]] on this stream will free the buffer. If a stream's data 99 // is transferred via [[buffer]], the stream shouldn't be closed as long as the 100 // data is used. 101 export fn dynamic(mode: io::mode) memstream = dynamic_from([], mode); 102 103 // Like [[dynamic]], but takes an existing slice as input. Writes are appended 104 // to it and reads consume bytes from the initial buffer, plus any additional 105 // writes. Like [[dynamic]], calling [[io::close]] will free the buffer. 106 export fn dynamic_from(in: []u8, mode: io::mode) memstream = { 107 let s = memstream { 108 stream = &memstream_vt_r, 109 buf = in, 110 pos = 0, 111 }; 112 if (mode & io::mode::RDWR == io::mode::RDWR) { 113 s.stream = &dynamic_vt_rw; 114 } else if (mode & io::mode::WRITE == io::mode::WRITE) { 115 s.stream = &dynamic_vt_w; 116 } else if (mode & io::mode::READ == io::mode::READ) { 117 s.stream = &memstream_vt_r; 118 }; 119 return s; 120 }; 121 122 // Returns the current buffer of a [[fixed]] or [[dynamic]] stream. 123 export fn buffer(in: *memstream) []u8 = { 124 return in.buf; 125 }; 126 127 // Resets the dynamic buffer's length to zero, but keeps the allocated memory 128 // around for future writes. 129 export fn reset(in: *memstream) void = { 130 in.pos = 0; 131 in.buf = in.buf[..0]; 132 }; 133 134 // Truncates the dynamic buffer, freeing memory associated with it and setting 135 // its length to zero. 136 export fn truncate(in: *memstream) (void | errors::unsupported) = { 137 in.pos = 0; 138 delete(in.buf[..]); 139 }; 140 141 // Reads data from a [[dynamic]] or [[fixed]] stream and returns a slice 142 // borrowed from the internal buffer. 143 export fn borrowedread(st: *memstream, amt: size) ([]u8 | io::EOF) = { 144 if (len(st.buf) - st.pos < amt) { 145 return io::EOF; 146 }; 147 let buf = st.buf[st.pos..st.pos + amt]; 148 st.pos += len(buf); 149 return buf; 150 }; 151 152 fn dynamic_write(s: *io::stream, buf: const []u8) (size | io::error) = { 153 let s = s: *memstream; 154 insert(s.buf[s.pos], buf...); 155 s.pos += len(buf); 156 return len(buf); 157 }; 158 159 fn dynamic_close(s: *io::stream) (void | io::error) = { 160 const s = s: *memstream; 161 free(s.buf); 162 }; 163 164 fn read(s: *io::stream, buf: []u8) (size | io::EOF | io::error) = { 165 let s = s: *memstream; 166 if (len(s.buf) == s.pos && len(buf) != 0) { 167 return io::EOF; 168 }; 169 const n = if (len(s.buf) - s.pos < len(buf)) { 170 yield len(s.buf) - s.pos; 171 } else { 172 yield len(buf); 173 }; 174 assert(s.pos + n <= len(s.buf)); 175 buf[..n] = s.buf[s.pos..s.pos + n]; 176 s.pos += n; 177 return n; 178 }; 179 180 fn seek( 181 s: *io::stream, 182 off: io::off, 183 w: io::whence 184 ) (io::off | io::error) = { 185 let s = s: *memstream; 186 switch (w) { 187 case io::whence::SET => 188 if (len(s.buf) < off: size) { 189 abort("invalid offset"); 190 }; 191 s.pos = off: size; 192 case io::whence::CUR => 193 if (s.pos + off: size > len(s.buf)) { 194 abort("invalid offset"); 195 }; 196 s.pos += off: size; 197 case io::whence::END => 198 if (len(s.buf) - (-off): size < len(s.buf)) { 199 abort("invalid offset"); 200 }; 201 s.pos = len(s.buf) - (-off): size; 202 }; 203 return s.pos: io::off; 204 }; 205 206 fn copy(dest: *io::stream, src: *io::stream) (size | io::error) = { 207 if (src.reader != &read || src.writer == null) { 208 return errors::unsupported; 209 }; 210 let src = src: *memstream; 211 return (dest.writer: *io::writer)(dest, src.buf[src.pos..]); 212 }; 213 214 @test fn dynamic() void = { 215 let s = dynamic(io::mode::RDWR); 216 assert(io::writeall(&s, [1, 2, 3]) as size == 3); 217 assert(bytes::equal(buffer(&s), [1, 2, 3])); 218 assert(io::writeall(&s, [4, 5]) as size == 2); 219 assert(bytes::equal(buffer(&s), [1, 2, 3, 4, 5])); 220 let buf: [2]u8 = [0...]; 221 assert(io::seek(&s, 0, io::whence::SET) as io::off == 0: io::off); 222 assert(io::read(&s, buf[..]) as size == 2 && bytes::equal(buf, [1, 2])); 223 assert(io::read(&s, buf[..]) as size == 2 && bytes::equal(buf, [3, 4])); 224 assert(io::read(&s, buf[..]) as size == 1 && buf[0] == 5); 225 assert(io::read(&s, buf[..]) is io::EOF); 226 assert(io::writeall(&s, [6, 7, 8]) as size == 3); 227 assert(bytes::equal(buffer(&s), [1, 2, 3, 4, 5, 6, 7, 8])); 228 reset(&s); 229 assert(len(buffer(&s)) == 0); 230 assert(io::writeall(&s, [1, 2, 3]) as size == 3); 231 assert(truncate(&s) is void); 232 assert(len(buffer(&s)) == 0); 233 234 let sl: []u8 = alloc([1, 2, 3]); 235 let s = dynamic_from(sl, io::mode::WRITE); 236 assert(io::writeall(&s, [0, 0]) as size == 2); 237 assert(io::seek(&s, 0, io::whence::END) as io::off == 5: io::off); 238 assert(io::writeall(&s, [4, 5, 6]) as size == 3); 239 assert(bytes::equal(buffer(&s), [0, 0, 1, 2, 3, 4, 5, 6])); 240 assert(io::read(&s, buf[..]) as io::error is errors::unsupported); 241 io::close(&s)!; 242 243 sl = alloc([1, 2]); 244 let s = dynamic_from(sl, io::mode::READ); 245 assert(io::read(&s, buf[..1]) as size == 1 && buf[0] == 1); 246 assert(io::seek(&s, 1, io::whence::CUR) as io::off == 2: io::off); 247 assert(io::read(&s, buf[..]) is io::EOF); 248 assert(io::write(&s, [1, 2]) as io::error is errors::unsupported); 249 io::close(&s)!; 250 assert(io::writeall(&s, [1, 2]) as io::error is errors::unsupported); 251 io::close(&s)!; 252 253 let in: [6]u8 = [0, 1, 2, 3, 4, 5]; 254 let source = dynamic_from(in, io::mode::READ); 255 let sink = dynamic(io::mode::WRITE); 256 io::copy(&sink, &source)!; 257 assert(bytes::equal(in, buffer(&sink))); 258 259 let in: [6]u8 = [0, 1, 2, 3, 4, 5]; 260 let source = dynamic_from(in, io::mode::READ); 261 const borrowed = borrowedread(&source, len(in)-1) as []u8; 262 assert(bytes::equal(borrowed, [0, 1, 2, 3, 4])); 263 let source = dynamic_from(in, io::mode::READ); 264 const borrowed = borrowedread(&source, len(in)) as []u8; 265 assert(bytes::equal(borrowed, [0, 1, 2, 3, 4, 5])); 266 let source = dynamic_from(in, io::mode::READ); 267 assert(borrowedread(&source, len(in)+1) is io::EOF); 268 }; 269 270 @test fn fixed() void = { 271 let buf: [1024]u8 = [0...]; 272 let stream = fixed(buf, io::mode::WRITE); 273 defer io::close(&stream)!; 274 275 let n = 0z; 276 n += io::writeall(&stream, strings::toutf8("hello ")) as size; 277 n += io::writeall(&stream, strings::toutf8("world")) as size; 278 assert(bytes::equal(buf[..n], strings::toutf8("hello world"))); 279 assert(io::seek(&stream, 6, io::whence::SET) as io::off == 6: io::off); 280 io::writeall(&stream, strings::toutf8("asdf")) as size; 281 assert(bytes::equal(buf[..n], strings::toutf8("hello asdfd"))); 282 283 let out: [2]u8 = [0...]; 284 let s = fixed([1u8, 2u8], io::mode::READ); 285 defer io::close(&s)!; 286 assert(io::read(&s, out[..1]) as size == 1 && out[0] == 1); 287 assert(io::seek(&s, 1, io::whence::CUR) as io::off == 2: io::off); 288 assert(io::read(&s, buf[..]) is io::EOF); 289 assert(io::writeall(&s, [1, 2]) as io::error is errors::unsupported); 290 291 let in: [6]u8 = [0, 1, 2, 3, 4, 5]; 292 let out: [6]u8 = [0...]; 293 let source = fixed(in, io::mode::READ); 294 let sink = fixed(out, io::mode::WRITE); 295 io::copy(&sink, &source)!; 296 assert(bytes::equal(in, out)); 297 298 assert(io::write(&sink, [])! == 0); 299 };