1#[cfg(feature = "py-bindings")]
2use chia_py_streamable_macro::{PyJsonDict, PyStreamable};
3#[cfg(feature = "py-bindings")]
4use pyo3::{
5 Bound, IntoPyObject, PyAny, PyResult, Python,
6 buffer::PyBuffer,
7 pyclass, pymethods,
8 types::{PyDict, PyDictMethods, PyListMethods, PyType},
9};
10
11use crate::merkle::iterators::{BreadthFirstIterator, LeftChildFirstIterator, ParentFirstIterator};
12use crate::merkle::{
13 deltas, format, proof_of_inclusion,
14 util::{sha256_bytes, sha256_num},
15};
16use crate::{
17 BLOCK_SIZE, Block, BlockBytes, Hash, InternalNode, KeyId, LeafNode, Node, NodeMetadata,
18 NodeType, Parent, TreeIndex, ValueId, merkle::error::Error,
19};
20use bitvec::prelude::BitVec;
21use chia_protocol::Bytes32;
22use chia_sha2::Sha256;
23use chia_streamable_macro::Streamable;
24#[cfg(feature = "py-bindings")]
25use chia_traits::Streamable;
26use indexmap::IndexSet;
27use std::cmp::Ordering;
28use std::collections::{HashMap, HashSet};
29use std::io::{Read, Write};
30#[cfg(feature = "py-bindings")]
31use std::iter::zip;
32use std::ops::Range;
33use std::path::PathBuf;
34
35pub fn zstd_decode_path(path: &PathBuf) -> Result<Vec<u8>, Error> {
40 let mut vector: Vec<u8> = Vec::new();
41 let file = std::fs::File::open(path)?;
42 let mut decoder = zstd::Decoder::new(file)?;
43 decoder.read_to_end(&mut vector)?;
44
45 Ok(vector)
46}
47
48pub fn internal_hash(left_hash: &Hash, right_hash: &Hash) -> Hash {
49 let mut hasher = Sha256::new();
50 hasher.update(b"\x02");
51 hasher.update(left_hash.0);
52 hasher.update(right_hash.0);
53
54 Hash(Bytes32::new(hasher.finalize()))
55}
56
57pub fn calculate_internal_hash(hash: &Hash, other_hash_side: Side, other_hash: &Hash) -> Hash {
58 match other_hash_side {
59 Side::Left => internal_hash(other_hash, hash),
60 Side::Right => internal_hash(hash, other_hash),
61 }
62}
63
64#[cfg_attr(feature = "py-bindings", derive(PyJsonDict, PyStreamable))]
65#[repr(u8)]
66#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Streamable)]
67pub enum Side {
68 Left = 0,
69 Right = 1,
70}
71
72#[cfg_attr(feature = "py-bindings", pyclass(from_py_object))]
73#[derive(Clone, Debug, Hash, Eq, PartialEq)]
74pub enum InsertLocation {
75 Auto {},
79 AsRoot {},
80 Leaf { index: TreeIndex, side: Side },
81}
82
83pub fn block_range(index: TreeIndex) -> Range<usize> {
84 let block_start = index.0 as usize * BLOCK_SIZE;
85 block_start..block_start + BLOCK_SIZE
86}
87
88#[cfg_attr(feature = "py-bindings", pyclass(from_py_object))]
89#[derive(Clone, Debug)]
90pub struct BlockStatusCache {
91 free_indexes: IndexSet<TreeIndex>,
92 key_to_index: HashMap<KeyId, TreeIndex>,
93 leaf_hash_to_index: HashMap<Hash, TreeIndex>,
94}
95
96impl BlockStatusCache {
97 fn new(blob: &[u8]) -> Result<Self, Error> {
98 let index_count = blob.len() / BLOCK_SIZE;
99
100 let mut seen_indexes: BitVec<u64, bitvec::order::Lsb0> = BitVec::repeat(false, index_count);
101 let mut key_to_index: HashMap<KeyId, TreeIndex> = HashMap::default();
102 let mut leaf_hash_to_index: HashMap<Hash, TreeIndex> = HashMap::default();
103
104 for item in LeftChildFirstIterator::new(blob, None) {
105 let (index, block) = item?;
106 seen_indexes.set(index.0 as usize, true);
107
108 if let Node::Leaf(leaf) = block.node {
109 if key_to_index.insert(leaf.key, index).is_some() {
110 return Err(Error::KeyAlreadyPresent());
111 }
112 if leaf_hash_to_index.insert(leaf.hash, index).is_some() {
113 return Err(Error::HashAlreadyPresent());
114 }
115 }
116 }
117
118 let mut free_indexes: IndexSet<TreeIndex> = IndexSet::new();
119 for (index, seen) in seen_indexes.iter().enumerate() {
120 if !seen {
121 free_indexes.insert(TreeIndex(index as u32));
122 }
123 }
124
125 Ok(Self {
126 free_indexes,
127 key_to_index,
128 leaf_hash_to_index,
129 })
130 }
131
132 fn iter_keys_indexes(&self) -> impl Iterator<Item = (&KeyId, &TreeIndex)> {
133 self.key_to_index.iter()
134 }
135
136 fn pop_free_index(&mut self) -> Option<TreeIndex> {
137 let maybe_index = self.free_indexes.iter().next().copied();
138 if let Some(index) = maybe_index {
139 self.free_indexes.shift_remove(&index);
140 }
141
142 maybe_index
143 }
144
145 fn get_index_by_key(&self, key: KeyId) -> Option<&TreeIndex> {
146 self.key_to_index.get(&key)
147 }
148
149 fn get_index_by_leaf_hash(&self, hash: &Hash) -> Option<&TreeIndex> {
150 self.leaf_hash_to_index.get(hash)
151 }
152
153 #[must_use]
154 fn is_index_free(&self, index: TreeIndex) -> bool {
155 self.free_indexes.contains(&index)
156 }
157
158 fn leaf_count(&self) -> usize {
159 self.key_to_index.len()
160 }
161
162 fn free_index_count(&self) -> usize {
163 self.free_indexes.len()
164 }
165
166 fn no_keys(&self) -> bool {
167 self.key_to_index.is_empty()
168 }
169
170 fn contains_key(&self, key: KeyId) -> bool {
171 self.key_to_index.contains_key(&key)
172 }
173
174 fn contains_leaf_hash(&self, hash: &Hash) -> bool {
175 self.leaf_hash_to_index.contains_key(hash)
176 }
177
178 fn clear(&mut self) {
179 self.key_to_index.clear();
180 self.free_indexes.clear();
181 self.leaf_hash_to_index.clear();
182 }
183
184 fn add_internal(&mut self, index: TreeIndex) {
185 self.free_indexes.shift_remove(&index);
186 }
187
188 fn add_leaf(&mut self, index: TreeIndex, leaf: LeafNode) {
189 self.free_indexes.shift_remove(&index);
190
191 self.key_to_index.insert(leaf.key, index);
192 self.leaf_hash_to_index.insert(leaf.hash, index);
193 }
194
195 fn remove_internal(&mut self, index: TreeIndex) {
196 self.free_indexes.insert(index);
197 }
198
199 fn remove_leaf(&mut self, node: &LeafNode) -> Result<(), Error> {
200 let Some(index) = self.key_to_index.remove(&node.key) else {
201 return Err(Error::UnknownKey(node.key));
202 };
203 self.leaf_hash_to_index.remove(&node.hash);
204
205 self.free_indexes.insert(index);
206
207 Ok(())
208 }
209
210 fn move_index(&mut self, source: TreeIndex, destination: TreeIndex) -> Result<(), Error> {
211 if self.free_indexes.contains(&source) {
214 return Err(Error::MoveSourceIndexNotInUse(source));
215 }
216 if self.free_indexes.contains(&destination) {
218 return Err(Error::MoveDestinationIndexNotInUse(destination));
219 }
220
221 self.free_indexes.insert(source);
222
223 Ok(())
224 }
225}
226
227pub type NodeHashToIndex = HashMap<Hash, TreeIndex>;
228pub type NodeHashToDeltaReaderNode = HashMap<Hash, deltas::DeltaReaderNode>;
229
230pub fn collect_and_return_from_merkle_blob(
231 path: &PathBuf,
232 hashes: &HashSet<Hash>,
233 known: impl Fn(&Hash) -> bool,
234) -> Result<(NodeHashToDeltaReaderNode, NodeHashToIndex), Error> {
235 let mut nodes = NodeHashToDeltaReaderNode::new();
236 let blob = zstd_decode_path(path)?;
237 let mut node_hash_to_index = NodeHashToIndex::new();
238
239 let mut index_to_hash: HashMap<TreeIndex, Hash> = HashMap::new();
240
241 let mut in_subtree: HashSet<Hash> = HashSet::new();
242 let mut index_stack: Vec<(TreeIndex, bool)> = Vec::new();
243 index_stack.push((TreeIndex(0), false));
244 while let Some((index, visited)) = index_stack.pop() {
245 let block = format::try_get_block(&blob, index)?;
246
247 let node_hash = block.node.hash();
248 index_to_hash.insert(index, node_hash);
249 if known(&node_hash) {
250 continue;
251 }
252
253 match block.node {
254 Node::Internal(InternalNode {
255 hash, left, right, ..
256 }) => {
257 if visited {
258 node_hash_to_index.insert(hash, index);
259 if !in_subtree.is_empty() {
260 nodes.insert(
261 hash,
262 deltas::DeltaReaderNode::Internal {
263 left: *index_to_hash.get(&left).unwrap(),
264 right: *index_to_hash.get(&right).unwrap(),
265 },
266 );
267 }
268
269 in_subtree.remove(&hash);
270 } else {
271 if hashes.contains(&hash) {
272 in_subtree.insert(hash);
273 }
274
275 index_stack.push((index, true));
276 index_stack.push((right, false));
277 index_stack.push((left, false));
278 }
279 }
280 Node::Leaf(LeafNode {
281 hash, key, value, ..
282 }) => {
283 if !in_subtree.is_empty() || hashes.contains(&hash) {
284 nodes.insert(hash, deltas::DeltaReaderNode::Leaf { key, value });
285 }
286
287 node_hash_to_index.insert(hash, index);
288 }
289 }
290 }
291
292 Ok((nodes, node_hash_to_index))
293}
294
295pub type InternalNodesMap = HashMap<Hash, (Hash, Hash)>;
296pub type LeafNodesMap = HashMap<Hash, (KeyId, ValueId)>;
297
298#[cfg_attr(feature = "py-bindings", pyclass(get_all, from_py_object))]
306#[derive(Clone, Debug)]
307pub struct MerkleBlob {
308 pub(crate) blob: Vec<u8>,
309 block_status_cache: BlockStatusCache,
310 pub check_integrity_on_drop: bool,
313}
314
315impl MerkleBlob {
316 pub fn new(blob: Vec<u8>) -> Result<Self, Error> {
317 let length = blob.len();
318 let remainder = length % BLOCK_SIZE;
319 if remainder != 0 {
320 return Err(Error::InvalidBlobLength(remainder));
321 }
322
323 let block_status_cache = BlockStatusCache::new(&blob)?;
324
325 let self_ = Self {
326 blob,
327 block_status_cache,
328 check_integrity_on_drop: cfg!(test),
329 };
330
331 Ok(self_)
332 }
333
334 pub fn from_path(path: &PathBuf) -> Result<Self, Error> {
335 let vector = zstd_decode_path(path)?;
336
337 Self::new(vector)
338 }
339
340 pub fn to_path(&self, path: &PathBuf) -> Result<(), Error> {
341 let directory = path.parent().ok_or(std::io::Error::new(
342 std::io::ErrorKind::IsADirectory,
343 format!(
344 "path must be a file, root directory given: {}",
345 path.display()
346 ),
347 ))?;
348 std::fs::create_dir_all(directory)?;
349 let file = std::fs::File::create(path)?;
350 let mut encoder = zstd::Encoder::new(file, 0)?;
351 encoder.write_all(&self.blob)?;
352 encoder.finish()?;
353
354 Ok(())
355 }
356
357 fn clear(&mut self) {
358 self.blob.clear();
359 self.block_status_cache.clear();
360 }
361
362 pub fn insert(
363 &mut self,
364 key: KeyId,
365 value: ValueId,
366 hash: &Hash,
367 insert_location: InsertLocation,
368 ) -> Result<TreeIndex, Error> {
369 if self.block_status_cache.contains_key(key) {
370 return Err(Error::KeyAlreadyPresent());
371 }
372 if self.block_status_cache.contains_leaf_hash(hash) {
373 return Err(Error::HashAlreadyPresent());
374 }
375
376 let insert_location = match insert_location {
377 InsertLocation::Auto {} => self.get_random_insert_location_by_key_id(key)?,
378 _ => insert_location,
379 };
380
381 match insert_location {
382 InsertLocation::Auto {} => {
383 unreachable!("this should have been caught and processed above")
384 }
385 InsertLocation::AsRoot {} => {
386 if !self.block_status_cache.no_keys() {
387 return Err(Error::UnableToInsertAsRootOfNonEmptyTree());
388 }
389 self.insert_first(key, value, hash)
390 }
391 InsertLocation::Leaf { index, side } => {
392 let old_leaf = self.get_node(index)?.try_into_leaf()?;
393
394 let internal_node_hash = match side {
395 Side::Left => internal_hash(hash, &old_leaf.hash),
396 Side::Right => internal_hash(&old_leaf.hash, hash),
397 };
398
399 let node = LeafNode {
400 parent: Parent(None),
401 hash: *hash,
402 key,
403 value,
404 };
405
406 if self.block_status_cache.leaf_count() == 1 {
407 self.insert_second(node, &old_leaf, &internal_node_hash, side)
408 } else {
409 self.insert_third_or_later(node, &old_leaf, index, &internal_node_hash, side)
410 }
411 }
412 }
413 }
414
415 fn insert_first(
416 &mut self,
417 key: KeyId,
418 value: ValueId,
419 hash: &Hash,
420 ) -> Result<TreeIndex, Error> {
421 let new_leaf_block = Block {
422 metadata: NodeMetadata {
423 node_type: NodeType::Leaf,
424 dirty: false,
425 },
426 node: Node::Leaf(LeafNode {
427 parent: Parent(None),
428 key,
429 value,
430 hash: *hash,
431 }),
432 };
433
434 let index = self.extend_index();
435 self.insert_entry_to_blob(index, &new_leaf_block)?;
436
437 Ok(index)
438 }
439
440 fn insert_second(
441 &mut self,
442 mut node: LeafNode,
443 old_leaf: &LeafNode,
444 internal_node_hash: &Hash,
445 side: Side,
446 ) -> Result<TreeIndex, Error> {
447 self.clear();
448 let root_index = self.get_new_index();
449 let left_index = self.get_new_index();
450 let right_index = self.get_new_index();
451
452 let new_internal_block = Block {
453 metadata: NodeMetadata {
454 node_type: NodeType::Internal,
455 dirty: false,
456 },
457 node: Node::Internal(InternalNode {
458 parent: Parent(None),
459 left: left_index,
460 right: right_index,
461 hash: *internal_node_hash,
462 }),
463 };
464
465 self.insert_entry_to_blob(root_index, &new_internal_block)?;
466
467 node.parent = Parent(Some(TreeIndex(0)));
468
469 let nodes = [
470 (
471 match side {
472 Side::Left => right_index,
473 Side::Right => left_index,
474 },
475 LeafNode {
476 parent: Parent(Some(TreeIndex(0))),
477 key: old_leaf.key,
478 value: old_leaf.value,
479 hash: old_leaf.hash,
480 },
481 ),
482 (
483 match side {
484 Side::Left => left_index,
485 Side::Right => right_index,
486 },
487 node,
488 ),
489 ];
490
491 for (index, node) in nodes {
492 let block = Block {
493 metadata: NodeMetadata {
494 node_type: NodeType::Leaf,
495 dirty: false,
496 },
497 node: Node::Leaf(node),
498 };
499
500 self.insert_entry_to_blob(index, &block)?;
501 }
502
503 Ok(nodes[1].0)
504 }
505
506 fn insert_third_or_later(
507 &mut self,
508 mut node: LeafNode,
509 old_leaf: &LeafNode,
510 old_leaf_index: TreeIndex,
511 internal_node_hash: &Hash,
512 side: Side,
513 ) -> Result<TreeIndex, Error> {
514 let new_leaf_index = self.get_new_index();
515 let new_internal_node_index = self.get_new_index();
516
517 node.parent = Parent(Some(new_internal_node_index));
518
519 let new_leaf_block = Block {
520 metadata: NodeMetadata {
521 node_type: NodeType::Leaf,
522 dirty: false,
523 },
524 node: Node::Leaf(node),
525 };
526 self.insert_entry_to_blob(new_leaf_index, &new_leaf_block)?;
527
528 let (left_index, right_index) = match side {
529 Side::Left => (new_leaf_index, old_leaf_index),
530 Side::Right => (old_leaf_index, new_leaf_index),
531 };
532 let new_internal_block = Block {
533 metadata: NodeMetadata {
534 node_type: NodeType::Internal,
535 dirty: false,
536 },
537 node: Node::Internal(InternalNode {
538 parent: old_leaf.parent,
539 left: left_index,
540 right: right_index,
541 hash: *internal_node_hash,
542 }),
543 };
544 self.insert_entry_to_blob(new_internal_node_index, &new_internal_block)?;
545
546 let old_parent_index = old_leaf.parent.0.expect("root found when not expected");
547
548 self.update_parent(old_leaf_index, Some(new_internal_node_index))?;
549
550 let mut old_parent_block = self.get_block(old_parent_index)?;
551 if let Node::Internal(ref mut internal_node, ..) = old_parent_block.node {
552 if old_leaf_index == internal_node.left {
553 internal_node.left = new_internal_node_index;
554 } else if old_leaf_index == internal_node.right {
555 internal_node.right = new_internal_node_index;
556 } else {
557 panic!("child not a child of its parent");
558 }
559 } else {
560 panic!("expected internal node but found leaf");
561 }
562
563 self.insert_entry_to_blob(old_parent_index, &old_parent_block)?;
564
565 self.mark_lineage_as_dirty(old_parent_index)?;
566
567 Ok(new_leaf_index)
568 }
569
570 pub fn batch_insert(
571 &mut self,
572 mut keys_values_hashes: Vec<((KeyId, ValueId), Hash)>,
573 ) -> Result<(), Error> {
574 let mut indexes = vec![];
577
578 if self.block_status_cache.leaf_count() <= 1 {
579 for _ in 0..2 {
580 let Some(((key, value), hash)) = keys_values_hashes.pop() else {
581 return Ok(());
582 };
583 self.insert(key, value, &hash, InsertLocation::Auto {})?;
584 }
585 }
586
587 for ((key, value), hash) in keys_values_hashes {
588 let new_leaf_index = self.get_new_index();
589 let new_block = Block {
590 metadata: NodeMetadata {
591 node_type: NodeType::Leaf,
592 dirty: false,
593 },
594 node: Node::Leaf(LeafNode {
595 parent: Parent(None),
596 hash,
597 key,
598 value,
599 }),
600 };
601 self.insert_entry_to_blob(new_leaf_index, &new_block)?;
602 indexes.push(new_leaf_index);
603 }
604
605 while indexes.len() > 1 {
608 let mut new_indexes = vec![];
609
610 for chunk in indexes.chunks(2) {
611 let [index_1, index_2] = match chunk {
612 [index] => {
613 new_indexes.push(*index);
614 continue;
615 }
616 [index_1, index_2] => [*index_1, *index_2],
617 _ => unreachable!(
618 "chunk should always be either one or two long and be handled above"
619 ),
620 };
621
622 let new_internal_node_index = self.get_new_index();
623
624 let mut hashes = vec![];
625 for index in [index_1, index_2] {
626 let block = self.update_parent(index, Some(new_internal_node_index))?;
627 hashes.push(block.node.hash());
628 }
629
630 let new_block = Block {
631 metadata: NodeMetadata {
632 node_type: NodeType::Internal,
633 dirty: false,
634 },
635 node: Node::Internal(InternalNode {
636 parent: Parent(None),
637 hash: internal_hash(&hashes[0], &hashes[1]),
638 left: index_1,
639 right: index_2,
640 }),
641 };
642
643 self.insert_entry_to_blob(new_internal_node_index, &new_block)?;
644 new_indexes.push(new_internal_node_index);
645 }
646
647 indexes = new_indexes;
648 }
649
650 if indexes.len() == 1 {
651 let min_height_leaf = self.get_min_height_leaf()?;
653 self.insert_subtree_at_key(min_height_leaf.key, indexes[0], Side::Left)?;
654 }
655
656 Ok(())
657 }
658
659 fn insert_subtree_at_key(
660 &mut self,
661 old_leaf_key: KeyId,
662 new_index: TreeIndex,
663 side: Side,
664 ) -> Result<(), Error> {
665 struct Stuff {
668 index: TreeIndex,
669 hash: Hash,
670 }
671
672 let new_internal_node_index = self.get_new_index();
673 let (old_leaf_index, old_leaf, _old_block) = self.get_leaf_by_key(old_leaf_key)?;
674 let new_node = self.get_node(new_index)?;
675
676 let new_stuff = Stuff {
677 index: new_index,
678 hash: new_node.hash(),
679 };
680 let old_stuff = Stuff {
681 index: old_leaf_index,
682 hash: old_leaf.hash,
683 };
684 let (left, right) = match side {
685 Side::Left => (new_stuff, old_stuff),
686 Side::Right => (old_stuff, new_stuff),
687 };
688 let internal_node_hash = internal_hash(&left.hash, &right.hash);
689
690 let block = Block {
691 metadata: NodeMetadata {
692 node_type: NodeType::Internal,
693 dirty: false,
694 },
695 node: Node::Internal(InternalNode {
696 parent: old_leaf.parent,
697 hash: internal_node_hash,
698 left: left.index,
699 right: right.index,
700 }),
701 };
702 self.insert_entry_to_blob(new_internal_node_index, &block)?;
703 self.update_parent(new_index, Some(new_internal_node_index))?;
704
705 let Some(old_leaf_parent) = old_leaf.parent.0 else {
706 return Err(Error::LeafCannotBeRootWhenInsertingSubtree());
707 };
708
709 let mut parent = self.get_block(old_leaf_parent)?;
710 if let Node::Internal(ref mut internal) = parent.node {
711 match old_leaf_index {
712 x if x == internal.left => internal.left = new_internal_node_index,
713 x if x == internal.right => internal.right = new_internal_node_index,
714 _ => panic!("parent not a child a grandparent"),
715 }
716 } else {
717 panic!("not handling this case now...")
718 }
719 self.insert_entry_to_blob(old_leaf_parent, &parent)?;
720 self.mark_lineage_as_dirty(old_leaf_parent)?;
721 self.update_parent(old_leaf_index, Some(new_internal_node_index))?;
722
723 Ok(())
724 }
725
726 fn get_min_height_leaf(&self) -> Result<LeafNode, Error> {
727 let (_index, block) = BreadthFirstIterator::new(&self.blob, None)
728 .next()
729 .ok_or(Error::UnableToFindALeaf())??;
730
731 Ok(block
732 .node
733 .expect_leaf("unexpectedly found internal node first: <<self>>"))
734 }
735
736 pub fn delete(&mut self, key: KeyId) -> Result<(), Error> {
737 let (leaf_index, leaf, _leaf_block) = self.get_leaf_by_key(key)?;
738 self.block_status_cache.remove_leaf(&leaf)?;
739
740 let Some(parent_index) = leaf.parent.0 else {
741 self.clear();
742 return Ok(());
743 };
744
745 let maybe_parent = self.get_node(parent_index)?;
746 let Node::Internal(parent) = maybe_parent else {
747 panic!("parent node not internal: {maybe_parent:?}")
748 };
749 let sibling_index = parent.sibling_index(leaf_index)?;
750 let mut sibling_block = self.get_block(sibling_index)?;
751
752 let Some(grandparent_index) = parent.parent.0 else {
753 sibling_block.node.set_parent(Parent(None));
754 let destination = TreeIndex(0);
755 if let Node::Internal(node) = sibling_block.node {
756 for child_index in [node.left, node.right] {
757 self.update_parent(child_index, Some(destination))?;
758 }
759 }
760
761 self.insert_entry_to_blob(destination, &sibling_block)?;
762 self.block_status_cache
763 .move_index(sibling_index, destination)?;
764
765 return Ok(());
766 };
767
768 self.block_status_cache.remove_internal(parent_index);
769 let mut grandparent_block = self.get_block(grandparent_index)?;
770
771 sibling_block
772 .node
773 .set_parent(Parent(Some(grandparent_index)));
774 self.insert_entry_to_blob(sibling_index, &sibling_block)?;
775
776 if let Node::Internal(ref mut internal) = grandparent_block.node {
777 match parent_index {
778 x if x == internal.left => internal.left = sibling_index,
779 x if x == internal.right => internal.right = sibling_index,
780 _ => panic!("parent not a child a grandparent"),
781 }
782 } else {
783 panic!("grandparent not an internal node")
784 }
785 self.insert_entry_to_blob(grandparent_index, &grandparent_block)?;
786
787 self.mark_lineage_as_dirty(grandparent_index)?;
788
789 Ok(())
790 }
791
792 pub fn upsert(&mut self, key: KeyId, value: ValueId, new_hash: &Hash) -> Result<(), Error> {
793 let Ok((leaf_index, mut leaf, mut block)) = self.get_leaf_by_key(key) else {
794 self.insert(key, value, new_hash, InsertLocation::Auto {})?;
795 return Ok(());
796 };
797
798 self.block_status_cache.remove_leaf(&leaf)?;
799 leaf.hash.clone_from(new_hash);
800 leaf.value = value;
801 block.node = Node::Leaf(leaf);
803 self.insert_entry_to_blob(leaf_index, &block)?;
804
805 if let Some(parent) = block.node.parent().0 {
806 self.mark_lineage_as_dirty(parent)?;
807 }
808
809 Ok(())
810 }
811
812 pub fn check_integrity(&self) -> Result<(), Error> {
813 self.check_just_integrity()?;
814
815 let mut clone = self.clone();
816 clone.check_integrity_on_drop = false;
817 clone.calculate_lazy_hashes()?;
818 clone.check_just_integrity()
819 }
820
821 fn check_just_integrity(&self) -> Result<(), Error> {
822 let mut leaf_count: usize = 0;
823 let mut internal_count: usize = 0;
824 let mut child_to_parent: HashMap<TreeIndex, TreeIndex> = HashMap::new();
825
826 for item in ParentFirstIterator::new(&self.blob, None) {
827 let (index, block) = item?;
828 if let Some(parent) = block.node.parent().0 {
829 if child_to_parent.remove(&index) != Some(parent) {
830 return Err(Error::IntegrityParentChildMismatch(index));
831 }
832 }
833 match block.node {
834 Node::Internal(node) => {
835 internal_count += 1;
836 child_to_parent.insert(node.left, index);
837 child_to_parent.insert(node.right, index);
838 }
839 Node::Leaf(node) => {
840 leaf_count += 1;
841 let cached_index = self
842 .block_status_cache
843 .get_index_by_key(node.key)
844 .ok_or(Error::IntegrityKeyNotInCache(node.key))?;
845 if *cached_index != index {
846 return Err(Error::IntegrityKeyToIndexCacheIndex(
847 node.key,
848 index,
849 *cached_index,
850 ));
851 }
852 assert!(
853 !self.block_status_cache.is_index_free(index),
854 "{}",
855 format!("active index found in free index list: {index:?}")
856 );
857 }
858 }
859 }
860
861 let key_to_index_cache_length = self.block_status_cache.key_to_index.len();
862 if leaf_count != key_to_index_cache_length {
863 return Err(Error::IntegrityKeyToIndexCacheLength(
864 leaf_count,
865 key_to_index_cache_length,
866 ));
867 }
868 let leaf_hash_to_index_cache_length = self.block_status_cache.leaf_hash_to_index.len();
869 if leaf_count != leaf_hash_to_index_cache_length {
870 return Err(Error::IntegrityLeafHashToIndexCacheLength(
871 leaf_count,
872 leaf_hash_to_index_cache_length,
873 ));
874 }
875 let total_count = leaf_count + internal_count + self.block_status_cache.free_index_count();
876 let extend_index = self.extend_index();
877 if total_count != extend_index.0 as usize {
878 return Err(Error::IntegrityTotalNodeCount(extend_index, total_count));
879 }
880 if !child_to_parent.is_empty() {
881 return Err(Error::IntegrityUnmatchedChildParentRelationships(
882 child_to_parent.len(),
883 ));
884 }
885
886 Ok(())
887 }
888
889 fn update_parent(
890 &mut self,
891 index: TreeIndex,
892 parent: Option<TreeIndex>,
893 ) -> Result<Block, Error> {
894 let mut block = self.get_block(index)?;
895 block.node.set_parent(Parent(parent));
896 self.insert_entry_to_blob(index, &block)?;
897
898 Ok(block)
899 }
900
901 fn mark_lineage_as_dirty(&mut self, index: TreeIndex) -> Result<(), Error> {
902 let mut next_index = Some(index);
903
904 while let Some(this_index) = next_index {
905 let mut block = Block::from_bytes(self.get_block_bytes(this_index)?)?;
906
907 if block.metadata.dirty {
908 break;
909 }
910
911 block.metadata.dirty = true;
912 self.insert_entry_to_blob(this_index, &block)?;
913 next_index = block.node.parent().0;
914 }
915
916 Ok(())
917 }
918
919 fn get_new_index(&mut self) -> TreeIndex {
920 match self.block_status_cache.pop_free_index() {
921 None => {
922 let index = self.extend_index();
923 self.blob.extend_from_slice(&[0; BLOCK_SIZE]);
924 index
928 }
929 Some(new_index) => new_index,
930 }
931 }
932
933 fn get_random_insert_location_by_seed(
934 &self,
935 seed_bytes: &[u8],
936 ) -> Result<InsertLocation, Error> {
937 let mut seed_bytes = Vec::from(seed_bytes);
938
939 if self.blob.is_empty() {
940 return Ok(InsertLocation::AsRoot {});
941 }
942
943 let final_side = if (seed_bytes
945 .first()
946 .ok_or(Error::ZeroLengthSeedNotAllowed())?
947 & (1 << 7))
948 == 0
949 {
950 Side::Left
951 } else {
952 Side::Right
953 };
954
955 let mut next_index = TreeIndex(0);
956 let mut node = self.get_node(next_index)?;
957
958 seed_bytes.reverse();
959 loop {
960 for byte in &seed_bytes {
961 for bit_index in 0..8 {
962 match node {
963 Node::Leaf { .. } => {
964 return Ok(InsertLocation::Leaf {
965 index: next_index,
966 side: final_side,
967 });
968 }
969 Node::Internal(internal) => {
970 let bit = byte & (1 << bit_index) != 0;
971 next_index = if bit { internal.right } else { internal.left };
972 node = self.get_node(next_index)?;
973 }
974 }
975 }
976 }
977
978 seed_bytes = sha256_bytes(&seed_bytes).0.into();
979 }
980 }
981
982 pub fn get_hash_at_index(&self, index: TreeIndex) -> Result<Option<Hash>, Error> {
983 if self.block_status_cache.no_keys() {
984 return Ok(None);
985 }
986
987 let block = self.get_block(index)?;
988 if block.metadata.dirty {
989 return Err(Error::Dirty(index));
990 }
991
992 Ok(Some(block.node.hash()))
993 }
994
995 fn get_random_insert_location_by_key_id(&self, seed: KeyId) -> Result<InsertLocation, Error> {
996 let seed = sha256_num(&seed.0);
997
998 self.get_random_insert_location_by_seed(&seed.0)
999 }
1000
1001 fn extend_index(&self) -> TreeIndex {
1002 let blob_length = self.blob.len();
1003 let index: TreeIndex = TreeIndex((blob_length / BLOCK_SIZE) as u32);
1004 let remainder = blob_length % BLOCK_SIZE;
1005 assert_eq!(
1006 remainder, 0,
1007 "blob length {blob_length:?} not a multiple of {BLOCK_SIZE:?}, remainder: {remainder:?}"
1008 );
1009
1010 index
1011 }
1012
1013 fn insert_entry_to_blob(&mut self, index: TreeIndex, block: &Block) -> Result<(), Error> {
1014 let new_block_bytes = block.to_bytes()?;
1015 let extend_index = self.extend_index();
1016 match index.cmp(&extend_index) {
1017 Ordering::Greater => return Err(Error::BlockIndexOutOfBounds(index)),
1018 Ordering::Equal => self.blob.extend_from_slice(&new_block_bytes),
1019 Ordering::Less => {
1020 self.blob[block_range(index)].copy_from_slice(&new_block_bytes);
1021 }
1022 }
1023
1024 match block.node {
1025 Node::Leaf(leaf) => self.block_status_cache.add_leaf(index, leaf),
1026 Node::Internal(..) => self.block_status_cache.add_internal(index),
1027 }
1028
1029 Ok(())
1030 }
1031
1032 fn get_block(&self, index: TreeIndex) -> Result<Block, Error> {
1033 Block::from_bytes(self.get_block_bytes(index)?)
1034 }
1035
1036 pub(crate) fn get_hash(&self, index: TreeIndex) -> Result<Hash, Error> {
1037 Ok(self.get_block(index)?.node.hash())
1038 }
1039
1040 fn get_block_bytes(&self, index: TreeIndex) -> Result<BlockBytes, Error> {
1041 Ok(self
1042 .blob
1043 .get(block_range(index))
1044 .ok_or(Error::BlockIndexOutOfBounds(index))?
1045 .try_into()
1046 .unwrap_or_else(|e| panic!("failed getting block {index}: {e}")))
1047 }
1048
1049 pub fn get_node(&self, index: TreeIndex) -> Result<Node, Error> {
1050 Ok(self.get_block(index)?.node)
1051 }
1052
1053 pub fn get_leaf_by_key(&self, key: KeyId) -> Result<(TreeIndex, LeafNode, Block), Error> {
1054 let index = *self
1055 .block_status_cache
1056 .get_index_by_key(key)
1057 .ok_or(Error::UnknownKey(key))?;
1058 let block = self.get_block(index)?;
1059 let leaf = block.node.expect_leaf(&format!(
1060 "expected leaf for index from key cache: {index} -> <<self>>"
1061 ));
1062
1063 Ok((index, leaf, block))
1064 }
1065
1066 pub fn get_parent_index(&self, index: TreeIndex) -> Result<Parent, Error> {
1067 Ok(self.get_block(index)?.node.parent())
1068 }
1069
1070 pub fn get_lineage_blocks_with_indexes(
1071 &self,
1072 index: TreeIndex,
1073 ) -> Result<Vec<(TreeIndex, Block)>, Error> {
1074 let mut next_index = Some(index);
1075 let mut lineage = vec![];
1076
1077 while let Some(this_index) = next_index {
1078 let block = self.get_block(this_index)?;
1079 next_index = block.node.parent().0;
1080 lineage.push((this_index, block));
1081 }
1082
1083 Ok(lineage)
1084 }
1085
1086 pub fn get_lineage_with_indexes(
1087 &self,
1088 index: TreeIndex,
1089 ) -> Result<Vec<(TreeIndex, Node)>, Error> {
1090 Ok(self
1091 .get_lineage_blocks_with_indexes(index)?
1092 .iter()
1093 .map(|(index, block)| (*index, block.node))
1094 .collect())
1095 }
1096
1097 pub fn get_lineage_indexes(&self, index: TreeIndex) -> Result<Vec<TreeIndex>, Error> {
1098 Ok(self
1099 .get_lineage_blocks_with_indexes(index)?
1100 .iter()
1101 .map(|(index, _block)| *index)
1102 .collect())
1103 }
1104
1105 pub fn calculate_lazy_hashes(&mut self) -> Result<(), Error> {
1110 for item in LeftChildFirstIterator::new_with_block_predicate(
1112 &self.blob,
1113 None,
1114 Some(|block: &Block| block.metadata.dirty),
1115 )
1116 .collect::<Vec<_>>()
1117 {
1118 let (index, mut block) = item?;
1119 assert!(block.metadata.dirty);
1120
1121 let Node::Internal(ref leaf) = block.node else {
1122 panic!("leaves should not be dirty")
1123 };
1124 let left_hash = self.get_hash(leaf.left)?;
1127 let right_hash = self.get_hash(leaf.right)?;
1128 block.update_hash(&left_hash, &right_hash);
1129 self.insert_entry_to_blob(index, &block)?;
1130 }
1131
1132 Ok(())
1133 }
1134
1135 pub fn get_keys_values(&self) -> Result<HashMap<KeyId, ValueId>, Error> {
1136 let mut map = HashMap::new();
1137 for (key, index) in self.block_status_cache.iter_keys_indexes() {
1138 let node = self.get_node(*index)?;
1139 let leaf = node.expect_leaf(
1140 "key was just retrieved from the key to index mapping, must be a leaf",
1141 );
1142 map.insert(*key, leaf.value);
1143 }
1144
1145 Ok(map)
1146 }
1147
1148 pub fn get_key_index(&self, key: KeyId) -> Result<TreeIndex, Error> {
1149 self.block_status_cache
1150 .get_index_by_key(key)
1151 .copied()
1152 .ok_or(Error::UnknownKey(key))
1153 }
1154
1155 pub fn get_proof_of_inclusion(
1156 &self,
1157 key: KeyId,
1158 ) -> Result<proof_of_inclusion::ProofOfInclusion, Error> {
1159 let mut index = *self
1160 .block_status_cache
1161 .get_index_by_key(key)
1162 .ok_or(Error::UnknownKey(key))?;
1163
1164 let node = self
1165 .get_node(index)?
1166 .expect_leaf("key to index mapping should only have leaves");
1167
1168 let parents = self.get_lineage_blocks_with_indexes(index)?;
1169 let mut layers: Vec<proof_of_inclusion::ProofOfInclusionLayer> = Vec::new();
1170 let mut parents_iter = parents.iter();
1171 parents_iter.next();
1173 for (next_index, block) in parents_iter {
1174 if block.metadata.dirty {
1175 return Err(Error::Dirty(*next_index));
1176 }
1177 let parent = block
1178 .node
1179 .expect_internal("all nodes after the first should be internal");
1180 let sibling_index = parent.sibling_index(index)?;
1181 let sibling_block = self.get_block(sibling_index)?;
1182 let sibling = sibling_block.node;
1183 let layer = proof_of_inclusion::ProofOfInclusionLayer {
1184 other_hash_side: parent.get_sibling_side(index)?,
1185 other_hash: sibling.hash(),
1186 combined_hash: parent.hash,
1187 };
1188 layers.push(layer);
1189 index = *next_index;
1190 }
1191
1192 Ok(proof_of_inclusion::ProofOfInclusion {
1193 node_hash: node.hash,
1194 layers,
1195 })
1196 }
1197
1198 pub fn get_node_by_hash(&self, node_hash: Hash) -> Result<(KeyId, ValueId), Error> {
1199 let Some(index) = self.block_status_cache.get_index_by_leaf_hash(&node_hash) else {
1200 return Err(Error::LeafHashNotFound(node_hash));
1201 };
1202
1203 let node = self
1204 .get_node(*index)?
1205 .expect_leaf("should only have leaves in the leaf hash to index cache");
1206
1207 Ok((node.key, node.value))
1208 }
1209
1210 pub fn get_hashes(&self) -> Result<HashSet<Hash>, Error> {
1211 let mut hashes = HashSet::<Hash>::new();
1212
1213 if self.blob.is_empty() {
1214 return Ok(hashes);
1215 }
1216
1217 for item in ParentFirstIterator::new(&self.blob, None) {
1218 let (_, block) = item?;
1219 hashes.insert(block.node.hash());
1220 }
1221
1222 Ok(hashes)
1223 }
1224
1225 pub fn get_hashes_indexes(&self, leafs_only: bool) -> Result<HashMap<Hash, TreeIndex>, Error> {
1226 let mut hash_to_index = HashMap::new();
1227
1228 if self.blob.is_empty() {
1229 return Ok(hash_to_index);
1230 }
1231
1232 for item in ParentFirstIterator::new(&self.blob, None) {
1233 let (index, block) = item?;
1234
1235 if leafs_only && block.metadata.node_type != NodeType::Leaf {
1236 continue;
1237 }
1238
1239 hash_to_index.insert(block.node.hash(), index);
1240 }
1241
1242 Ok(hash_to_index)
1243 }
1244
1245 pub fn build_blob_from_node_list(
1246 nodes: &NodeHashToDeltaReaderNode,
1247 node_hash: Hash,
1248 interested_hashes: &HashSet<Hash>,
1249 all_used_hashes: &mut HashSet<Hash>,
1250 ) -> Result<Self, Error> {
1251 let mut hashes_and_indexes: Vec<(Hash, TreeIndex)> = Vec::new();
1252 let mut visited: HashSet<Hash> = HashSet::new();
1253 let mut merkle_blob = Self::new(Vec::new())?;
1254 merkle_blob.inner_build_blob_from_node_list(
1255 nodes,
1256 node_hash,
1257 interested_hashes,
1258 &mut hashes_and_indexes,
1259 all_used_hashes,
1260 &mut visited,
1261 0,
1262 )?;
1263
1264 Ok(merkle_blob)
1265 }
1266
1267 #[allow(clippy::too_many_arguments)]
1268 fn inner_build_blob_from_node_list(
1269 &mut self,
1270 nodes: &NodeHashToDeltaReaderNode,
1271 node_hash: Hash,
1272 interested_hashes: &HashSet<Hash>,
1273 hashes_and_indexes: &mut Vec<(Hash, TreeIndex)>,
1274 all_used_hashes: &mut HashSet<Hash>,
1275 visited: &mut HashSet<Hash>,
1276 depth: usize,
1277 ) -> Result<TreeIndex, Error> {
1278 const MAX_RECURSION_DEPTH: usize = 64;
1279 if depth > MAX_RECURSION_DEPTH {
1280 return Err(Error::RecursionDepthExceeded());
1281 }
1282
1283 let node = nodes
1284 .get(&node_hash)
1285 .ok_or(Error::NodeHashNotInNodeMaps(node_hash))?;
1286
1287 if !visited.insert(node_hash) {
1288 return Err(Error::CycleFound());
1289 }
1290
1291 match node {
1292 deltas::DeltaReaderNode::Leaf { key, value } => {
1293 let index = self.get_new_index();
1294 self.insert_entry_to_blob(
1295 index,
1296 &Block {
1297 metadata: NodeMetadata {
1298 node_type: NodeType::Leaf,
1299 dirty: false,
1300 },
1301 node: Node::Leaf(LeafNode {
1302 hash: node_hash,
1303 parent: Parent(None),
1304 key: *key,
1305 value: *value,
1306 }),
1307 },
1308 )?;
1309
1310 if interested_hashes.contains(&node_hash) {
1311 hashes_and_indexes.push((node_hash, index));
1312 }
1313 all_used_hashes.insert(node_hash);
1314
1315 Ok(index)
1316 }
1317 deltas::DeltaReaderNode::Internal { left, right } => {
1318 let index = self.get_new_index();
1319
1320 let left_index = self.inner_build_blob_from_node_list(
1321 nodes,
1322 *left,
1323 interested_hashes,
1324 hashes_and_indexes,
1325 all_used_hashes,
1326 visited,
1327 depth + 1,
1328 )?;
1329 let right_index = self.inner_build_blob_from_node_list(
1330 nodes,
1331 *right,
1332 interested_hashes,
1333 hashes_and_indexes,
1334 all_used_hashes,
1335 visited,
1336 depth + 1,
1337 )?;
1338
1339 for child_index in [left_index, right_index] {
1340 self.update_parent(child_index, Some(index))?;
1341 }
1342 let block = Block {
1343 metadata: NodeMetadata {
1344 node_type: NodeType::Internal,
1345 dirty: false,
1346 },
1347 node: Node::Internal(InternalNode {
1348 hash: node_hash,
1349 parent: Parent(None),
1350 left: left_index,
1351 right: right_index,
1352 }),
1353 };
1354 self.insert_entry_to_blob(index, &block)?;
1355
1356 if interested_hashes.contains(&node_hash) {
1357 hashes_and_indexes.push((node_hash, index));
1358 }
1359 all_used_hashes.insert(node_hash);
1360
1361 Ok(index)
1362 }
1363 }
1364 }
1365
1366 pub fn read_blob(&self) -> &Vec<u8> {
1367 &self.blob
1368 }
1369}
1370
1371#[cfg(feature = "py-bindings")]
1372#[pymethods]
1373impl MerkleBlob {
1374 #[allow(clippy::needless_pass_by_value)]
1375 #[new]
1376 pub fn py_init(blob: PyBuffer<u8>) -> PyResult<Self> {
1377 assert!(
1378 blob.is_c_contiguous(),
1379 "from_bytes() must be called with a contiguous buffer"
1380 );
1381 #[allow(unsafe_code)]
1382 let slice =
1383 unsafe { std::slice::from_raw_parts(blob.buf_ptr() as *const u8, blob.len_bytes()) };
1384
1385 Ok(Self::new(Vec::from(slice))?)
1386 }
1387
1388 #[allow(clippy::needless_pass_by_value)]
1389 #[classmethod]
1390 #[pyo3(name = "from_path")]
1391 pub fn py_from_path(_cls: &Bound<'_, PyType>, path: PathBuf) -> PyResult<Self> {
1392 Ok(Self::from_path(&path)?)
1393 }
1394
1395 #[allow(clippy::needless_pass_by_value)]
1396 #[pyo3(name = "to_path")]
1397 pub fn py_to_path(&self, path: PathBuf) -> PyResult<()> {
1398 Ok(self.to_path(&path)?)
1399 }
1400
1401 #[allow(unused_variables)]
1403 #[must_use]
1404 #[pyo3(name = "__deepcopy__")]
1405 pub fn py_deepcopy(&self, memo: &pyo3::Bound<'_, pyo3::PyAny>) -> Self {
1406 self.clone()
1407 }
1408
1409 #[pyo3(name = "insert", signature = (key, value, hash, reference_kid = None, side = None))]
1410 pub fn py_insert(
1411 &mut self,
1412 key: KeyId,
1413 value: ValueId,
1414 hash: Hash,
1415 reference_kid: Option<KeyId>,
1416 side: Option<u8>,
1418 ) -> PyResult<()> {
1419 let insert_location = match (reference_kid, side) {
1420 (None, None) => InsertLocation::Auto {},
1421 (Some(key), Some(side)) => InsertLocation::Leaf {
1422 index: *self
1423 .block_status_cache
1424 .get_index_by_key(key)
1425 .ok_or(Error::UnknownKey(key))?,
1426 side: Side::from_bytes(&[side])?,
1427 },
1428 _ => Err(Error::IncompleteInsertLocationParameters())?,
1429 };
1430 self.insert(key, value, &hash, insert_location)?;
1431
1432 Ok(())
1433 }
1434
1435 #[pyo3(name = "upsert")]
1436 pub fn py_upsert(&mut self, key: KeyId, value: ValueId, new_hash: Hash) -> PyResult<()> {
1437 self.upsert(key, value, &new_hash)?;
1438
1439 Ok(())
1440 }
1441
1442 #[pyo3(name = "delete")]
1443 pub fn py_delete(&mut self, key: KeyId) -> PyResult<()> {
1444 Ok(self.delete(key)?)
1445 }
1446
1447 #[pyo3(name = "get_raw_node")]
1448 pub fn py_get_raw_node(&mut self, index: TreeIndex) -> PyResult<Node> {
1449 Ok(self.get_node(index)?)
1450 }
1451
1452 #[pyo3(name = "calculate_lazy_hashes")]
1453 pub fn py_calculate_lazy_hashes(&mut self) -> PyResult<()> {
1454 Ok(self.calculate_lazy_hashes()?)
1455 }
1456
1457 #[pyo3(name = "get_lineage_with_indexes")]
1458 pub fn py_get_lineage_with_indexes<'py>(
1459 &self,
1460 index: TreeIndex,
1461 py: Python<'py>,
1462 ) -> PyResult<Bound<'py, PyAny>> {
1463 let list = pyo3::types::PyList::empty(py);
1464
1465 for (index, node) in self.get_lineage_with_indexes(index)? {
1466 list.append((index.into_pyobject(py)?, node.into_pyobject(py)?))?;
1467 }
1468
1469 Ok(list.into_any())
1470 }
1471
1472 #[pyo3(name = "get_nodes_with_indexes", signature = (index=None))]
1473 pub fn py_get_nodes_with_indexes<'py>(
1474 &self,
1475 index: Option<TreeIndex>,
1476 py: Python<'py>,
1477 ) -> PyResult<Bound<'py, PyAny>> {
1478 let list = pyo3::types::PyList::empty(py);
1479
1480 for item in ParentFirstIterator::new(&self.blob, index) {
1481 let (index, block) = item?;
1482 list.append((index.into_pyobject(py)?, block.node.into_pyobject(py)?))?;
1483 }
1484
1485 Ok(list.into_any())
1486 }
1487
1488 #[pyo3(name = "empty")]
1489 pub fn py_empty(&self) -> PyResult<bool> {
1490 Ok(self.block_status_cache.no_keys())
1491 }
1492
1493 #[pyo3(name = "get_root_hash")]
1494 pub fn py_get_root_hash(&self) -> PyResult<Option<Hash>> {
1495 self.py_get_hash_at_index(TreeIndex(0))
1496 }
1497
1498 #[pyo3(name = "get_hash_at_index")]
1499 pub fn py_get_hash_at_index(&self, index: TreeIndex) -> PyResult<Option<Hash>> {
1500 Ok(self.get_hash_at_index(index)?)
1501 }
1502
1503 #[pyo3(name = "batch_insert")]
1504 pub fn py_batch_insert(
1505 &mut self,
1506 keys_values: Vec<(KeyId, ValueId)>,
1507 hashes: Vec<Hash>,
1508 ) -> PyResult<()> {
1509 if keys_values.len() != hashes.len() {
1510 Err(Error::UnmatchedKeysAndValues(
1511 keys_values.len(),
1512 hashes.len(),
1513 ))?;
1514 }
1515
1516 self.batch_insert(zip(keys_values, hashes).collect())?;
1517
1518 Ok(())
1519 }
1520
1521 #[pyo3(name = "__len__")]
1522 pub fn py_len(&self) -> PyResult<usize> {
1523 Ok(self.blob.len())
1524 }
1525
1526 #[pyo3(name = "get_keys_values")]
1527 pub fn py_get_keys_values<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
1528 let map = self.get_keys_values()?;
1529 let dict = PyDict::new(py);
1530 for (key, value) in map {
1531 dict.set_item(key, value)?;
1532 }
1533
1534 Ok(dict.into_any())
1535 }
1536
1537 #[pyo3(name = "get_key_index")]
1538 pub fn py_get_key_index(&self, key: KeyId) -> PyResult<TreeIndex> {
1539 Ok(self.get_key_index(key)?)
1540 }
1541
1542 #[pyo3(name = "get_proof_of_inclusion")]
1543 pub fn py_get_proof_of_inclusion(
1544 &self,
1545 key: KeyId,
1546 ) -> PyResult<proof_of_inclusion::ProofOfInclusion> {
1547 Ok(self.get_proof_of_inclusion(key)?)
1548 }
1549
1550 #[pyo3(name = "get_node_by_hash")]
1551 pub fn py_get_node_by_hash(&self, node_hash: Hash) -> PyResult<(KeyId, ValueId)> {
1552 Ok(self.get_node_by_hash(node_hash)?)
1553 }
1554
1555 #[pyo3(name = "get_hashes_indexes", signature = (leafs_only=false))]
1556 pub fn py_get_hashes_indexes(&self, leafs_only: bool) -> PyResult<HashMap<Hash, TreeIndex>> {
1557 Ok(self.get_hashes_indexes(leafs_only)?)
1558 }
1559
1560 #[pyo3(name = "get_random_leaf_node")]
1561 pub fn py_get_random_leaf_node(&self, seed: &[u8]) -> PyResult<LeafNode> {
1562 let insert_location = self.get_random_insert_location_by_seed(seed)?;
1563 let InsertLocation::Leaf { index, side: _ } = insert_location else {
1564 Err(Error::UnableToFindALeaf())?
1565 };
1566
1567 Ok(self.get_node(index)?.expect_leaf("matched leaf above"))
1568 }
1569
1570 #[pyo3(name = "check_integrity")]
1571 pub fn py_check_integrity(&mut self) -> PyResult<()> {
1572 Ok(self.check_integrity()?)
1573 }
1574}
1575
1576pub fn get_internal_terminal(
1577 blob: &[u8],
1578 indexes: &Vec<TreeIndex>,
1579) -> Result<HashMap<Hash, (TreeIndex, deltas::DeltaReaderNode)>, Error> {
1580 let mut nodes: HashMap<Hash, (TreeIndex, deltas::DeltaReaderNode)> = HashMap::new();
1581 let mut index_to_hash: HashMap<TreeIndex, Hash> = HashMap::new();
1582
1583 for subroot_index in indexes {
1584 for item in LeftChildFirstIterator::new(blob, Some(*subroot_index)) {
1585 let (index, block) = item?;
1586 match block.node {
1587 Node::Internal(node) => {
1588 index_to_hash.insert(index, node.hash);
1589 nodes.insert(
1590 node.hash,
1591 (
1592 index,
1593 deltas::DeltaReaderNode::Internal {
1594 left: *index_to_hash.get(&node.left).unwrap(),
1595 right: *index_to_hash.get(&node.right).unwrap(),
1596 },
1597 ),
1598 );
1599 }
1600 Node::Leaf(node) => {
1601 index_to_hash.insert(index, node.hash);
1602 nodes.insert(
1603 node.hash,
1604 (
1605 index,
1606 deltas::DeltaReaderNode::Leaf {
1607 key: node.key,
1608 value: node.value,
1609 },
1610 ),
1611 );
1612 }
1613 }
1614 }
1615 }
1616
1617 Ok(nodes)
1618}
1619
1620#[cfg(any(test, debug_assertions))]
1621impl Drop for MerkleBlob {
1622 fn drop(&mut self) {
1623 if self.check_integrity_on_drop {
1624 self.check_integrity()
1625 .expect("integrity check failed while dropping merkle blob");
1626 }
1627 }
1628}
1629
1630#[cfg(test)]
1631mod tests {
1632 use super::*;
1633 use crate::merkle::test_util::{
1634 HASH_ONE, HASH_ZERO, generate_hash, open_dot, small_blob, traversal_blob,
1635 };
1636 use crate::merkle::util::sha256_num;
1637 use chia_traits::Streamable;
1638 use expect_test::expect;
1639 use rstest::rstest;
1640 use std::iter::zip;
1641 use std::time::{Duration, Instant};
1642
1643 fn blob_tree_equality(this: &MerkleBlob, that: &MerkleBlob) -> bool {
1644 for item in zip(
1646 LeftChildFirstIterator::new(&this.blob, None),
1647 LeftChildFirstIterator::new(&that.blob, None),
1648 ) {
1649 let (Ok((_, self_block)), Ok((_, other_block))) = item else {
1650 return false;
1651 };
1652 if (self_block.metadata.dirty || other_block.metadata.dirty)
1653 || self_block.node.hash() != other_block.node.hash()
1654 {
1655 return false;
1656 }
1657 match self_block.node {
1658 Node::Internal(..) => {}
1660 Node::Leaf(..) => return self_block.node == other_block.node,
1661 }
1662 }
1663
1664 true
1665 }
1666
1667 #[test]
1668 fn test_node_type_serialized_values() {
1669 assert_eq!(NodeType::Internal as u8, 0);
1670 assert_eq!(NodeType::Leaf as u8, 1);
1671
1672 for node_type in [NodeType::Internal, NodeType::Leaf] {
1673 assert_eq!(
1674 Streamable::to_bytes(&node_type).unwrap()[0],
1675 node_type as u8,
1676 );
1677 assert_eq!(
1678 format::streamable_from_bytes_ignore_extra_bytes::<NodeType>(&[node_type as u8])
1679 .unwrap(),
1680 node_type,
1681 );
1682 }
1683 }
1684
1685 #[test]
1686 fn test_inner_build_blob_from_node_list_depth_limit() {
1687 let mut nodes = NodeHashToDeltaReaderNode::new();
1688 let internal_hashes: Vec<Hash> = (0..=65).map(generate_hash).collect();
1689
1690 for d in 0..=64usize {
1691 let right_leaf_hash = generate_hash(10_000 + d as i32);
1692 nodes.insert(
1693 right_leaf_hash,
1694 deltas::DeltaReaderNode::Leaf {
1695 key: KeyId(d as i64),
1696 value: ValueId(d as i64),
1697 },
1698 );
1699 nodes.insert(
1700 internal_hashes[d],
1701 deltas::DeltaReaderNode::Internal {
1702 left: internal_hashes[d + 1],
1703 right: right_leaf_hash,
1704 },
1705 );
1706 }
1707
1708 let mut blob = MerkleBlob::new(Vec::new()).unwrap();
1709 blob.check_integrity_on_drop = false;
1710 let mut hashes_and_indexes: Vec<(Hash, TreeIndex)> = Vec::new();
1711 let mut all_used_hashes: HashSet<Hash> = HashSet::new();
1712 let mut visited: HashSet<Hash> = HashSet::new();
1713
1714 let err = blob
1715 .inner_build_blob_from_node_list(
1716 &nodes,
1717 internal_hashes[0],
1718 &HashSet::new(),
1719 &mut hashes_and_indexes,
1720 &mut all_used_hashes,
1721 &mut visited,
1722 0,
1723 )
1724 .unwrap_err();
1725
1726 assert!(matches!(err, Error::RecursionDepthExceeded()));
1727 }
1728
1729 #[test]
1730 fn test_inner_build_blob_from_node_list_revisit_node_hash() {
1731 let a = generate_hash(42);
1732 let b = generate_hash(43);
1733 let leaf1 = generate_hash(44);
1734 let leaf2 = generate_hash(45);
1735
1736 let mut nodes = NodeHashToDeltaReaderNode::new();
1737 nodes.insert(
1738 leaf1,
1739 deltas::DeltaReaderNode::Leaf {
1740 key: KeyId(1),
1741 value: ValueId(1),
1742 },
1743 );
1744 nodes.insert(
1745 leaf2,
1746 deltas::DeltaReaderNode::Leaf {
1747 key: KeyId(2),
1748 value: ValueId(2),
1749 },
1750 );
1751 nodes.insert(
1752 a,
1753 deltas::DeltaReaderNode::Internal {
1754 left: b,
1755 right: leaf1,
1756 },
1757 );
1758 nodes.insert(
1759 b,
1760 deltas::DeltaReaderNode::Internal {
1761 left: a,
1762 right: leaf2,
1763 },
1764 );
1765
1766 let mut blob = MerkleBlob::new(Vec::new()).unwrap();
1767 blob.check_integrity_on_drop = false;
1768 let mut hashes_and_indexes: Vec<(Hash, TreeIndex)> = Vec::new();
1769 let mut all_used_hashes: HashSet<Hash> = HashSet::new();
1770 let mut visited: HashSet<Hash> = HashSet::new();
1771
1772 let err = blob
1773 .inner_build_blob_from_node_list(
1774 &nodes,
1775 a,
1776 &HashSet::new(),
1777 &mut hashes_and_indexes,
1778 &mut all_used_hashes,
1779 &mut visited,
1780 0,
1781 )
1782 .unwrap_err();
1783
1784 assert!(matches!(err, Error::CycleFound()));
1785 }
1786
1787 #[test]
1788 fn test_internal_hash() {
1789 let left = Hash((0u8..32).collect::<Vec<_>>().try_into().unwrap());
1792 let right = Hash((32u8..64).collect::<Vec<_>>().try_into().unwrap());
1793
1794 assert_eq!(
1795 internal_hash(&left, &right),
1796 Hash(Bytes32::new(
1797 clvm_utils::tree_hash_pair(
1798 clvm_utils::TreeHash::new(left.0.to_bytes()),
1799 clvm_utils::TreeHash::new(right.0.to_bytes()),
1800 )
1801 .to_bytes()
1802 )),
1803 );
1804 }
1805
1806 #[rstest]
1807 fn test_node_metadata_from_to(
1808 #[values(false, true)] dirty: bool,
1809 #[values(NodeType::Internal, NodeType::Leaf)] node_type: NodeType,
1810 ) {
1811 let bytes: [u8; 2] = [Streamable::to_bytes(&node_type).unwrap()[0], dirty as u8];
1812 let object = NodeMetadata::from_bytes(&bytes).unwrap();
1813 assert_eq!(object, NodeMetadata { node_type, dirty },);
1814 assert_eq!(object.to_bytes().unwrap(), bytes);
1815 }
1816
1817 #[rstest]
1818 fn test_get_lineage(small_blob: MerkleBlob) {
1819 let lineage = small_blob.get_lineage_with_indexes(TreeIndex(2)).unwrap();
1820 for (_, node) in &lineage {
1821 println!("{node:?}");
1822 }
1823 assert_eq!(lineage.len(), 2);
1824 let (_, last_node) = lineage.last().unwrap();
1825 assert_eq!(last_node.parent(), Parent(None));
1826 }
1827
1828 #[rstest]
1829 #[case::right(0, TreeIndex(1), Side::Left)]
1830 #[case::left(0xff, TreeIndex(2), Side::Right)]
1831 fn test_get_random_insert_location_by_seed(
1832 #[case] seed: u8,
1833 #[case] expected_index: TreeIndex,
1834 #[case] expected_side: Side,
1835 small_blob: MerkleBlob,
1836 ) {
1837 let location = small_blob
1838 .get_random_insert_location_by_seed(&[seed; 32])
1839 .unwrap();
1840
1841 assert_eq!(
1842 location,
1843 InsertLocation::Leaf {
1844 index: expected_index,
1845 side: expected_side
1846 },
1847 );
1848 }
1849
1850 #[test]
1851 fn test_get_random_insert_location_by_seed_with_seed_too_short() {
1852 let mut blob = MerkleBlob::new(vec![]).unwrap();
1853 let seed = [0xff];
1854 let layer_count = 8 * seed.len() + 10;
1855
1856 for n in 0..layer_count {
1857 let n = (n + 100) as i64;
1858 let key = KeyId(n);
1859 let value = ValueId(n);
1860 let hash = sha256_num(&key.0);
1861 let insert_location = blob.get_random_insert_location_by_seed(&seed).unwrap();
1862 blob.insert(key, value, &hash, insert_location).unwrap();
1863 }
1864
1865 let location = blob.get_random_insert_location_by_seed(&seed).unwrap();
1866
1867 let InsertLocation::Leaf { index, .. } = location else {
1868 panic!()
1869 };
1870 let lineage = blob.get_lineage_indexes(index).unwrap();
1871
1872 assert_eq!(lineage.len(), layer_count);
1873 assert!(lineage.len() > seed.len() * 8);
1874 }
1875
1876 #[rstest]
1877 fn test_just_insert_a_bunch(
1878 #[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)] n: i64,
1880 ) {
1881 let mut merkle_blob = MerkleBlob::new(vec![]).unwrap();
1882
1883 let mut total_time = Duration::new(0, 0);
1884
1885 let count = 10_000;
1886 let m = count * n;
1887 for i in m..(m + count) {
1888 let start = Instant::now();
1889 merkle_blob
1890 .insert(
1892 KeyId(i),
1893 ValueId(i),
1894 &sha256_num(&i),
1895 InsertLocation::Auto {},
1896 )
1897 .unwrap();
1898 let end = Instant::now();
1899 total_time += end.duration_since(start);
1900 }
1901
1902 println!("total time: {total_time:?}");
1903
1904 merkle_blob.calculate_lazy_hashes().unwrap();
1905 }
1906
1907 #[test]
1908 fn test_delete_in_reverse_creates_matching_trees() {
1909 const COUNT: usize = 10;
1910 let mut dots = vec![];
1911
1912 let mut merkle_blob = MerkleBlob::new(vec![]).unwrap();
1913 let mut reference_blobs = vec![];
1914
1915 let key_value_ids: [i64; COUNT] = core::array::from_fn(|i| i as i64);
1916
1917 for key_value_id in key_value_ids {
1918 let hash: Hash = sha256_num(&key_value_id);
1919
1920 println!("inserting: {key_value_id}");
1921 merkle_blob.calculate_lazy_hashes().unwrap();
1922 reference_blobs.push(MerkleBlob::new(merkle_blob.blob.clone()).unwrap());
1923 merkle_blob
1924 .insert(
1925 KeyId(key_value_id),
1926 ValueId(key_value_id),
1927 &hash,
1928 InsertLocation::Auto {},
1929 )
1930 .unwrap();
1931 dots.push(merkle_blob.to_dot().unwrap().dump());
1932 }
1933
1934 merkle_blob.check_integrity().unwrap();
1935
1936 for key_value_id in key_value_ids.iter().rev() {
1937 println!("deleting: {key_value_id}");
1938 merkle_blob.delete(KeyId(*key_value_id)).unwrap();
1939 merkle_blob.calculate_lazy_hashes().unwrap();
1940 assert!(blob_tree_equality(
1941 &merkle_blob,
1942 &reference_blobs[*key_value_id as usize]
1943 ));
1944 dots.push(merkle_blob.to_dot().unwrap().dump());
1945 }
1946 }
1947
1948 #[test]
1949 fn test_insert_first() {
1950 let mut merkle_blob = MerkleBlob::new(vec![]).unwrap();
1951
1952 let key_value_id = 1;
1953 open_dot(merkle_blob.to_dot().unwrap().set_note("empty"));
1954 merkle_blob
1955 .insert(
1956 KeyId(key_value_id),
1957 ValueId(key_value_id),
1958 &sha256_num(&key_value_id),
1959 InsertLocation::Auto {},
1960 )
1961 .unwrap();
1962 open_dot(merkle_blob.to_dot().unwrap().set_note("first after"));
1963
1964 assert_eq!(merkle_blob.block_status_cache.leaf_count(), 1);
1965 }
1966
1967 #[rstest]
1968 fn test_insert_choosing_side(
1969 #[values(Side::Left, Side::Right)] side: Side,
1970 #[values(1, 2)] pre_count: usize,
1971 ) {
1972 let mut merkle_blob = MerkleBlob::new(vec![]).unwrap();
1973
1974 let mut last_key: KeyId = KeyId(0);
1975 for i in 1..=pre_count {
1976 let key_value = i as i64;
1977 open_dot(merkle_blob.to_dot().unwrap().set_note("empty"));
1978 merkle_blob
1979 .insert(
1980 KeyId(key_value),
1981 ValueId(key_value),
1982 &sha256_num(&key_value),
1983 InsertLocation::Auto {},
1984 )
1985 .unwrap();
1986 last_key = KeyId(key_value);
1987 }
1988
1989 let key_value_id = (pre_count + 1) as i64;
1990 open_dot(merkle_blob.to_dot().unwrap().set_note("first after"));
1991 merkle_blob
1992 .insert(
1993 KeyId(key_value_id),
1994 ValueId(key_value_id),
1995 &sha256_num(&key_value_id),
1996 InsertLocation::Leaf {
1997 index: *merkle_blob
1998 .block_status_cache
1999 .get_index_by_key(last_key)
2000 .unwrap(),
2001 side,
2002 },
2003 )
2004 .unwrap();
2005 open_dot(merkle_blob.to_dot().unwrap().set_note("first after"));
2006
2007 let sibling = merkle_blob
2008 .get_node(
2009 *merkle_blob
2010 .block_status_cache
2011 .get_index_by_key(last_key)
2012 .unwrap(),
2013 )
2014 .unwrap();
2015 let parent = merkle_blob.get_node(sibling.parent().0.unwrap()).unwrap();
2016 let Node::Internal(internal) = parent else {
2017 panic!()
2018 };
2019
2020 let left = merkle_blob
2021 .get_node(internal.left)
2022 .unwrap()
2023 .expect_leaf("<<self>>");
2024 let right = merkle_blob
2025 .get_node(internal.right)
2026 .unwrap()
2027 .expect_leaf("<<self>>");
2028
2029 let expected_keys: [KeyId; 2] = match side {
2030 Side::Left => [KeyId(pre_count as i64 + 1), KeyId(pre_count as i64)],
2031 Side::Right => [KeyId(pre_count as i64), KeyId(pre_count as i64 + 1)],
2032 };
2033 assert_eq!([left.key, right.key], expected_keys);
2034 }
2035
2036 #[test]
2037 fn test_delete_last() {
2038 let mut merkle_blob = MerkleBlob::new(vec![]).unwrap();
2039
2040 let key_value_id = 1;
2041 open_dot(merkle_blob.to_dot().unwrap().set_note("empty"));
2042 merkle_blob
2043 .insert(
2044 KeyId(key_value_id),
2045 ValueId(key_value_id),
2046 &sha256_num(&key_value_id),
2047 InsertLocation::Auto {},
2048 )
2049 .unwrap();
2050 open_dot(merkle_blob.to_dot().unwrap().set_note("first after"));
2051 merkle_blob.check_integrity().unwrap();
2052
2053 merkle_blob.delete(KeyId(key_value_id)).unwrap();
2054
2055 assert_eq!(merkle_blob.block_status_cache.leaf_count(), 0);
2056 }
2057
2058 #[rstest]
2059 fn test_delete_frees_index(mut small_blob: MerkleBlob) {
2060 let key = KeyId(0x0001_0203_0405_0607);
2061 let index = *small_blob.block_status_cache.get_index_by_key(key).unwrap();
2062 small_blob.delete(key).unwrap();
2063
2064 assert_eq!(
2065 small_blob.block_status_cache.free_indexes,
2066 IndexSet::from([index, TreeIndex(1)])
2067 );
2068 }
2069
2070 #[rstest]
2071 fn test_delete_with_internal_sibling(mut small_blob: MerkleBlob) {
2072 let key_to_delete = KeyId(0x0001_0203_0405_0607);
2073 let (other_key_index, _, _) = small_blob.get_leaf_by_key(key_to_delete).unwrap();
2074
2075 small_blob
2076 .insert(
2077 KeyId(0x4041_4243_4445_4647),
2078 ValueId(0x5051_5253_5455_5657),
2079 &sha256_num(&0x4050),
2080 InsertLocation::Leaf {
2081 index: other_key_index,
2082 side: Side::Left,
2083 },
2084 )
2085 .unwrap();
2086
2087 small_blob.delete(key_to_delete).unwrap();
2088
2089 let keys_values = small_blob.get_keys_values().unwrap();
2090 #[allow(clippy::needless_raw_string_hashes)]
2091 let expected = expect![[r#"
2092 [
2093 (
2094 KeyId(
2095 2315169217770759719,
2096 ),
2097 ValueId(
2098 3472611983179986487,
2099 ),
2100 ),
2101 (
2102 KeyId(
2103 4630054748589213255,
2104 ),
2105 ValueId(
2106 5787497513998440023,
2107 ),
2108 ),
2109 ]
2110 "#]];
2111 let mut keys_values = keys_values.iter().collect::<Vec<_>>();
2112 keys_values.sort();
2113 expected.assert_debug_eq(&keys_values);
2114 }
2115
2116 #[rstest]
2117 fn test_get_new_index_with_free_index(mut small_blob: MerkleBlob) {
2118 open_dot(small_blob.to_dot().unwrap().set_note("initial"));
2119 let key = KeyId(0x0001_0203_0405_0607);
2120 let _ = small_blob.block_status_cache.get_index_by_key(key).unwrap();
2121 small_blob.delete(key).unwrap();
2122 open_dot(small_blob.to_dot().unwrap().set_note("after delete"));
2123
2124 let expected = IndexSet::from([TreeIndex(1), TreeIndex(2)]);
2125 assert_eq!(small_blob.block_status_cache.free_indexes, expected);
2126 }
2127
2128 #[rstest]
2129 fn test_dump_small_blob_bytes(small_blob: MerkleBlob) {
2130 println!("{}", hex::encode(small_blob.blob.clone()));
2131 }
2132
2133 #[test]
2134 fn test_node_type_from_u8_invalid() {
2135 let invalid_value = 2;
2136 let actual =
2137 format::streamable_from_bytes_ignore_extra_bytes::<NodeType>(&[invalid_value as u8]);
2138 actual.expect_err("invalid node type value should fail");
2139 }
2140
2141 #[test]
2142 fn test_node_specific_sibling_index_panics_for_unknown_sibling() {
2143 let node = InternalNode {
2144 parent: Parent(None),
2145 hash: sha256_num(&0),
2146 left: TreeIndex(0),
2147 right: TreeIndex(1),
2148 };
2149 let index = TreeIndex(2);
2150 node.sibling_index(TreeIndex(2))
2151 .expect_err(&Error::IndexIsNotAChild(index).to_string());
2152 }
2153
2154 #[rstest]
2155 fn test_get_free_indexes(small_blob: MerkleBlob) {
2156 let mut blob = small_blob.blob.clone();
2157 let expected_free_index = TreeIndex((blob.len() / BLOCK_SIZE) as u32);
2158 blob.extend_from_slice(&[0; BLOCK_SIZE]);
2159 let block_status_cache = BlockStatusCache::new(&blob).unwrap();
2160 assert_eq!(
2161 block_status_cache.free_indexes,
2162 IndexSet::from([expected_free_index])
2163 );
2164 }
2165
2166 #[test]
2167 fn test_merkle_blob_new_errs_for_nonmultiple_of_block_length() {
2168 MerkleBlob::new(vec![1]).expect_err("invalid length should fail");
2169 }
2170
2171 #[rstest]
2172 fn test_upsert_inserts(small_blob: MerkleBlob) {
2173 let key = KeyId(1234);
2174 assert!(!small_blob.block_status_cache.contains_key(key));
2175 let value = ValueId(5678);
2176
2177 let mut insert_blob = MerkleBlob::new(small_blob.blob.clone()).unwrap();
2178 insert_blob
2179 .insert(key, value, &sha256_num(&key.0), InsertLocation::Auto {})
2180 .unwrap();
2181 open_dot(insert_blob.to_dot().unwrap().set_note("first after"));
2182
2183 let mut upsert_blob = MerkleBlob::new(small_blob.blob.clone()).unwrap();
2184 upsert_blob.upsert(key, value, &sha256_num(&key.0)).unwrap();
2185 open_dot(upsert_blob.to_dot().unwrap().set_note("first after"));
2186
2187 assert_eq!(insert_blob.blob, upsert_blob.blob);
2188 }
2189
2190 #[rstest]
2191 fn test_upsert_upserts(mut small_blob: MerkleBlob) {
2192 let before_blocks = LeftChildFirstIterator::new(&small_blob.blob, None).collect::<Vec<_>>();
2193 let (key, index) = small_blob
2194 .block_status_cache
2195 .iter_keys_indexes()
2196 .next()
2197 .unwrap();
2198 let original = small_blob.get_node(*index).unwrap().expect_leaf("<<self>>");
2199 let new_value = ValueId(original.value.0 + 1);
2200
2201 small_blob.upsert(*key, new_value, &original.hash).unwrap();
2202
2203 let after_blocks = LeftChildFirstIterator::new(&small_blob.blob, None).collect::<Vec<_>>();
2204
2205 assert_eq!(before_blocks.len(), after_blocks.len());
2206 for item in zip(before_blocks, after_blocks) {
2207 let ((before_index, before_block), (after_index, after_block)) =
2208 (item.0.unwrap(), item.1.unwrap());
2209 assert_eq!(before_block.node.parent(), after_block.node.parent());
2210 assert_eq!(before_index, after_index);
2211 let before: LeafNode = match before_block.node {
2212 Node::Leaf(leaf) => leaf,
2213 Node::Internal(internal) => {
2214 let Node::Internal(after) = after_block.node else {
2215 panic!()
2216 };
2217 assert_eq!(internal.left, after.left);
2218 assert_eq!(internal.right, after.right);
2219 continue;
2220 }
2221 };
2222 let Node::Leaf(after) = after_block.node else {
2223 panic!()
2224 };
2225 assert_eq!(before.key, after.key);
2226 if before.key == original.key {
2227 assert_eq!(after.value, new_value);
2228 } else {
2229 assert_eq!(before.value, after.value);
2230 }
2231 }
2232 }
2233
2234 #[test]
2235 fn test_double_insert_fails() {
2236 let mut blob = MerkleBlob::new(vec![]).unwrap();
2237 let kv = 0;
2238 blob.insert(
2239 KeyId(kv),
2240 ValueId(kv),
2241 &Hash(Bytes32::new([0u8; 32])),
2242 InsertLocation::Auto {},
2243 )
2244 .unwrap();
2245 blob.insert(
2246 KeyId(kv),
2247 ValueId(kv),
2248 &Hash(Bytes32::new([0u8; 32])),
2249 InsertLocation::Auto {},
2250 )
2251 .expect_err("");
2252 }
2253
2254 #[rstest]
2255 fn test_batch_insert(
2256 #[values(0, 1, 2, 10)] pre_inserts: usize,
2257 #[values(0, 1, 2, 8, 9)] count: usize,
2258 ) {
2259 let mut blob = MerkleBlob::new(vec![]).unwrap();
2260 for i in 0..pre_inserts {
2261 let i = i as i64;
2262 blob.insert(
2263 KeyId(i),
2264 ValueId(i),
2265 &sha256_num(&i),
2266 InsertLocation::Auto {},
2267 )
2268 .unwrap();
2269 }
2270 open_dot(blob.to_dot().unwrap().set_note("initial"));
2271
2272 let mut batch: Vec<((KeyId, ValueId), Hash)> = vec![];
2273
2274 let mut batch_map: HashMap<KeyId, ValueId> = HashMap::new();
2275 for i in pre_inserts..(pre_inserts + count) {
2276 let i = i as i64;
2277 batch.push(((KeyId(i), ValueId(i)), sha256_num(&i)));
2278 batch_map.insert(KeyId(i), ValueId(i));
2279 }
2280
2281 let before = blob.get_keys_values().unwrap();
2282 blob.batch_insert(batch).unwrap();
2283 let after = blob.get_keys_values().unwrap();
2284
2285 open_dot(
2286 blob.to_dot()
2287 .unwrap()
2288 .set_note(&format!("after batch insert of {count} values")),
2289 );
2290
2291 let mut expected = before.clone();
2292 expected.extend(batch_map);
2293
2294 assert_eq!(after, expected);
2295 }
2296
2297 #[rstest]
2298 fn test_root_insert_location_when_not_empty(mut small_blob: MerkleBlob) {
2299 small_blob
2300 .insert(
2301 KeyId(0),
2302 ValueId(0),
2303 &sha256_num(&0),
2304 InsertLocation::AsRoot {},
2305 )
2306 .expect_err("tree not empty so inserting to root should fail");
2307 }
2308
2309 #[rstest]
2310 fn test_free_index_reused(mut small_blob: MerkleBlob) {
2311 let count = 5;
2313 for n in 0..count {
2314 small_blob
2315 .insert(
2316 KeyId(n),
2317 ValueId(n),
2318 &sha256_num(&n),
2319 InsertLocation::Auto {},
2320 )
2321 .unwrap();
2322 }
2323 let (key, index) = {
2324 let (key, index) = small_blob
2325 .block_status_cache
2326 .iter_keys_indexes()
2327 .next()
2328 .unwrap();
2329 (*key, *index)
2330 };
2331 let expected_length = small_blob.blob.len();
2332 assert!(!small_blob.block_status_cache.is_index_free(index));
2333 small_blob.delete(key).unwrap();
2334 assert!(small_blob.block_status_cache.is_index_free(index));
2335 let free_indexes = small_blob.block_status_cache.free_indexes.clone();
2336 assert_eq!(free_indexes.len(), 2);
2337 let new_index = small_blob
2338 .insert(
2339 KeyId(count),
2340 ValueId(count),
2341 &sha256_num(&count),
2342 InsertLocation::Auto {},
2343 )
2344 .unwrap();
2345 assert_eq!(small_blob.blob.len(), expected_length);
2346 assert!(free_indexes.contains(&new_index));
2347 assert_eq!(small_blob.block_status_cache.free_index_count(), 0);
2348 }
2349
2350 #[rstest]
2351 fn test_writing_to_free_block_that_contained_an_active_key(small_blob: MerkleBlob) {
2352 let key = KeyId(0x0001_0203_0405_0607);
2353 let Some(index) = small_blob.block_status_cache.get_index_by_key(key).copied() else {
2354 panic!("maybe the test key needs to be updated?")
2355 };
2356 let mut prepared_bytes = small_blob.blob.clone();
2357 prepared_bytes.extend_from_slice(&small_blob.get_block_bytes(index).unwrap());
2358 let mut prepared_blob = MerkleBlob::new(prepared_bytes).unwrap();
2359 prepared_blob.check_integrity().unwrap();
2360 prepared_blob
2361 .insert(
2362 KeyId(1),
2363 ValueId(2),
2364 &generate_hash(3),
2365 InsertLocation::Auto {},
2366 )
2367 .unwrap();
2368 assert!(prepared_blob.block_status_cache.contains_key(key));
2369 }
2370
2371 #[test]
2372 fn test_node_expect_leaf_passes() {
2373 Node::Leaf(LeafNode {
2374 hash: Hash(Bytes32::default()),
2375 parent: Parent(None),
2376 key: KeyId(0),
2377 value: ValueId(0),
2378 })
2379 .expect_leaf("panic message");
2380 }
2381
2382 #[test]
2383 #[should_panic(expected = "panic message")]
2384 fn test_node_expect_leaf_panics() {
2385 Node::Internal(InternalNode {
2386 hash: Hash(Bytes32::default()),
2387 parent: Parent(None),
2388 left: TreeIndex(0),
2389 right: TreeIndex(0),
2390 })
2391 .expect_leaf("panic message");
2392 }
2393
2394 #[test]
2395 fn test_node_try_into_leaf_passes() {
2396 Node::Leaf(LeafNode {
2397 hash: Hash(Bytes32::default()),
2398 parent: Parent(None),
2399 key: KeyId(0),
2400 value: ValueId(0),
2401 })
2402 .try_into_leaf()
2403 .expect("should pass since it is a leaf");
2404 }
2405
2406 #[test]
2407 fn test_node_try_into_leaf_fails() {
2408 Node::Internal(InternalNode {
2409 hash: Hash(Bytes32::default()),
2410 parent: Parent(None),
2411 left: TreeIndex(0),
2412 right: TreeIndex(0),
2413 })
2414 .try_into_leaf()
2415 .expect_err("should fail since it is not a leaf");
2416 }
2417
2418 #[test]
2419 fn test_node_expect_internal_passes() {
2420 Node::Internal(InternalNode {
2421 hash: Hash(Bytes32::default()),
2422 parent: Parent(None),
2423 left: TreeIndex(0),
2424 right: TreeIndex(0),
2425 })
2426 .expect_internal("panic message");
2427 }
2428
2429 #[test]
2430 #[should_panic(expected = "panic message")]
2431 fn test_node_expect_internal_panics() {
2432 Node::Leaf(LeafNode {
2433 hash: Hash(Bytes32::default()),
2434 parent: Parent(None),
2435 key: KeyId(0),
2436 value: ValueId(0),
2437 })
2438 .expect_internal("panic message");
2439 }
2440
2441 #[test]
2442 fn test_internal_node_get_sibling_side_fails_for_non_sibling() {
2443 let node = InternalNode {
2444 hash: HASH_ZERO,
2445 parent: Parent(None),
2446 left: TreeIndex(1),
2447 right: TreeIndex(2),
2448 };
2449 node.get_sibling_side(TreeIndex(0))
2450 .expect_err("should fail");
2451 }
2452
2453 #[rstest]
2454 fn test_merkle_blob_to_from_path(traversal_blob: MerkleBlob) {
2455 let dir_path = tempfile::tempdir().unwrap();
2456 let file_path = dir_path.path().join("blob");
2457 traversal_blob.to_path(&file_path).unwrap();
2458 let loaded = MerkleBlob::from_path(&file_path).unwrap();
2459
2460 assert!(blob_tree_equality(&traversal_blob, &loaded));
2461 assert_eq!(traversal_blob.blob, loaded.blob);
2462 }
2463
2464 #[rstest]
2465 fn test_get_node_by_hash(small_blob: MerkleBlob) {
2466 let node = small_blob.get_node_by_hash(sha256_num(&0x1020)).unwrap();
2467
2468 #[allow(clippy::needless_raw_string_hashes)]
2469 let expected = expect![[r#"
2470 (
2471 KeyId(
2472 283686952306183,
2473 ),
2474 ValueId(
2475 1157726452361532951,
2476 ),
2477 )
2478 "#]];
2479
2480 expected.assert_debug_eq(&node);
2481 }
2482
2483 #[rstest]
2484 fn test_get_node_by_hash_fails_not_found(small_blob: MerkleBlob) {
2485 let result = small_blob.get_node_by_hash(sha256_num(&27));
2486
2487 #[allow(clippy::needless_raw_string_hashes)]
2488 let expected = expect![[r#"
2489 Err(
2490 LeafHashNotFound(
2491 Hash(
2492 688e94a51ee508a95e761294afb7a6004b432c15d9890c80ddf23bde8caa4c26,
2493 ),
2494 ),
2495 )
2496 "#]];
2497
2498 expected.assert_debug_eq(&result);
2499 }
2500
2501 #[rstest]
2502 fn test_get_hashes_indexes(small_blob: MerkleBlob) {
2503 let hashes_indexes = small_blob.get_hashes_indexes(false).unwrap();
2504
2505 let mut expected = HashMap::new();
2506 let one = sha256_num(&0x2030);
2507 let two = sha256_num(&0x1020);
2508 let zero = internal_hash(&one, &two);
2509 expected.insert(zero, TreeIndex(0));
2510 expected.insert(one, TreeIndex(1));
2511 expected.insert(two, TreeIndex(2));
2512
2513 assert_eq!(hashes_indexes, expected);
2514 }
2515
2516 #[rstest]
2517 fn test_get_hashes_indexes_leafs_only(small_blob: MerkleBlob) {
2518 let hashes_indexes = small_blob.get_hashes_indexes(true).unwrap();
2519
2520 let mut expected = HashMap::new();
2521 let one = sha256_num(&0x2030);
2522 let two = sha256_num(&0x1020);
2523 expected.insert(one, TreeIndex(1));
2524 expected.insert(two, TreeIndex(2));
2525
2526 assert_eq!(hashes_indexes, expected);
2527 }
2528
2529 #[rstest]
2530 fn test_get_hashes_indexes_empty() {
2531 let blob = MerkleBlob::new(Vec::new()).unwrap();
2532 let result = blob.get_hashes_indexes(false).unwrap();
2533
2534 assert_eq!(result, HashMap::new());
2535 }
2536
2537 #[rstest]
2538 fn test_node_set_hash(
2539 #[values(
2540 Node::Internal(InternalNode{hash: HASH_ZERO, parent: Parent(None), left: TreeIndex(0), right: TreeIndex(1)}),
2541 Node::Leaf(LeafNode{hash:HASH_ZERO, parent: Parent(None), key: KeyId(0), value: ValueId(0)}),
2542 )]
2543 mut node: Node,
2544 ) {
2545 assert_eq!(node.hash(), HASH_ZERO);
2546 node.set_hash(HASH_ONE);
2547 assert_eq!(node.hash(), HASH_ONE);
2548 }
2549
2550 #[rstest]
2551 fn test_remove_not_present_leaf_from_block_status_cache(mut small_blob: MerkleBlob) {
2552 let key = KeyId(10948);
2553 let leaf = LeafNode {
2554 hash: HASH_ZERO,
2555 parent: Parent(None),
2556 key,
2557 value: ValueId(0),
2558 };
2559 let result = small_blob.block_status_cache.remove_leaf(&leaf);
2560
2561 #[allow(clippy::needless_raw_string_hashes)]
2562 let expected = expect![[r#"
2563 Err(
2564 UnknownKey(
2565 KeyId(
2566 10948,
2567 ),
2568 ),
2569 )
2570 "#]];
2571
2572 expected.assert_debug_eq(&result);
2573 }
2574
2575 #[rstest]
2576 fn test_insert_past_extend_entry_fails(mut small_blob: MerkleBlob) {
2577 let index = TreeIndex(small_blob.extend_index().0 + 1);
2578 let block = Block {
2579 metadata: NodeMetadata {
2580 node_type: NodeType::Leaf,
2581 dirty: true,
2582 },
2583 node: Node::Internal(InternalNode {
2584 hash: HASH_ZERO,
2585 parent: Parent(None),
2586 left: TreeIndex(0),
2587 right: TreeIndex(0),
2588 }),
2589 };
2590 let error = small_blob.insert_entry_to_blob(index, &block);
2591
2592 #[allow(clippy::needless_raw_string_hashes)]
2593 let expected = expect![[r#"
2594 Err(
2595 BlockIndexOutOfBounds(
2596 TreeIndex(
2597 4,
2598 ),
2599 ),
2600 )
2601 "#]];
2602 expected.assert_debug_eq(&error);
2603 }
2604
2605 #[rstest]
2606 fn test_get_key_index(small_blob: MerkleBlob) {
2607 let key = KeyId(0x0001_0203_0405_0607);
2608 let index = small_blob.get_key_index(key).unwrap();
2609 assert_eq!(index, TreeIndex(2));
2610 }
2611
2612 #[rstest]
2613 fn test_block_status_cache_move_index_invalid_source(mut traversal_blob: MerkleBlob) {
2614 let key = KeyId(307);
2615 let index = traversal_blob.get_key_index(key).unwrap();
2616 traversal_blob.delete(key).unwrap();
2617 assert!(
2618 traversal_blob
2619 .block_status_cache
2620 .free_indexes
2621 .contains(&index)
2622 );
2623 let result = traversal_blob.block_status_cache.move_index(index, index);
2624 #[allow(clippy::needless_raw_string_hashes)]
2625 let expected = expect![[r#"
2626 Err(
2627 MoveSourceIndexNotInUse(
2628 TreeIndex(
2629 5,
2630 ),
2631 ),
2632 )
2633 "#]];
2634
2635 expected.assert_debug_eq(&result);
2636 }
2637
2638 #[rstest]
2639 fn test_block_status_cache_move_index_invalid_destination(mut traversal_blob: MerkleBlob) {
2640 let key = KeyId(307);
2641 let index = traversal_blob.get_key_index(key).unwrap();
2642 traversal_blob.delete(key).unwrap();
2643 assert!(
2644 traversal_blob
2645 .block_status_cache
2646 .free_indexes
2647 .contains(&index)
2648 );
2649 let result = traversal_blob
2650 .block_status_cache
2651 .move_index(TreeIndex(0), index);
2652 #[allow(clippy::needless_raw_string_hashes)]
2653 let expected = expect![[r#"
2654 Err(
2655 MoveDestinationIndexNotInUse(
2656 TreeIndex(
2657 5,
2658 ),
2659 ),
2660 )
2661 "#]];
2662
2663 expected.assert_debug_eq(&result);
2664 }
2665
2666 #[rstest]
2667 fn test_moved_sibling_retains_hash(mut small_blob: MerkleBlob) {
2668 let key_to_delete = KeyId(0x0001_0203_0405_0607);
2669 let remaining_hash = sha256_num(&0x2030);
2670 assert_ne!(small_blob.get_hash(TreeIndex(0)).unwrap(), remaining_hash);
2671 small_blob.delete(key_to_delete).unwrap();
2672 assert_eq!(small_blob.get_hash(TreeIndex(0)).unwrap(), remaining_hash);
2673 }
2674}