src/BitTree.bend

mylsm-lsm-store@0.6.0.1 · 261 lines · 28 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 bit count in the packed bit-tree operations.
84def bit_count(tree: BitTree) -> Nat:
85  Nat.mul(word_count(tree), 32n)
86
87# Check whether bit holds for the packed bit-tree operations.
88def valid_bit(tree: BitTree, bit: Nat) -> Bool:
89  Nat.is_lt(bit, bit_count(tree))
90
91# Build the root-to-leaf path by consuming low bits while prepending them.
92def path_bit(rem: Nat) -> Bool:
93  match rem:
94    case 0n:
95      False{}
96    case 1n+r:
97      True{}
98
99# Handle index path in the packed bit-tree operations.
100def index_path(height: Nat, +index: Nat, acc: List<&2, Bool>) -> List<&2, Bool>:
101  match height:
102    case 0n:
103      acc
104    case 1n+h:
105      index_path(h, Nat.div(index, 2n), Con{path_bit(Nat.mod(index, 2n)), acc})
106
107# Handle tree word in the packed bit-tree operations.
108def tree_word(path: List<&2, Bool>, tree: WordTree) -> U32:
109  match path tree:
110    case Nil{} Empty{}:
111      0
112    case Nil{} Leaf{word}:
113      word
114    case Nil{} Fork{left, right}:
115      0
116    case Con{side, rest} Empty{}:
117      0
118    case Con{side, rest} Leaf{word}:
119      0
120    case Con{False{}, rest} Fork{left, right}:
121      tree_word(rest, left)
122    case Con{True{}, rest} Fork{left, right}:
123      tree_word(rest, right)
124
125# Handle tree set word in the packed bit-tree operations.
126def tree_set_word(path: List<&2, Bool>, value: U32, tree: WordTree) -> WordTree:
127  match path tree:
128    case Nil{} Empty{}:
129      Empty{}
130    case Nil{} Leaf{word}:
131      Leaf{value}
132    case Nil{} Fork{left, right}:
133      Fork{left, right}
134    case Con{side, rest} Empty{}:
135      Empty{}
136    case Con{side, rest} Leaf{word}:
137      Leaf{word}
138    case Con{False{}, rest} Fork{left, right}:
139      Fork{tree_set_word(rest, value, left), right}
140    case Con{True{}, rest} Fork{left, right}:
141      Fork{left, tree_set_word(rest, value, right)}
142
143# Handle word at if in the packed bit-tree operations.
144def word_at_if(inside: Bool, tree: BitTree, index: Nat) -> U32:
145  match inside:
146    case False{}:
147      0
148    case True{}:
149      match tree:
150        case BT{root, words, capacity, height}:
151          tree_word(index_path(height, index, Nil{}), root)
152
153# Safe logical-word lookup; padding leaves are not externally addressable.
154def word_at(+tree: BitTree, +index: Nat) -> U32:
155  word_at_if(Nat.is_lt(index, word_count(tree)), tree, index)
156
157# Replace word if for the packed bit-tree operations.
158def replace_word_if(inside: Bool, tree: BitTree, index: Nat, value: U32) -> BitTree:
159  match inside:
160    case False{}:
161      tree
162    case True{}:
163      match tree:
164        case BT{root, words, capacity, +height}:
165          BT{tree_set_word(index_path(height, index, Nil{}), value, root), words, capacity, height}
166
167# Safe logical-word update. This helper makes metadata preservation and
168# unaffected-word laws direct to state.
169def replace_word(+tree: BitTree, +index: Nat, value: U32) -> BitTree:
170  replace_word_if(Nat.is_lt(index, word_count(tree)), tree, index, value)
171
172# Update if for the packed bit-tree operations.
173def set_if(inside: Bool, +tree: BitTree, +bit: Nat) -> BitTree:
174  match inside:
175    case False{}:
176      tree
177    case True{}:
178      +index = Nat.div(bit, 32n)
179      +old = word_at(tree, index)
180      replace_word(tree, index, BitWords.word_put(True{}, old, Nat.mod(bit, 32n)))
181
182# Set one logical bit. Out-of-range indexes leave the tree unchanged.
183def set(+tree: BitTree, +bit: Nat) -> BitTree:
184  set_if(valid_bit(tree, bit), tree, bit)
185
186# Handle test if in the packed bit-tree operations.
187def test_if(inside: Bool, +tree: BitTree, +bit: Nat) -> Bool:
188  match inside:
189    case False{}:
190      False{}
191    case True{}:
192      BitWords.word_get(word_at(tree, Nat.div(bit, 32n)), Nat.mod(bit, 32n))
193
194# Test one logical bit. Out-of-range indexes, including every index in the
195# zero-word tree, return False{}.
196def test(+tree: BitTree, +bit: Nat) -> Bool:
197  test_if(valid_bit(tree, bit), tree, bit)
198
199# Word-wise union for chunked Bloom builds. Both sides always share the exact
200# shape (same height from the same bit count); mismatched arms are unreachable
201# by construction and keep the Fork side.
202def union_bits(ta: WordTree, tb: WordTree) -> WordTree:
203  match ta tb:
204    case Empty{} Empty{}:
205      Empty{}
206    case Empty{} Leaf{wb}:
207      Leaf{wb}
208    case Empty{} Fork{lb, rb}:
209      Fork{lb, rb}
210    case Leaf{wa} Empty{}:
211      Leaf{wa}
212    case Leaf{wa} Leaf{wb}:
213      Leaf{U32.or(wa, wb)}
214    case Leaf{wa} Fork{lb, rb}:
215      Fork{lb, rb}
216    case Fork{la, ra} Empty{}:
217      Fork{la, ra}
218    case Fork{la, ra} Leaf{wb}:
219      Fork{la, ra}
220    case Fork{la, ra} Fork{lb, rb}:
221      Fork{union_bits(la, lb), union_bits(ra, rb)}
222
223# Structural checker for the exact perfect-tree shape described by metadata.
224def shape_ok(+height: Nat, tree: WordTree) -> Bool:
225  match height tree:
226    case 0n Empty{}:
227      False{}
228    case 0n Leaf{word}:
229      True{}
230    case 0n Fork{left, right}:
231      False{}
232    case 1n+h Empty{}:
233      False{}
234    case 1n+h Leaf{word}:
235      False{}
236    case 1n+h Fork{left, right}:
237      a b = shape_ok(h, left) shape_ok(h, right)
238      Bool.and(a, b)
239
240# Check whether whether metadata ok holds for the packed bit-tree operations.
241def empty_metadata_ok(words: Nat, capacity: Nat, height: Nat) -> Bool:
242  Bool.and(Nat.is_eq(words, 0n), Bool.and(Nat.is_eq(capacity, 0n), Nat.is_eq(height, 0n)))
243
244# Handle nonempty metadata ok in the packed bit-tree operations.
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)