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	type WordTree is Data:
18	  Empty{}
19	  Leaf{word: U32}
20	  Fork{left: WordTree, right: WordTree}
21	
22	type BitTree is Data:
23	  BT{root: WordTree, words: Nat, capacity: Nat, height: Nat}
24	
25	type 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.
30	def 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	
42	def 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.
50	def 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	
58	def 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	
63	def 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.
72	def word_count(tree: BitTree) -> Nat:
73	  match tree:
74	    case BT{root, words, capacity, height}:
75	      words
76	
77	def capacity_words(tree: BitTree) -> Nat:
78	  match tree:
79	    case BT{root, words, capacity, height}:
80	      capacity
81	
82	def tree_height(tree: BitTree) -> Nat:
83	  match tree:
84	    case BT{root, words, capacity, height}:
85	      height
86	
87	def bit_count(tree: BitTree) -> Nat:
88	  Nat.mul(word_count(tree), 32n)
89	
90	def valid_bit(tree: BitTree, bit: Nat) -> Bool:
91	  Nat.is_lt(bit, bit_count(tree))
92	
93	def 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.
101	def path_bit(rem: Nat) -> Bool:
102	  match rem:
103	    case 0n:
104	      False{}
105	    case 1n+r:
106	      True{}
107	
108	def 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	
115	def 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	
132	def 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	
149	def 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.
159	def word_at(+tree: BitTree, +index: Nat) -> U32:
160	  word_at_if(Nat.is_lt(index, word_count(tree)), tree, index)
161	
162	def 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.
173	def replace_word(+tree: BitTree, +index: Nat, value: U32) -> BitTree:
174	  replace_word_if(Nat.is_lt(index, word_count(tree)), tree, index, value)
175	
176	def 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.
186	def set(+tree: BitTree, +bit: Nat) -> BitTree:
187	  set_if(valid_bit(tree, bit), tree, bit)
188	
189	def 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{}.
198	def 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.
204	def 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.
226	def 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	
242	def 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	
245	def 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.
256	def 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)