1use sha2::{Digest, Sha256};
2use tracing::debug;
3
4use crate::{
5 NonExistenceProof, Proof,
6 bintrie_error::BinTrieError,
7 bintrie_types::{Hash, Pubkey},
8};
9
10#[derive(Debug, Clone, PartialEq)]
12pub struct BinTriePair {
13 pub pubkey: Pubkey,
14 pub value_hash: Hash,
15 pub is_sibling_hash: bool,
16}
17
18impl BinTriePair {
19 pub fn new(pubkey: Pubkey, value_hash: Hash) -> Self {
20 Self {
21 pubkey,
22 value_hash,
23 is_sibling_hash: false,
24 }
25 }
26
27 pub fn new_sibling_hash(sibling_hash: Hash) -> Self {
28 Self {
29 pubkey: Pubkey::default(),
30 value_hash: sibling_hash,
31 is_sibling_hash: true,
32 }
33 }
34}
35
36#[derive(Debug, Clone, PartialEq)]
38pub struct BinTrieLeaf {
39 pub hash: Hash,
40 pub pair: BinTriePair,
41}
42
43impl BinTrieLeaf {
44 pub fn new(pair: BinTriePair) -> Self {
45 let mut leaf = Self {
46 hash: Hash::default(),
47 pair,
48 };
49 leaf.compute_hash();
50 leaf
51 }
52
53 pub fn compute_hash(&mut self) {
55 if self.pair.is_sibling_hash {
56 self.hash = self.pair.value_hash;
57 } else {
58 let mut hasher = Sha256::new();
59 hasher.update(&[0x00]); hasher.update(self.pair.pubkey.as_bytes());
61 hasher.update(self.pair.value_hash.as_bytes());
62 let result = hasher.finalize();
63 self.hash = Hash::new(result.into());
64 }
65 }
66}
67
68#[derive(Debug, Clone, PartialEq)]
70pub struct BinTrieNode {
71 pub hash: Hash,
72 pub bit_idx: u8,
73 pub left: Option<Box<BinTrieElement>>,
74 pub right: Option<Box<BinTrieElement>>,
75}
76
77impl BinTrieNode {
78 pub fn new(bit_idx: u8) -> Self {
79 Self {
80 hash: Hash::default(),
81 bit_idx,
82 left: None,
83 right: None,
84 }
85 }
86
87 pub fn compute_hash(&mut self) {
89 let left_hash = self
90 .left
91 .as_ref()
92 .map(|e| e.hash())
93 .unwrap_or(Hash::default());
94 let right_hash = self
95 .right
96 .as_ref()
97 .map(|e| e.hash())
98 .unwrap_or(Hash::default());
99
100 let mut hasher = Sha256::new();
101 hasher.update(&[0x01]); hasher.update(left_hash.as_bytes());
103 hasher.update(right_hash.as_bytes());
104 let result = hasher.finalize();
105 self.hash = Hash::new(result.into());
106 }
107}
108
109#[derive(Debug, Clone, PartialEq)]
111pub enum BinTrieElement {
112 Node(BinTrieNode),
113 Leaf(BinTrieLeaf),
114}
115
116impl BinTrieElement {
117 pub fn hash(&self) -> Hash {
118 match self {
119 BinTrieElement::Node(node) => node.hash,
120 BinTrieElement::Leaf(leaf) => leaf.hash,
121 }
122 }
123
124 pub fn compute_hash(&mut self) {
125 match self {
126 BinTrieElement::Node(node) => node.compute_hash(),
127 BinTrieElement::Leaf(leaf) => leaf.compute_hash(),
128 }
129 }
130
131 pub fn is_leaf(&self) -> bool {
132 matches!(self, BinTrieElement::Leaf(_))
133 }
134
135 pub fn is_node(&self) -> bool {
136 matches!(self, BinTrieElement::Node(_))
137 }
138}
139
140#[derive(Debug, Clone, PartialEq)]
147pub struct BinTrie {
148 root: Option<Box<BinTrieElement>>,
149}
150
151impl Default for BinTrie {
152 fn default() -> Self {
153 Self::new()
154 }
155}
156
157impl BinTrie {
158 pub fn new() -> Self {
160 Self { root: None }
161 }
162
163 pub fn is_empty(&self) -> bool {
165 self.root.is_none()
166 }
167
168 pub fn state_root(&self) -> Hash {
170 self.root
171 .as_ref()
172 .map(|r| r.hash())
173 .unwrap_or(Hash::default())
174 }
175
176 pub fn query(&self, pubkey: &Pubkey) -> Option<&BinTriePair> {
178 let mut current = self.root.as_ref()?;
179
180 loop {
182 match current.as_ref() {
183 BinTrieElement::Leaf(leaf) => {
184 if leaf.pair.pubkey == *pubkey {
185 return Some(&leaf.pair);
186 } else {
187 return None;
188 }
189 }
190 BinTrieElement::Node(node) => {
191 let go_right = pubkey.get_bit(node.bit_idx);
192 current = if go_right {
193 node.right.as_ref()?
194 } else {
195 node.left.as_ref()?
196 };
197 }
198 }
199 }
200 }
201
202 pub fn insert(&mut self, pubkey: Pubkey, value_hash: Hash) -> Result<(), BinTrieError> {
204 if self.root.is_none() {
206 let pair = BinTriePair::new(pubkey, value_hash);
207 self.root = Some(Box::new(BinTrieElement::Leaf(BinTrieLeaf::new(pair))));
208 return Ok(());
209 }
210
211 if self.query(&pubkey).is_some() {
213 return Err(BinTrieError::KeyExists);
214 }
215
216 let mut path_nodes: Vec<(u8, bool)> = Vec::new(); let existing_pubkey = {
219 let mut current = self.root.as_ref().unwrap();
220 loop {
221 match current.as_ref() {
222 BinTrieElement::Leaf(leaf) => break leaf.pair.pubkey,
223 BinTrieElement::Node(node) => {
224 let go_right = pubkey.get_bit(node.bit_idx);
225 path_nodes.push((node.bit_idx, go_right));
226 current = if go_right {
227 node.right.as_ref().unwrap()
228 } else {
229 node.left.as_ref().unwrap()
230 };
231 }
232 }
233 }
234 };
235
236 let mut split_idx = 0u8;
238 for bit_idx in 0..=255 {
239 if pubkey.get_bit(bit_idx) != existing_pubkey.get_bit(bit_idx) {
240 split_idx = bit_idx;
241 break;
242 }
243 }
244
245 let mut parent_depth = None;
248 for (i, &(bit_idx, _)) in path_nodes.iter().enumerate().rev() {
249 if bit_idx < split_idx {
250 parent_depth = Some(i);
251 break;
252 }
253 }
254
255 let new_pair = BinTriePair::new(pubkey, value_hash);
257 let new_leaf = BinTrieLeaf::new(new_pair);
258
259 let mut new_node = BinTrieNode::new(split_idx);
261 let new_key_goes_right = pubkey.get_bit(split_idx);
262
263 if let Some(parent_idx) = parent_depth {
264 let (_, parent_go_right) = path_nodes[parent_idx];
266
267 let child_subtree = {
269 let mut current = self.root.as_mut().unwrap();
270 for &(_, go_right) in &path_nodes[0..parent_idx] {
271 current = match current.as_mut() {
272 BinTrieElement::Node(node) => {
273 if go_right {
274 node.right.as_mut().unwrap()
275 } else {
276 node.left.as_mut().unwrap()
277 }
278 }
279 _ => unreachable!(),
280 };
281 }
282
283 match current.as_mut() {
285 BinTrieElement::Node(parent_node) => {
286 if parent_go_right {
287 parent_node.right.take().unwrap()
288 } else {
289 parent_node.left.take().unwrap()
290 }
291 }
292 _ => unreachable!(),
293 }
294 };
295
296 if new_key_goes_right {
298 new_node.left = Some(child_subtree);
299 new_node.right = Some(Box::new(BinTrieElement::Leaf(new_leaf)));
300 } else {
301 new_node.left = Some(Box::new(BinTrieElement::Leaf(new_leaf)));
302 new_node.right = Some(child_subtree);
303 }
304 new_node.compute_hash();
305
306 {
308 let mut current = self.root.as_mut().unwrap();
309 for &(_, go_right) in &path_nodes[0..parent_idx] {
310 current = match current.as_mut() {
311 BinTrieElement::Node(node) => {
312 if go_right {
313 node.right.as_mut().unwrap()
314 } else {
315 node.left.as_mut().unwrap()
316 }
317 }
318 _ => unreachable!(),
319 };
320 }
321
322 match current.as_mut() {
324 BinTrieElement::Node(parent_node) => {
325 if parent_go_right {
326 parent_node.right = Some(Box::new(BinTrieElement::Node(new_node)));
327 } else {
328 parent_node.left = Some(Box::new(BinTrieElement::Node(new_node)));
329 }
330 }
331 _ => unreachable!(),
332 }
333 }
334
335 for depth in (0..=parent_idx).rev() {
337 let mut current = self.root.as_mut().unwrap();
338 for &(_, go_right) in &path_nodes[0..depth] {
339 current = match current.as_mut() {
340 BinTrieElement::Node(node) => {
341 if go_right {
342 node.right.as_mut().unwrap()
343 } else {
344 node.left.as_mut().unwrap()
345 }
346 }
347 _ => unreachable!(),
348 };
349 }
350
351 match current.as_mut() {
352 BinTrieElement::Node(node) => {
353 node.compute_hash();
354 }
355 _ => break,
356 }
357 }
358 } else {
359 let existing_tree = self.root.take().unwrap();
361
362 if new_key_goes_right {
363 new_node.left = Some(existing_tree);
364 new_node.right = Some(Box::new(BinTrieElement::Leaf(new_leaf)));
365 } else {
366 new_node.left = Some(Box::new(BinTrieElement::Leaf(new_leaf)));
367 new_node.right = Some(existing_tree);
368 }
369 new_node.compute_hash();
370 self.root = Some(Box::new(BinTrieElement::Node(new_node)));
371 }
372
373 Ok(())
374 }
375
376 pub fn update_hash(
378 &mut self,
379 pubkey: &Pubkey,
380 new_value_hash: Hash,
381 ) -> Result<(), BinTrieError> {
382 Self::update_hash_recursive(&mut self.root, pubkey, new_value_hash)
383 }
384
385 fn update_hash_recursive(
386 current: &mut Option<Box<BinTrieElement>>,
387 pubkey: &Pubkey,
388 new_value_hash: Hash,
389 ) -> Result<(), BinTrieError> {
390 let current_element = current.as_mut().ok_or(BinTrieError::KeyNotFound)?;
391
392 match current_element.as_mut() {
393 BinTrieElement::Leaf(leaf) => {
394 if leaf.pair.pubkey == *pubkey {
395 leaf.pair.value_hash = new_value_hash;
396 leaf.compute_hash();
397 Ok(())
398 } else {
399 Err(BinTrieError::KeyNotFound)
400 }
401 }
402 BinTrieElement::Node(node) => {
403 let go_right = pubkey.get_bit(node.bit_idx);
404 let child = if go_right {
405 &mut node.right
406 } else {
407 &mut node.left
408 };
409
410 Self::update_hash_recursive(child, pubkey, new_value_hash)?;
411 node.compute_hash();
412 Ok(())
413 }
414 }
415 }
416
417 pub fn prove_existence(&self, pubkey: &Pubkey) -> Result<(Proof, Hash), BinTrieError> {
419 let mut proof = Proof::new();
420 let mut current = self.root.as_ref().ok_or(BinTrieError::KeyNotFound)?;
421
422 loop {
424 match current.as_ref() {
425 BinTrieElement::Leaf(leaf) => {
426 if leaf.pair.pubkey == *pubkey {
427 return Ok((proof, leaf.pair.value_hash));
428 } else {
429 return Err(BinTrieError::KeyNotFound);
430 }
431 }
432 BinTrieElement::Node(node) => {
433 let go_right = pubkey.get_bit(node.bit_idx);
434 proof.proof_indices.push(node.bit_idx);
435
436 if go_right {
437 let sibling_hash = node
439 .left
440 .as_ref()
441 .map(|e| e.hash())
442 .unwrap_or(Hash::default());
443 proof.sibling_hashes.push(sibling_hash);
444 current = node.right.as_ref().ok_or(BinTrieError::KeyNotFound)?;
445 } else {
446 let sibling_hash = node
448 .right
449 .as_ref()
450 .map(|e| e.hash())
451 .unwrap_or(Hash::default());
452 proof.sibling_hashes.push(sibling_hash);
453 current = node.left.as_ref().ok_or(BinTrieError::KeyNotFound)?;
454 }
455 }
456 }
457 }
458 }
459
460 pub fn prove_non_existence(&self, pubkey: &Pubkey) -> Result<NonExistenceProof, BinTrieError> {
462 if self.root.is_none() {
464 return Ok(NonExistenceProof {
465 proof: Proof::new(),
466 existing_pubkey: Pubkey::default(),
467 existing_hash: Hash::default(),
468 });
469 }
470
471 let mut proof = Proof::new();
472 let mut current = self.root.as_ref().unwrap();
473
474 loop {
476 match current.as_ref() {
477 BinTrieElement::Leaf(leaf) => {
478 if leaf.pair.pubkey == *pubkey {
479 return Err(BinTrieError::KeyExists);
480 } else {
481 return Ok(NonExistenceProof {
482 proof,
483 existing_pubkey: leaf.pair.pubkey,
484 existing_hash: leaf.pair.value_hash,
485 });
486 }
487 }
488 BinTrieElement::Node(node) => {
489 let bit_idx = node.bit_idx;
491 let go_right = pubkey.get_bit(bit_idx);
492 proof.proof_indices.push(bit_idx);
493
494 if go_right {
495 let sibling_hash = node
496 .left
497 .as_ref()
498 .map(|e| e.hash())
499 .unwrap_or(Hash::default());
500 proof.sibling_hashes.push(sibling_hash);
501 current = node.right.as_ref().unwrap();
502 } else {
503 let sibling_hash = node
504 .right
505 .as_ref()
506 .map(|e| e.hash())
507 .unwrap_or(Hash::default());
508 proof.sibling_hashes.push(sibling_hash);
509 current = node.left.as_ref().unwrap();
510 }
511 }
512 }
513 }
514 }
515
516 pub fn insert_with_proof(
518 &mut self,
519 pubkey: Pubkey,
520 value_hash: Hash,
521 proof: &Proof,
522 ) -> Result<(), BinTrieError> {
523 if self.root.is_none() && proof.proof_indices.is_empty() {
525 return self.insert(pubkey, value_hash);
526 }
527
528 self.insert_with_proof_recursive(pubkey, value_hash, proof, 0)
530 }
531
532 fn insert_with_proof_recursive(
533 &mut self,
534 pubkey: Pubkey,
535 value_hash: Hash,
536 proof: &Proof,
537 depth: usize,
538 ) -> Result<(), BinTrieError> {
539 if depth >= proof.proof_indices.len() {
540 let pair = BinTriePair::new(pubkey, value_hash);
542 self.root = Some(Box::new(BinTrieElement::Leaf(BinTrieLeaf::new(pair))));
543 return Ok(());
544 }
545
546 let bit_idx = proof.proof_indices[depth];
547 let sibling_hash = proof.sibling_hashes[depth];
548 let go_right = pubkey.get_bit(bit_idx);
549
550 let mut node = BinTrieNode::new(bit_idx);
552
553 let sibling_pair = BinTriePair::new_sibling_hash(sibling_hash);
555 let sibling_leaf = BinTrieLeaf::new(sibling_pair);
556
557 if depth + 1 < proof.proof_indices.len() {
558 let mut subtrie = BinTrie::new();
560 subtrie.insert_with_proof_recursive(pubkey, value_hash, proof, depth + 1)?;
561
562 if go_right {
563 node.left = Some(Box::new(BinTrieElement::Leaf(sibling_leaf)));
564 node.right = subtrie.root;
565 } else {
566 node.left = subtrie.root;
567 node.right = Some(Box::new(BinTrieElement::Leaf(sibling_leaf)));
568 }
569 } else {
570 let pair = BinTriePair::new(pubkey, value_hash);
572 let new_leaf = BinTrieLeaf::new(pair);
573
574 if go_right {
575 node.left = Some(Box::new(BinTrieElement::Leaf(sibling_leaf)));
576 node.right = Some(Box::new(BinTrieElement::Leaf(new_leaf)));
577 } else {
578 node.left = Some(Box::new(BinTrieElement::Leaf(new_leaf)));
579 node.right = Some(Box::new(BinTrieElement::Leaf(sibling_leaf)));
580 }
581 }
582
583 node.compute_hash();
584 self.root = Some(Box::new(BinTrieElement::Node(node)));
585 Ok(())
586 }
587
588 pub fn insert_with_creation_proof(
590 &mut self,
591 pubkey: Pubkey,
592 value_hash: Hash,
593 existing_pubkey: Pubkey,
594 existing_value_hash: Hash,
595 proof: &Proof,
596 ) -> Result<(), BinTrieError> {
597 if self.root.is_none() && proof.proof_indices.is_empty() {
599 self.insert(existing_pubkey, existing_value_hash)?;
601 self.insert(pubkey, value_hash)?;
603 return Ok(());
604 }
605
606 if self.root.is_some() || !proof.proof_indices.is_empty() {
608 match self.insert_with_proof(existing_pubkey, existing_value_hash, proof) {
609 Ok(()) => {}
610 Err(BinTrieError::KeyExists) => {} Err(e) => return Err(e),
612 }
613 }
614
615 self.insert(pubkey, value_hash)
617 }
618
619 pub fn print(&self) {
621 println!("Binary Trie:");
622 if let Some(root) = &self.root {
623 self.print_element(root, 1);
624 } else {
625 println!(" (Empty trie)");
626 }
627 }
628
629 pub fn print_verbose(&self) {
631 println!("Binary Trie (Verbose):");
632 println!(" Root hash: {}", self.state_root());
633 if let Some(root) = &self.root {
634 self.print_element_verbose(root, 1);
635 } else {
636 println!(" (Empty trie)");
637 }
638 }
639
640 fn print_element(&self, element: &BinTrieElement, depth: usize) {
641 let indent = " ".repeat(depth);
642
643 match element {
644 BinTrieElement::Leaf(leaf) => {
645 if leaf.pair.is_sibling_hash {
646 println!("{}S {}", indent, leaf.hash);
647 } else {
648 println!("{}L {}", indent, leaf.pair.pubkey);
649 }
650 }
651 BinTrieElement::Node(node) => {
652 println!("{}N bit={}", indent, node.bit_idx);
653 if let Some(left) = &node.left {
654 self.print_element(left, depth + 1);
655 }
656 if let Some(right) = &node.right {
657 self.print_element(right, depth + 1);
658 }
659 }
660 }
661 }
662
663 fn print_element_verbose(&self, element: &BinTrieElement, depth: usize) {
664 let indent = " ".repeat(depth);
665
666 match element {
667 BinTrieElement::Leaf(leaf) => {
668 if leaf.pair.is_sibling_hash {
669 println!("{}Leaf: SIBLING HASH: {}", indent, leaf.hash);
670 } else {
671 println!(
672 "{}Leaf: KEY: {:?}, VALUE HASH: {}, LEAF HASH: {}",
673 indent, leaf.pair.pubkey, leaf.pair.value_hash, leaf.hash
674 );
675 }
676 }
677 BinTrieElement::Node(node) => {
678 println!(
679 "{}Node: bit_idx={}, HASH: {}",
680 indent, node.bit_idx, node.hash
681 );
682 println!("{}Left:", indent);
683 if let Some(left) = &node.left {
684 self.print_element_verbose(left, depth + 1);
685 } else {
686 println!("{} (null)", indent);
687 }
688 println!("{}Right:", indent);
689 if let Some(right) = &node.right {
690 self.print_element_verbose(right, depth + 1);
691 } else {
692 println!("{} (null)", indent);
693 }
694 }
695 }
696 }
697
698 pub fn print_log_verbose(&self) {
700 debug!("Binary Trie (Verbose):");
701 debug!(" Root hash: {}", self.state_root());
702 if let Some(root) = &self.root {
703 self.print_element_log_verbose(root, 1);
704 } else {
705 debug!(" (Empty trie)");
706 }
707 }
708
709 fn print_element_log_verbose(&self, element: &BinTrieElement, depth: usize) {
710 let indent = " ".repeat(depth);
711
712 match element {
713 BinTrieElement::Leaf(leaf) => {
714 if leaf.pair.is_sibling_hash {
715 debug!("{}Leaf: SIBLING HASH: {}", indent, leaf.hash);
716 } else {
717 debug!(
718 "{}Leaf: KEY: {}, VALUE HASH: {}, LEAF HASH: {}",
719 indent, leaf.pair.pubkey, leaf.pair.value_hash, leaf.hash
720 );
721 }
722 }
723 BinTrieElement::Node(node) => {
724 debug!(
725 "{}Node: bit_idx={}, HASH: {}",
726 indent, node.bit_idx, node.hash
727 );
728 debug!("{}Left:", indent);
729 if let Some(left) = &node.left {
730 self.print_element_log_verbose(left, depth + 1);
731 } else {
732 debug!("{} (null)", indent);
733 }
734 debug!("{}Right:", indent);
735 if let Some(right) = &node.right {
736 self.print_element_log_verbose(right, depth + 1);
737 } else {
738 debug!("{} (null)", indent);
739 }
740 }
741 }
742 }
743}
744
745pub mod test_helpers {
747 use super::*;
748
749 pub fn hash_leaf(pubkey: &Pubkey, value_hash: &Hash) -> Hash {
751 let mut hasher = Sha256::new();
752 hasher.update(&[0x00]);
753 hasher.update(pubkey.as_bytes());
754 hasher.update(value_hash.as_bytes());
755 let result = hasher.finalize();
756 Hash::new(result.into())
757 }
758
759 pub fn hash_node(left_hash: &Hash, right_hash: &Hash) -> Hash {
761 let mut hasher = Sha256::new();
762 hasher.update(&[0x01]);
763 hasher.update(left_hash.as_bytes());
764 hasher.update(right_hash.as_bytes());
765 let result = hasher.finalize();
766 Hash::new(result.into())
767 }
768}