Skip to main content

chia_datalayer/merkle/
blob.rs

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
35// assumptions
36// - root is at index 0
37// - any case with no keys will have a zero length blob
38
39pub 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    // error: Unit variant `Auto` is not yet supported in a complex enum
76    // = help: change to a struct variant with no fields: `Auto { }`
77    // = note: the enum is complex because of non-unit variant `Leaf`
78    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        // to be called _after_ having written to the destination index
212        // TODO: not checking it is within bounds of the present blob
213        if self.free_indexes.contains(&source) {
214            return Err(Error::MoveSourceIndexNotInUse(source));
215        }
216        // TODO: not checking it is within bounds of the present blob
217        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/// Stores a DataLayer merkle tree in bytes and provides serialization on each access so that only
299/// the parts presently in use are stored in active objects.  The bytes are grouped as blocks of
300/// equal size regardless of being internal vs. external nodes so that block indexes can be used
301/// for references to particular nodes and readily converted to byte indexes.  The leaf nodes
302/// do not hold the DataLayer key and value data but instead an id for each of the key and value
303/// such that the code using a merkle blob can store the key and value as they see fit.  Each node
304/// stores the hash for the merkle aspect of the tree.
305#[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    // TODO: used by fuzzing, some cleaner way?  making it cfg-dependent is annoying with
311    //       the type stubs
312    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        // OPT: perhaps go back to taking an iterator?
575        // OPT: would it be worthwhile to hold the entire blocks?
576        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        // OPT: can we insert the top node first?  maybe more efficient to update it's children
606        //      than to update the parents of the children when traversing leaf to sub-root?
607        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            // OPT: can we avoid this extra min height leaf traversal?
652            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        // TODO: seems like this ought to be fairly similar to regular insert
666
667        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        // OPT: maybe just edit in place?
802        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                // NOTE: explicitly not marking index as free since that would hazard two
925                //       sequential calls to this function through this path to both return
926                //       the same index
927                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        // NOTE: zero means left here but right below
944        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 iter(&self) -> MerkleBlobLeftChildFirstIterator<'_> {
1106    //     <&Self as IntoIterator>::into_iter(self)
1107    // }
1108
1109    pub fn calculate_lazy_hashes(&mut self) -> Result<(), Error> {
1110        // OPT: yeah, storing the whole set of blocks via collect is not great
1111        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            // OPT: obviously inefficient to re-get/deserialize these blocks inside
1125            //      an iteration that's already doing that
1126            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        // first in the lineage is the index itself, second is the first parent
1172        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    // it is known that memo is unused here, but is part of the interface of deepcopy
1402    #[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        // TODO: should be a Side, but python has a different Side right now
1417        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        // NOTE: this is checking tree structure equality, not serialized bytes equality
1645        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                // NOTE: this is effectively checked by the controlled overall traversal
1659                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        // in Python: Program.to((left_hash, right_hash)).get_tree_hash_precalc(left_hash, right_hash)
1790
1791        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        // just allowing parallelism of testing 100,000 inserts total
1879        #[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                // NOTE: yeah this hash is garbage
1891                .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        // there must be enough nodes to avoid the few-node insertion methods that clear the blob
2312        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}