1 # Byte buffers packed four bytes to a U32, with bounds-checked access. Source: https://github.com/paymog/bend-kit/tree/main/bytes 2 import Base 3 4 # Bytes{len, buf}: byte i is bits 8*(i%4) of buf[i/4]. buf has the fewest 2^d words 5 # that hold len bytes. Every byte at or past len is 0, so equal buffers have equal words. 6 # Out-of-range reads answer None and out-of-range writes do nothing; Array alone would wrap. 7 # import ./bytes/bytes.bend as Bytes 8 9 type Bytes is Type: 10 Bytes{len: U32, buf: Array} 11 12 def b8(+x: U32) -> U32: 13 (x .&. 255 : U32) 14 15 def shift(+i: U32) -> Nat: 16 U32.to_nat(((i .&. 3 : U32) * 8 : U32)) 17 18 def words(+n: U32) -> U32: 19 ((n + 3 : U32) >> 2n : U32) 20 21 # The fewest d with 2^d >= w. 22 def depth.go(f: Nat, more: Bool, +w: U32, +d: Nat, +cap: U32) -> Nat: 23 match f: 24 case 0n: 25 d 26 case 1n+p: 27 match more: 28 case False{}: 29 d 30 case True{}: 31 depth.go(p, U32.is_lt((cap * 2 : U32), w), w, 1n+d, (cap * 2 : U32)) 32 33 def depth(+w: U32) -> Nat: 34 depth.go(32n, U32.is_lt(1, w), w, 0n, 1) 35 36 def alloc(+n: U32) -> Array: 37 Array.new(U32, depth(words(n)), 0) 38 39 # Unchecked byte read and write. Callers keep i below len. 40 def peek.of(+sh: Nat, r: Array & U32) -> Array & U32: 41 (a, w) = r 42 (a, b8(U32.shrn(w, sh))) 43 44 def peek(a: Array, +i: U32) -> Array & U32: 45 peek.of(shift(i), Array.get(U32, a, (i >> 2n : U32))) 46 47 def poke.of(+i: U32, +v: U32, r: Array & U32) -> Array: 48 (a, w) = r 49 +sh = shift(i) 50 Array.set(U32, a, (i >> 2n : U32), ((w .&. ((255 << sh) .^. 4294967295)) .|. (b8(v) << sh) : U32)) 51 52 def poke(a: Array, +i: U32, +v: U32) -> Array: 53 poke.of(i, v, Array.get(U32, a, (i >> 2n : U32))) 54 55 # n bytes from src[s..] to dst[d..], one byte at a time. 56 def copy.go(n: Nat, r: Array & U32, dst: Array, +s: U32, +d: U32) -> Array & Array: 57 match n: 58 case 0n: 59 (a, v) = r 60 (a, dst) 61 case 1n+p: 62 (a, v) = r 63 copy.go(p, peek(a, (s + 1 : U32)), poke(dst, d, v), (s + 1 : U32), (d + 1 : U32)) 64 65 def copy.bytes(+m: U32, src: Array, dst: Array, +s: U32, +d: U32) -> Array & Array: 66 copy.go(U32.to_nat(m), peek(src, s), dst, s, d) 67 68 # n words from src[s..] to dst[d..] (word indexes). 69 def copy.words(n: Nat, r: Array & U32, dst: Array, +s: U32, +d: U32) -> Array & Array: 70 match n: 71 case 0n: 72 (a, w) = r 73 (a, dst) 74 case 1n+p: 75 (a, +w) = r 76 copy.words(p, Array.get(U32, a, (s + 1 : U32)), Array.set(U32, dst, d, w), (s + 1 : U32), (d + 1 : U32)) 77 78 # m bytes to a word-aligned d from an unaligned s: each out word joins two source words. 79 # lo = 8*(s%4) and hi = 32 - lo; r holds source word k, prev the word before it. 80 def copy.shift(n: Nat, +lo: Nat, +hi: Nat, +prev: U32, r: Array & U32, dst: Array, +k: U32, +d: U32) -> Array & Array: 81 match n: 82 case 0n: 83 (a, w) = r 84 (a, dst) 85 case 1n+p: 86 (a, +w) = r 87 copy.shift(p, lo, hi, w, Array.get(U32, a, (k + 1 : U32)), Array.set(U32, dst, d, (U32.shrn(prev, lo) .|. U32.shln(w, hi) : U32)), (k + 1 : U32), (d + 1 : U32)) 88 89 def copy.shift.at(+w: U32, +lo: U32, +k: U32, +d: U32, r: Array & U32, dst: Array) -> Array & Array: 90 (a, prev) = r 91 copy.shift(U32.to_nat(w), U32.to_nat(lo), U32.to_nat((32 - lo : U32)), prev, Array.get(U32, a, (k + 1 : U32)), dst, (k + 1 : U32), d) 92 93 def copy.tail(+m: U32, +s: U32, +d: U32, r: Array & Array) -> Array & Array: 94 (a, b) = r 95 copy.bytes(m, a, b, s, d) 96 97 def copy.pick(aligned: Bool, +m: U32, src: Array, dst: Array, +s: U32, +d: U32) -> Array & Array: 98 match aligned: 99 case True{}: 100 +w = (m >> 2n : U32) 101 +k = (w * 4 : U32) 102 copy.tail((m - k : U32), (s + k : U32), (d + k : U32), copy.words(U32.to_nat(w), Array.get(U32, src, (s >> 2n : U32)), dst, (s >> 2n : U32), (d >> 2n : U32))) 103 case False{}: 104 +w = (m >> 2n : U32) 105 +k = (w * 4 : U32) 106 copy.tail((m - k : U32), (s + k : U32), (d + k : U32), copy.shift.at(w, ((s .&. 3 : U32) * 8 : U32), (s >> 2n : U32), (d >> 2n : U32), Array.get(U32, src, (s >> 2n : U32)), dst)) 107 108 def copy.dst(daligned: Bool, +m: U32, src: Array, dst: Array, +s: U32, +d: U32) -> Array & Array: 109 match daligned: 110 case False{}: 111 copy.bytes(m, src, dst, s, d) 112 case True{}: 113 copy.pick(U32.is_eq((s .&. 3 : U32), 0), m, src, dst, s, d) 114 115 # m bytes from src[s..] to dst[d..]. Whole words when d is word-aligned, shifted when s is not. 116 # ponytail: an unaligned d (append after an odd length) copies bytes; merge into d's first word if that gets hot. 117 def copy(+m: U32, src: Array, dst: Array, +s: U32, +d: U32) -> Array & Array: 118 copy.dst(U32.is_eq((d .&. 3 : U32), 0), m, src, dst, s, d) 119 120 # n zero bytes. 121 def new(+n: U32) -> Bytes: 122 Bytes{n, alloc(n)} 123 124 def length(b: Bytes) -> Bytes & U32: 125 Bytes{+len, buf} = b 126 (Bytes{len, buf}, len) 127 128 def get.some(+len: U32, r: Array & U32) -> Bytes & Maybe<&2, U32>: 129 (a, v) = r 130 (Bytes{len, a}, Some{v}) 131 132 def get.if(ok: Bool, +len: U32, buf: Array, +i: U32) -> Bytes & Maybe<&2, U32>: 133 match ok: 134 case True{}: 135 get.some(len, peek(buf, i)) 136 case False{}: 137 (Bytes{len, buf}, None{}) 138 139 # Byte i, or None when i >= len. 140 def get(b: Bytes, +i: U32) -> Bytes & Maybe<&2, U32>: 141 Bytes{+len, buf} = b 142 get.if(U32.is_lt(i, len), len, buf, i) 143 144 def set.if(ok: Bool, +len: U32, buf: Array, +i: U32, +v: U32) -> Bytes: 145 match ok: 146 case True{}: 147 Bytes{len, poke(buf, i, v)} 148 case False{}: 149 Bytes{len, buf} 150 151 # Byte i becomes v & 255. Nothing changes when i >= len. 152 def set(b: Bytes, +i: U32, +v: U32) -> Bytes: 153 Bytes{+len, buf} = b 154 set.if(U32.is_lt(i, len), len, buf, i, v) 155 156 # String.length counts in Nat, which costs more than the walk itself. 157 def count(s: String, +n: U32) -> U32: 158 match s: 159 case SNil{}: 160 n 161 case SCon{c, t}: 162 count(t, (n + 1 : U32)) 163 164 # Bytes enter at the top of w and shift down, so a full word has byte 0 lowest. 165 def flush(full: Bool, a: Array, +k: U32, +w: U32) -> Array: 166 match full: 167 case True{}: 168 Array.set(U32, a, k, w) 169 case False{}: 170 a 171 172 def from.go(s: String, a: Array, +i: U32, +w: U32) -> Array: 173 match s: 174 case SNil{}: 175 +r = (i .&. 3 : U32) 176 flush(U32.is_ne(r, 0), a, (i >> 2n : U32), U32.shrn(w, U32.to_nat((((4 - r) .&. 3) * 8 : U32)))) 177 case SCon{Chr{+c}, t}: 178 +w2 = ((w >> 8n) .|. (b8(c) << 24n) : U32) 179 +full = U32.is_eq((i .&. 3 : U32), 3) 180 from.go(t, flush(full, a, (i >> 2n : U32), w2), (i + 1 : U32), Bool.pick(U32, full, 0, w2)) 181 182 # A byte string (one Char per octet, as Wire and Http use) to Bytes. Each Char keeps its low 8 bits. 183 def from_string(+s: String) -> Bytes: 184 +n = count(s, 0) 185 Bytes{n, from.go(s, alloc(n), 0, 0)} 186 187 def to.go(n: Nat, r: Array & U32, +i: U32, acc: String) -> String: 188 match n: 189 case 0n: 190 acc 191 case 1n+p: 192 (a, v) = r 193 to.go(p, peek(a, (i - 1 : U32)), (i - 1 : U32), SCon{Chr{v}, acc}) 194 195 # Bytes to a byte string, one Char per octet. 196 def to_string(b: Bytes) -> String: 197 Bytes{+len, buf} = b 198 to.go(U32.to_nat(len), peek(buf, (len - 1 : U32)), (len - 1 : U32), SNil{}) 199 200 def slice.fin(+len: U32, +m: U32, r: Array & Array) -> Bytes & Bytes: 201 (a, c) = r 202 (Bytes{len, a}, Bytes{m, c}) 203 204 def slice.at(+len: U32, buf: Array, +s: U32, +m: U32) -> Bytes & Bytes: 205 slice.fin(len, m, copy(m, buf, alloc(m), s, 0)) 206 207 # The buffer back, and a copy of up to n bytes from start. Both ends are clamped to len. 208 def slice(b: Bytes, +start: U32, +n: U32) -> Bytes & Bytes: 209 Bytes{+len, buf} = b 210 +s = U32.min(start, len) 211 slice.at(len, buf, s, U32.min(n, (len - s : U32))) 212 213 def append.fin(+n: U32, r: Array & Array) -> Bytes: 214 (x, out) = r 215 Bytes{n, out} 216 217 def append.two(+n: U32, +la: U32, +lb: U32, bb: Array, r: Array & Array) -> Bytes: 218 (x, out) = r 219 append.fin(n, copy(lb, bb, out, 0, la)) 220 221 # a then b, in a new buffer. 222 def append(a: Bytes, b: Bytes) -> Bytes: 223 Bytes{+la, ba} = a 224 Bytes{+lb, bb} = b 225 +n = (la + lb : U32) 226 append.two(n, la, lb, bb, copy(la, ba, alloc(n), 0, 0)) 227 228 def concat.total.con(x: Bytes, r: List<&1, Bytes> & U32) -> List<&1, Bytes> & U32: 229 (t, n) = r 230 (Con{x, t}, n) 231 232 def concat.total(xs: List<&1, Bytes>, +n: U32) -> List<&1, Bytes> & U32: 233 match xs: 234 case Nil{}: 235 (Nil{}, n) 236 case Con{x, t}: 237 Bytes{+l, buf} = x 238 concat.total.con(Bytes{l, buf}, concat.total(t, (n + l : U32))) 239 240 def concat.put(r: Array & Array) -> Array: 241 (src, out) = r 242 out 243 244 def concat.go(xs: List<&1, Bytes>, out: Array, +at: U32) -> Array: 245 match xs: 246 case Nil{}: 247 out 248 case Con{x, t}: 249 Bytes{+l, buf} = x 250 concat.go(t, concat.put(copy(l, buf, out, 0, at)), (at + l : U32)) 251 252 def concat.of(r: List<&1, Bytes> & U32) -> Bytes: 253 (xs, +n) = r 254 Bytes{n, concat.go(xs, alloc(n), 0)} 255 256 # The pieces in order, in one new buffer: one copy per byte, however many pieces. 257 def concat(xs: List<&1, Bytes>) -> Bytes: 258 concat.of(concat.total(xs, 0)) 259 260 # Does needle match at j? ok is the previous byte's result; the first mismatch stops the walk. 261 def at.go(needle: String, ok: Bool, r: Array & U32, +j: U32) -> Array & Bool: 262 match needle: 263 case SNil{}: 264 (a, v) = r 265 (a, ok) 266 case SCon{Chr{+c}, t}: 267 match ok: 268 case False{}: 269 (a, v) = r 270 (a, False{}) 271 case True{}: 272 (a, +v) = r 273 at.go(t, U32.is_eq(v, c), peek(a, (j + 1 : U32)), (j + 1 : U32)) 274 275 def at(a: Array, +needle: String, +j: U32) -> Array & Bool: 276 at.go(needle, True{}, peek(a, j), j) 277 278 # ponytail: naive search; a mismatch costs one byte read per position, but a repetitive needle is O(len * needle). Add a SWAR scan or two-way search if find gets hot. 279 def find.go(f: Nat, r: Array & Bool, +needle: String, +i: U32) -> Array & Maybe<&2, U32>: 280 match f: 281 case 0n: 282 (a, hit) = r 283 (a, Bool.pick(Maybe<&2, U32>, hit, Some{i}, None{})) 284 case 1n+p: 285 (a, hit) = r 286 match hit: 287 case True{}: 288 (a, Some{i}) 289 case False{}: 290 find.go(p, at(a, needle, (i + 1 : U32)), needle, (i + 1 : U32)) 291 292 def find.fin(+len: U32, r: Array & Maybe<&2, U32>) -> Bytes & Maybe<&2, U32>: 293 (a, x) = r 294 (Bytes{len, a}, x) 295 296 def find.at(big: Bool, +len: U32, buf: Array, +needle: String, +m: U32) -> Bytes & Maybe<&2, U32>: 297 match big: 298 case True{}: 299 (Bytes{len, buf}, None{}) 300 case False{}: 301 find.fin(len, find.go(U32.to_nat((len - m : U32)), at(buf, needle, 0), needle, 0)) 302 303 # Index of the first match of needle (a byte string), or None. An empty needle matches at 0. 304 def find(b: Bytes, +needle: String) -> Bytes & Maybe<&2, U32>: 305 Bytes{+len, buf} = b 306 +m = count(needle, 0) 307 find.at(U32.is_lt(len, m), len, buf, needle, m) 308 309 # Compares whole words; same tells whether the previous pair matched. 310 def eq.go(n: Nat, same: Bool, r: Array & U32, s: Array & U32, +i: U32) -> Array & Array & Bool: 311 match n: 312 case 0n: 313 (x, v) = r 314 (y, w) = s 315 (x, y, same) 316 case 1n+p: 317 match same: 318 case False{}: 319 (x, v) = r 320 (y, w) = s 321 (x, y, False{}) 322 case True{}: 323 (x, +v) = r 324 (y, +w) = s 325 eq.go(p, U32.is_eq(v, w), Array.get(U32, x, (i + 1 : U32)), Array.get(U32, y, (i + 1 : U32)), (i + 1 : U32)) 326 327 def eq.fin(+la: U32, +lb: U32, r: Array & Array & Bool) -> Bytes & Bytes & Bool: 328 (x, y, ok) = r 329 (Bytes{la, x}, Bytes{lb, y}, ok) 330 331 def eq.len(same: Bool, +la: U32, xa: Array, +lb: U32, ya: Array) -> Bytes & Bytes & Bool: 332 match same: 333 case False{}: 334 (Bytes{la, xa}, Bytes{lb, ya}, False{}) 335 case True{}: 336 eq.fin(la, lb, eq.go(U32.to_nat(words(la)), True{}, Array.get(U32, xa, 0), Array.get(U32, ya, 0), 0)) 337 338 # Both buffers back, and whether they hold the same bytes. 339 def eq(a: Bytes, b: Bytes) -> Bytes & Bytes & Bool: 340 Bytes{+la, xa} = a 341 Bytes{+lb, ya} = b 342 eq.len(U32.is_eq(la, lb), la, xa, lb, ya)