1// SPDX-License-Identifier: MIT
2
3package kv
4
5import (
6 "bytes"
7 "iter"
8 "slices"
9)
10
11// The in-memory index is a persistent (copy-on-write) B-tree of order 32:
12// a node holds at most maxEntries entries and, except for the root, at least
13// minEntries. Entries live in internal nodes too.
14const (
15 degree = 32
16 maxEntries = 2*degree - 1
17 minEntries = degree - 1
18)
19
20// entry is a key with its value and the seq of the batch that last wrote it.
21// The tree never copies or modifies key and value: callers pass owned slices.
22type entry struct {
23 key []byte
24 value []byte
25 version uint64
26}
27
28// owner identifies one builder epoch. It has a non-zero size so that every
29// new(owner) is a distinct pointer.
30type owner struct{ _ byte }
31
32// node is a B-tree node. A node without children is a leaf; otherwise it has
33// len(entries)+1 children.
34//
35// Invariant: a node is mutated only by the builder whose current owner token
36// equals node.own. freeze gives the builder a fresh token, so no live builder
37// holds the token of any node reachable from a frozen tree: published nodes
38// are never mutated, and a builder clones a node the first time it touches it.
39type node struct {
40 own *owner
41 entries []entry
42 children []*node
43}
44
45func (n *node) leaf() bool { return len(n.children) == 0 }
46
47// find returns the index of the first entry whose key is >= key and whether
48// that entry's key equals key.
49func (n *node) find(key []byte) (int, bool) {
50 return slices.BinarySearchFunc(n.entries, key, func(e entry, k []byte) int {
51 return bytes.Compare(e.key, k)
52 })
53}
54
55// lookup returns the entry stored under key, if any.
56func lookup(n *node, key []byte) (entry, bool) {
57 for n != nil {
58 i, found := n.find(key)
59 if found {
60 return n.entries[i], true
61 }
62 if n.leaf() {
63 break
64 }
65 n = n.children[i]
66 }
67 return entry{}, false
68}
69
70// tree is an immutable persistent B-tree; the zero value is empty. Its methods
71// never modify existing nodes, so a tree can be shared between goroutines.
72type tree struct {
73 root *node
74 n int
75}
76
77func (t tree) len() int { return t.n }
78
79func (t tree) get(key []byte) (Item, bool) {
80 e, ok := lookup(t.root, key)
81 return Item{Value: e.value, Version: e.version}, ok
82}
83
84// put returns a tree where key maps to value at version; t is untouched.
85func (t tree) put(key, value []byte, version uint64) tree {
86 b := t.builder()
87 b.put(key, value, version)
88 return b.freeze()
89}
90
91// delete returns a tree without key; t is untouched.
92func (t tree) delete(key []byte) tree {
93 b := t.builder()
94 b.delete(key)
95 return b.freeze()
96}
97
98// ascend yields the entries with start <= key < end in ascending order. An
99// empty start and a nil end are unbounded. It walks with an explicit stack and
100// does not allocate per yielded item.
101func (t tree) ascend(start, end []byte) iter.Seq2[[]byte, Item] {
102 return func(yield func([]byte, Item) bool) {
103 // A frame's i is the index of the next entry of n to yield; the
104 // child left of that entry has already been visited or skipped.
105 type frame struct {
106 n *node
107 i int
108 }
109 var buf [16]frame
110 stack := buf[:0]
111
112 // Seek to the first entry >= start.
113 for n := t.root; n != nil; {
114 i, found := n.find(start)
115 stack = append(stack, frame{n, i})
116 if found || n.leaf() {
117 break
118 }
119 n = n.children[i]
120 }
121
122 for len(stack) > 0 {
123 f := &stack[len(stack)-1]
124 if f.i >= len(f.n.entries) {
125 stack = stack[:len(stack)-1]
126 continue
127 }
128 n := f.n
129 e := &n.entries[f.i]
130 f.i++
131 if end != nil && bytes.Compare(e.key, end) >= 0 {
132 return
133 }
134 if !yield(e.key, Item{Value: e.value, Version: e.version}) {
135 return
136 }
137 // Descend to the leftmost leaf of the subtree right of e.
138 for c := n; !c.leaf(); {
139 c = c.children[f.i]
140 stack = append(stack, frame{c, 0})
141 f = &stack[len(stack)-1]
142 }
143 }
144 }
145}
146
147// builder applies many mutations to a tree while copying each touched node at
148// most once. It is not safe for concurrent use. Trees returned by freeze are
149// immutable and independent of later builder operations.
150type builder struct {
151 root *node
152 n int
153 own *owner
154}
155
156// builder returns a builder starting from t; t is untouched.
157func (t tree) builder() *builder {
158 return &builder{root: t.root, n: t.n, own: new(owner)}
159}
160
161// freeze returns the current state as an immutable tree. The builder remains
162// usable: it switches to a new owner token, so nodes of the returned tree are
163// cloned before any further mutation.
164func (b *builder) freeze() tree {
165 b.own = new(owner)
166 return tree{root: b.root, n: b.n}
167}
168
169func (b *builder) len() int { return b.n }
170
171func (b *builder) get(key []byte) (Item, bool) {
172 e, ok := lookup(b.root, key)
173 return Item{Value: e.value, Version: e.version}, ok
174}
175
176// mutable returns n if this builder owns it, else an owned shallow copy.
177func (b *builder) mutable(n *node) *node {
178 if n.own == b.own {
179 return n
180 }
181 c := &node{own: b.own, entries: slices.Clone(n.entries)}
182 if len(n.children) > 0 {
183 c.children = slices.Clone(n.children)
184 }
185 return c
186}
187
188// mutableChild makes n.children[i] owned and returns it. n must be owned.
189func (b *builder) mutableChild(n *node, i int) *node {
190 c := b.mutable(n.children[i])
191 n.children[i] = c
192 return c
193}
194
195// split cuts the owned node n at entry i: n keeps entries[:i], the returned
196// entry is entries[i], and the returned node holds the rest.
197func (b *builder) split(n *node, i int) (entry, *node) {
198 mid := n.entries[i]
199 r := &node{own: b.own, entries: slices.Clone(n.entries[i+1:])}
200 clear(n.entries[i:])
201 n.entries = n.entries[:i]
202 if !n.leaf() {
203 r.children = slices.Clone(n.children[i+1:])
204 clear(n.children[i+1:])
205 n.children = n.children[:i+1]
206 }
207 return mid, r
208}
209
210// put sets key to value at version, replacing any existing entry.
211func (b *builder) put(key, value []byte, version uint64) {
212 e := entry{key: key, value: value, version: version}
213 if b.root == nil {
214 b.root = &node{own: b.own, entries: []entry{e}}
215 b.n = 1
216 return
217 }
218 b.root = b.mutable(b.root)
219 if len(b.root.entries) >= maxEntries {
220 mid, r := b.split(b.root, maxEntries/2)
221 b.root = &node{
222 own: b.own,
223 entries: []entry{mid},
224 children: []*node{b.root, r},
225 }
226 }
227 if !b.insert(b.root, e) {
228 b.n++
229 }
230}
231
232// insert adds e below the owned, non-full node n and reports whether it
233// replaced an existing entry. Full children are split on the way down.
234func (b *builder) insert(n *node, e entry) (replaced bool) {
235 for {
236 i, found := n.find(e.key)
237 if found {
238 n.entries[i] = e
239 return true
240 }
241 if n.leaf() {
242 n.entries = slices.Insert(n.entries, i, e)
243 return false
244 }
245 c := b.mutableChild(n, i)
246 if len(c.entries) >= maxEntries {
247 mid, r := b.split(c, maxEntries/2)
248 n.entries = slices.Insert(n.entries, i, mid)
249 n.children = slices.Insert(n.children, i+1, r)
250 switch cmp := bytes.Compare(e.key, mid.key); {
251 case cmp == 0:
252 n.entries[i] = e
253 return true
254 case cmp > 0:
255 c = r
256 }
257 }
258 n = c
259 }
260}
261
262// delete removes key if present.
263func (b *builder) delete(key []byte) {
264 if _, ok := lookup(b.root, key); !ok {
265 return
266 }
267 b.root = b.mutable(b.root)
268 b.remove(b.root, key, false)
269 b.n--
270 switch {
271 case len(b.root.entries) > 0:
272 case b.root.leaf():
273 b.root = nil
274 default:
275 b.root = b.root.children[0]
276 }
277}
278
279// remove deletes key (or, with rightmost, the largest entry) from the subtree
280// of the owned node n and returns the removed entry. The key must exist: a
281// missing key would delete an unrelated leaf entry. Every child descended into
282// is first guaranteed to hold more than minEntries entries.
283func (b *builder) remove(n *node, key []byte, rightmost bool) entry {
284 for {
285 var i int
286 var found bool
287 switch {
288 case rightmost && n.leaf():
289 i, found = len(n.entries)-1, true
290 case rightmost:
291 i = len(n.entries)
292 default:
293 i, found = n.find(key)
294 }
295 if n.leaf() {
296 e := n.entries[i]
297 n.entries = slices.Delete(n.entries, i, i+1)
298 return e
299 }
300 if len(n.children[i].entries) <= minEntries {
301 b.growChild(n, i)
302 continue // the shape of n changed: search again
303 }
304 c := b.mutableChild(n, i)
305 if found {
306 // Replace the entry by its predecessor.
307 e := n.entries[i]
308 n.entries[i] = b.remove(c, nil, true)
309 return e
310 }
311 n = c
312 }
313}
314
315// growChild ensures n.children[i] has more than minEntries entries by
316// borrowing from a sibling or merging with one. n must be owned.
317func (b *builder) growChild(n *node, i int) {
318 switch {
319 case i > 0 && len(n.children[i-1].entries) > minEntries:
320 c := b.mutableChild(n, i)
321 l := b.mutableChild(n, i-1)
322 last := len(l.entries) - 1
323 c.entries = slices.Insert(c.entries, 0, n.entries[i-1])
324 n.entries[i-1] = l.entries[last]
325 l.entries[last] = entry{}
326 l.entries = l.entries[:last]
327 if !c.leaf() {
328 lc := len(l.children) - 1
329 c.children = slices.Insert(c.children, 0, l.children[lc])
330 l.children[lc] = nil
331 l.children = l.children[:lc]
332 }
333 case i < len(n.entries) && len(n.children[i+1].entries) > minEntries:
334 c := b.mutableChild(n, i)
335 r := b.mutableChild(n, i+1)
336 c.entries = append(c.entries, n.entries[i])
337 n.entries[i] = r.entries[0]
338 r.entries = slices.Delete(r.entries, 0, 1)
339 if !c.leaf() {
340 c.children = append(c.children, r.children[0])
341 r.children = slices.Delete(r.children, 0, 1)
342 }
343 default:
344 // Merge children i and i+1 around entry i (children i-1 and i
345 // around entry i-1 for the last child).
346 if i == len(n.entries) {
347 i--
348 }
349 c := b.mutableChild(n, i)
350 r := n.children[i+1] // only read
351 c.entries = append(c.entries, n.entries[i])
352 c.entries = append(c.entries, r.entries...)
353 c.children = append(c.children, r.children...)
354 n.entries = slices.Delete(n.entries, i, i+1)
355 n.children = slices.Delete(n.children, i+1, i+2)
356 }
357}