1	import Base
2	import 0x9ee2e9a299991dcc089fe22c7f3ceb5f/src/containers/bitset.bend as BitWords
3	import 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.
18	type 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.
24	type 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.
28	type 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.
33	def 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.
46	def 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.
54	def 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.
63	def 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.
69	def 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.
78	def 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.
84	def 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.
90	def 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.
96	def bit_count(tree: BitTree) -> Nat:
97	  Nat.mul(word_count(tree), 32n)
98	
99	# Check whether bit holds for the packed bit-tree operations.
100	def 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.
104	def 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.
112	def 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.
120	def 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.
128	def 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.
146	def 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.
164	def 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.
174	def 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.
178	def 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.
189	def 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.
193	def 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.
203	def 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.
207	def 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{}.
216	def 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.
222	def 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.
244	def 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.
261	def 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.
265	def 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.
276	def 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)