bytes.bend

bend-kit-bytes@0.2.0.0 · 342 lines · 0 definitions

1# Byte buffers packed four bytes to a U32, with bounds-checked access. Source: https://github.com/paymog/bend-kit/tree/main/bytes
2import 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
9type Bytes is Type:
10  Bytes{len: U32, buf: Array<U32>}
11
12def b8(+x: U32) -> U32:
13  (x .&. 255 : U32)
14
15def shift(+i: U32) -> Nat:
16  U32.to_nat(((i .&. 3 : U32) * 8 : U32))
17
18def words(+n: U32) -> U32:
19  ((n + 3 : U32) >> 2n : U32)
20
21# The fewest d with 2^d >= w.
22def 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
33def depth(+w: U32) -> Nat:
34  depth.go(32n, U32.is_lt(1, w), w, 0n, 1)
35
36def alloc(+n: U32) -> Array<U32>:
37  Array.new(U32, depth(words(n)), 0)
38
39# Unchecked byte read and write. Callers keep i below len.
40def peek.of(+sh: Nat, r: Array<U32> & U32) -> Array<U32> & U32:
41  (a, w) = r
42  (a, b8(U32.shrn(w, sh)))
43
44def peek(a: Array<U32>, +i: U32) -> Array<U32> & U32:
45  peek.of(shift(i), Array.get(U32, a, (i >> 2n : U32)))
46
47def poke.of(+i: U32, +v: U32, r: Array<U32> & U32) -> Array<U32>:
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
52def poke(a: Array<U32>, +i: U32, +v: U32) -> Array<U32>:
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.
56def copy.go(n: Nat, r: Array<U32> & U32, dst: Array<U32>, +s: U32, +d: U32) -> Array<U32> & Array<U32>:
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
65def copy.bytes(+m: U32, src: Array<U32>, dst: Array<U32>, +s: U32, +d: U32) -> Array<U32> & Array<U32>:
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).
69def copy.words(n: Nat, r: Array<U32> & U32, dst: Array<U32>, +s: U32, +d: U32) -> Array<U32> & Array<U32>:
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.
80def copy.shift(n: Nat, +lo: Nat, +hi: Nat, +prev: U32, r: Array<U32> & U32, dst: Array<U32>, +k: U32, +d: U32) -> Array<U32> & Array<U32>:
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
89def copy.shift.at(+w: U32, +lo: U32, +k: U32, +d: U32, r: Array<U32> & U32, dst: Array<U32>) -> Array<U32> & Array<U32>:
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
93def copy.tail(+m: U32, +s: U32, +d: U32, r: Array<U32> & Array<U32>) -> Array<U32> & Array<U32>:
94  (a, b) = r
95  copy.bytes(m, a, b, s, d)
96
97def copy.pick(aligned: Bool, +m: U32, src: Array<U32>, dst: Array<U32>, +s: U32, +d: U32) -> Array<U32> & Array<U32>:
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
108def copy.dst(daligned: Bool, +m: U32, src: Array<U32>, dst: Array<U32>, +s: U32, +d: U32) -> Array<U32> & Array<U32>:
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.
117def copy(+m: U32, src: Array<U32>, dst: Array<U32>, +s: U32, +d: U32) -> Array<U32> & Array<U32>:
118  copy.dst(U32.is_eq((d .&. 3 : U32), 0), m, src, dst, s, d)
119
120# n zero bytes.
121def new(+n: U32) -> Bytes:
122  Bytes{n, alloc(n)}
123
124def length(b: Bytes) -> Bytes & U32:
125  Bytes{+len, buf} = b
126  (Bytes{len, buf}, len)
127
128def get.some(+len: U32, r: Array<U32> & U32) -> Bytes & Maybe<&2, U32>:
129  (a, v) = r
130  (Bytes{len, a}, Some{v})
131
132def get.if(ok: Bool, +len: U32, buf: Array<U32>, +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.
140def 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
144def set.if(ok: Bool, +len: U32, buf: Array<U32>, +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.
152def 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.
157def 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.
165def flush(full: Bool, a: Array<U32>, +k: U32, +w: U32) -> Array<U32>:
166  match full:
167    case True{}:
168      Array.set(U32, a, k, w)
169    case False{}:
170      a
171
172def from.go(s: String, a: Array<U32>, +i: U32, +w: U32) -> Array<U32>:
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.
183def from_string(+s: String) -> Bytes:
184  +n = count(s, 0)
185  Bytes{n, from.go(s, alloc(n), 0, 0)}
186
187def to.go(n: Nat, r: Array<U32> & 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.
196def 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
200def slice.fin(+len: U32, +m: U32, r: Array<U32> & Array<U32>) -> Bytes & Bytes:
201  (a, c) = r
202  (Bytes{len, a}, Bytes{m, c})
203
204def slice.at(+len: U32, buf: Array<U32>, +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.
208def 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
213def append.fin(+n: U32, r: Array<U32> & Array<U32>) -> Bytes:
214  (x, out) = r
215  Bytes{n, out}
216
217def append.two(+n: U32, +la: U32, +lb: U32, bb: Array<U32>, r: Array<U32> & Array<U32>) -> Bytes:
218  (x, out) = r
219  append.fin(n, copy(lb, bb, out, 0, la))
220
221# a then b, in a new buffer.
222def 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
228def concat.total.con(x: Bytes, r: List<&1, Bytes> & U32) -> List<&1, Bytes> & U32:
229  (t, n) = r
230  (Con{x, t}, n)
231
232def 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
240def concat.put(r: Array<U32> & Array<U32>) -> Array<U32>:
241  (src, out) = r
242  out
243
244def concat.go(xs: List<&1, Bytes>, out: Array<U32>, +at: U32) -> Array<U32>:
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
252def 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.
257def 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.
261def at.go(needle: String, ok: Bool, r: Array<U32> & U32, +j: U32) -> Array<U32> & 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
275def at(a: Array<U32>, +needle: String, +j: U32) -> Array<U32> & 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.
279def find.go(f: Nat, r: Array<U32> & Bool, +needle: String, +i: U32) -> Array<U32> & 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
292def find.fin(+len: U32, r: Array<U32> & Maybe<&2, U32>) -> Bytes & Maybe<&2, U32>:
293  (a, x) = r
294  (Bytes{len, a}, x)
295
296def find.at(big: Bool, +len: U32, buf: Array<U32>, +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.
304def 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.
310def eq.go(n: Nat, same: Bool, r: Array<U32> & U32, s: Array<U32> & U32, +i: U32) -> Array<U32> & Array<U32> & 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
327def eq.fin(+la: U32, +lb: U32, r: Array<U32> & Array<U32> & Bool) -> Bytes & Bytes & Bool:
328  (x, y, ok) = r
329  (Bytes{la, x}, Bytes{lb, y}, ok)
330
331def eq.len(same: Bool, +la: U32, xa: Array<U32>, +lb: U32, ya: Array<U32>) -> 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.
339def 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)