diff --git a/piecrust/src/store.rs b/piecrust/src/store.rs index cee3a81d..3f45a7a9 100644 --- a/piecrust/src/store.rs +++ b/piecrust/src/store.rs @@ -478,7 +478,6 @@ impl Commit { Some(contract.page_indices.into_iter().map(move |page_index| { let tree_opening = self .contracts_merkle - .tree .opening(pos) .expect("There must be a leaf for the contract"); @@ -520,7 +519,6 @@ impl Commit { } self.contracts_merkle - .tree .insert(position_from_contract(&contract_id), *element.tree.root()); } @@ -530,7 +528,7 @@ impl Commit { } pub fn root(&self) -> Ref { - self.contracts_merkle.tree.root() + self.contracts_merkle.root() } } diff --git a/piecrust/src/store/tree.rs b/piecrust/src/store/tree.rs index 229fb896..6eecdf47 100644 --- a/piecrust/src/store/tree.rs +++ b/piecrust/src/store/tree.rs @@ -148,12 +148,39 @@ pub struct NewContractIndex { #[derive(Debug, Clone, Archive, Deserialize, Serialize)] #[archive_attr(derive(CheckBytes))] pub struct ContractsMerkle { - pub tree: Tree, + inner_tree: Tree, + dict: BTreeMap, } impl Default for ContractsMerkle { fn default() -> Self { - Self { tree: Tree::new() } + Self { + inner_tree: Tree::new(), + dict: BTreeMap::new(), + } + } +} + +impl ContractsMerkle { + pub fn insert(&mut self, pos: u64, hash: Hash) { + let new_pos = match self.dict.get(&pos) { + None => { + let new_pos = (self.dict.len() + 1) as u64; + self.dict.insert(pos, new_pos); + new_pos + } + Some(p) => *p, + }; + self.inner_tree.insert(new_pos, hash); + } + + pub fn opening(&self, pos: u64) -> Option { + let new_pos = self.dict.get(&pos).expect("pos should exist in dict"); + self.inner_tree.opening(*new_pos) + } + + pub fn root(&self) -> Ref { + self.inner_tree.root() } }