src/BitTree.bend
mylsm-lsm-store@0.3.1.0 · 261 lines · 0 definitions
1import Base
2import 0x9ee2e9a299991dcc089fe22c7f3ceb5f/src/containers/bitset.bend as BitWords
3import 0x9ee2e9a299991dcc089fe22c7f3ceb5f/src/math/pow2.bend as Pow2
4
5# Pure indexed bit tree.
6#
7# A non-empty BitTree is a perfect binary tree whose leaves are U32 words.
8# `words` is the logical word count; `capacity` is the power-of-two leaf
9# capacity, and `height` is log2(capacity). Leaves at indexes >= words are
10# padding. A word index follows its `height`-bit binary representation from
11# most-significant bit at the root to least-significant bit at the leaf.
12#
13# The empty value is represented by Empty{} with all metadata zero. Bit
14# indexes outside [0, words * 32) are total: test returns False{} and set is
15# the identity.
16
17type WordTree is Data:
18 Empty{}
19 Leaf{word: U32}
20 Fork{left: WordTree, right: WordTree}
21
22type BitTree is Data:
23 BT{root: WordTree, words: Nat, capacity: Nat, height: Nat}
24
25type Plan is Data:
26 Plan{capacity: Nat, height: Nat}
27
28# Find the least power of two >= need. `fuel = need` is a conservative
29# termination bound; each non-final step doubles capacity.
30def plan_go(fuel: Nat, +need: Nat, +capacity: Nat, +height: Nat, done: Bool) -> Plan:
31 match fuel:
32 case 0n:
33 Plan{capacity, height}
34 case 1n+f:
35 match done:
36 case True{}:
37 Plan{capacity, height}
38 case False{}:
39 +next = Nat.double(capacity)
40 plan_go(f, need, next, Nat.add(height, 1n), Nat.is_le(need, next))
41
42def plan(+need: Nat) -> Plan:
43 match need:
44 case 0n:
45 Plan{0n, 0n}
46 case 1n+n:
47 plan_go(need, need, 1n, 0n, Nat.is_le(need, 1n))
48
49# Independent subtrees are built in parallel.
50def zero_tree(+height: Nat) -> WordTree:
51 match height:
52 case 0n:
53 Leaf{0}
54 case 1n+h:
55 left right = zero_tree(h) zero_tree(h)
56 Fork{left, right}
57
58def new_from_plan(words: Nat, plan: Plan) -> BitTree:
59 match plan:
60 case Plan{capacity, +height}:
61 BT{zero_tree(height), words, capacity, height}
62
63def new(+words: Nat) -> BitTree:
64 match words:
65 case 0n:
66 BT{Empty{}, 0n, 0n, 0n}
67 case 1n+n:
68 new_from_plan(words, plan(words))
69
70# Metadata helpers intended to keep representation laws independent of field
71# layout.
72def word_count(tree: BitTree) -> Nat:
73 match tree:
74 case BT{root, words, capacity, height}:
75 words
76
77def capacity_words(tree: BitTree) -> Nat:
78 match tree:
79 case BT{root, words, capacity, height}:
80 capacity
81
82def tree_height(tree: BitTree) -> Nat:
83 match tree:
84 case BT{root, words, capacity, height}:
85 height
86
87def bit_count(tree: BitTree) -> Nat:
88 Nat.mul(word_count(tree), 32n)
89
90def valid_bit(tree: BitTree, bit: Nat) -> Bool:
91 Nat.is_lt(bit, bit_count(tree))
92
93def same_metadata(+ta: BitTree, +tb: BitTree) -> Bool:
94 Bool.and(
95 Nat.is_eq(word_count(ta), word_count(tb)),
96 Bool.and(
97 Nat.is_eq(capacity_words(ta), capacity_words(tb)),
98 Nat.is_eq(tree_height(ta), tree_height(tb))))
99
100# Build the root-to-leaf path by consuming low bits while prepending them.
101def path_bit(rem: Nat) -> Bool:
102 match rem:
103 case 0n:
104 False{}
105 case 1n+r:
106 True{}
107
108def index_path(height: Nat, +index: Nat, acc: List<&2, Bool>) -> List<&2, Bool>:
109 match height:
110 case 0n:
111 acc
112 case 1n+h:
113 index_path(h, Nat.div(index, 2n), Con{path_bit(Nat.mod(index, 2n)), acc})
114
115def tree_word(path: List<&2, Bool>, tree: WordTree) -> U32:
116 match path tree:
117 case Nil{} Empty{}:
118 0
119 case Nil{} Leaf{word}:
120 word
121 case Nil{} Fork{left, right}:
122 0
123 case Con{side, rest} Empty{}:
124 0
125 case Con{side, rest} Leaf{word}:
126 0
127 case Con{False{}, rest} Fork{left, right}:
128 tree_word(rest, left)
129 case Con{True{}, rest} Fork{left, right}:
130 tree_word(rest, right)
131
132def tree_set_word(path: List<&2, Bool>, value: U32, tree: WordTree) -> WordTree:
133 match path tree:
134 case Nil{} Empty{}:
135 Empty{}
136 case Nil{} Leaf{word}:
137 Leaf{value}
138 case Nil{} Fork{left, right}:
139 Fork{left, right}
140 case Con{side, rest} Empty{}:
141 Empty{}
142 case Con{side, rest} Leaf{word}:
143 Leaf{word}
144 case Con{False{}, rest} Fork{left, right}:
145 Fork{tree_set_word(rest, value, left), right}
146 case Con{True{}, rest} Fork{left, right}:
147 Fork{left, tree_set_word(rest, value, right)}
148
149def word_at_if(inside: Bool, tree: BitTree, index: Nat) -> U32:
150 match inside:
151 case False{}:
152 0
153 case True{}:
154 match tree:
155 case BT{root, words, capacity, height}:
156 tree_word(index_path(height, index, Nil{}), root)
157
158# Safe logical-word lookup; padding leaves are not externally addressable.
159def word_at(+tree: BitTree, +index: Nat) -> U32:
160 word_at_if(Nat.is_lt(index, word_count(tree)), tree, index)
161
162def replace_word_if(inside: Bool, tree: BitTree, index: Nat, value: U32) -> BitTree:
163 match inside:
164 case False{}:
165 tree
166 case True{}:
167 match tree:
168 case BT{root, words, capacity, +height}:
169 BT{tree_set_word(index_path(height, index, Nil{}), value, root), words, capacity, height}
170
171# Safe logical-word update. This helper makes metadata preservation and
172# unaffected-word laws direct to state.
173def replace_word(+tree: BitTree, +index: Nat, value: U32) -> BitTree:
174 replace_word_if(Nat.is_lt(index, word_count(tree)), tree, index, value)
175
176def set_if(inside: Bool, +tree: BitTree, +bit: Nat) -> BitTree:
177 match inside:
178 case False{}:
179 tree
180 case True{}:
181 +index = Nat.div(bit, 32n)
182 +old = word_at(tree, index)
183 replace_word(tree, index, BitWords.word_put(True{}, old, Nat.mod(bit, 32n)))
184
185# Set one logical bit. Out-of-range indexes leave the tree unchanged.
186def set(+tree: BitTree, +bit: Nat) -> BitTree:
187 set_if(valid_bit(tree, bit), tree, bit)
188
189def test_if(inside: Bool, +tree: BitTree, +bit: Nat) -> Bool:
190 match inside:
191 case False{}:
192 False{}
193 case True{}:
194 BitWords.word_get(word_at(tree, Nat.div(bit, 32n)), Nat.mod(bit, 32n))
195
196# Test one logical bit. Out-of-range indexes, including every index in the
197# zero-word tree, return False{}.
198def test(+tree: BitTree, +bit: Nat) -> Bool:
199 test_if(valid_bit(tree, bit), tree, bit)
200
201# Word-wise union for chunked Bloom builds. Both sides always share the exact
202# shape (same height from the same bit count); mismatched arms are unreachable
203# by construction and keep the Fork side.
204def union_bits(ta: WordTree, tb: WordTree) -> WordTree:
205 match ta tb:
206 case Empty{} Empty{}:
207 Empty{}
208 case Empty{} Leaf{wb}:
209 Leaf{wb}
210 case Empty{} Fork{lb, rb}:
211 Fork{lb, rb}
212 case Leaf{wa} Empty{}:
213 Leaf{wa}
214 case Leaf{wa} Leaf{wb}:
215 Leaf{U32.or(wa, wb)}
216 case Leaf{wa} Fork{lb, rb}:
217 Fork{lb, rb}
218 case Fork{la, ra} Empty{}:
219 Fork{la, ra}
220 case Fork{la, ra} Leaf{wb}:
221 Fork{la, ra}
222 case Fork{la, ra} Fork{lb, rb}:
223 Fork{union_bits(la, lb), union_bits(ra, rb)}
224
225# Structural checker for the exact perfect-tree shape described by metadata.
226def shape_ok(+height: Nat, tree: WordTree) -> Bool:
227 match height tree:
228 case 0n Empty{}:
229 False{}
230 case 0n Leaf{word}:
231 True{}
232 case 0n Fork{left, right}:
233 False{}
234 case 1n+h Empty{}:
235 False{}
236 case 1n+h Leaf{word}:
237 False{}
238 case 1n+h Fork{left, right}:
239 a b = shape_ok(h, left) shape_ok(h, right)
240 Bool.and(a, b)
241
242def empty_metadata_ok(words: Nat, capacity: Nat, height: Nat) -> Bool:
243 Bool.and(Nat.is_eq(words, 0n), Bool.and(Nat.is_eq(capacity, 0n), Nat.is_eq(height, 0n)))
244
245def nonempty_metadata_ok(root: WordTree, +words: Nat, +capacity: Nat, +height: Nat) -> Bool:
246 Bool.and(
247 Nat.is_gt(words, 0n),
248 Bool.and(
249 Nat.is_le(words, capacity),
250 Bool.and(
251 Nat.is_eq(capacity, Pow2.pow2t(height)),
252 shape_ok(height, root))))
253
254# Checks zero canonicality, logical bounds, power-of-two capacity, height, and
255# exact tree shape. It does not require padding words to remain zero.
256def well_formed(tree: BitTree) -> Bool:
257 match tree:
258 case BT{Empty{}, words, capacity, height}:
259 empty_metadata_ok(words, capacity, height)
260 case BT{root, words, capacity, height}:
261 nonempty_metadata_ok(root, words, capacity, height)