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