piecrust: introduced tree positions file

piecrust/src/store.rs

@@ -17,7 +17,7 @@ mod tree;
 use std::cell::Ref;
 use std::collections::btree_map::Entry::*;
 use std::collections::btree_map::Keys;
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::fmt::{Debug, Formatter};
 use std::fs::create_dir_all;
 use std::path::{Path, PathBuf};
@@ -32,6 +32,7 @@ use tree::{Hash, NewContractIndex};
 use crate::store::commit::Hulk;
 use crate::store::tree::{
     position_from_contract, BaseInfo, ContractIndexElement, ContractsMerkle,
+    TreePos,
 };
 pub use bytecode::Bytecode;
 pub use memory::{Memory, PAGE_SIZE};
@@ -44,6 +45,7 @@ const BYTECODE_DIR: &str = "bytecode";
 const MEMORY_DIR: &str = "memory";
 const LEAF_DIR: &str = "leaf";
 const BASE_FILE: &str = "base";
+const TREE_POS_FILE: &str = "tree_pos";
 const ELEMENT_FILE: &str = "element";
 const OBJECTCODE_EXTENSION: &str = "a";
 const METADATA_EXTENSION: &str = "m";
@@ -362,6 +364,16 @@ fn base_path_main<P: AsRef<Path>, S: AsRef<str>>(
     Ok(dir.join(BASE_FILE))
 }
 
+fn tree_pos_path_main<P: AsRef<Path>, S: AsRef<str>>(
+    main_dir: P,
+    commit_id: S,
+) -> io::Result<PathBuf> {
+    let commit_id = commit_id.as_ref();
+    let dir = main_dir.as_ref().join(commit_id);
+    fs::create_dir_all(&dir)?;
+    Ok(dir.join(TREE_POS_FILE))
+}
+
 fn commit_id_to_hash<S: AsRef<str>>(commit_id: S) -> Hash {
     let hash: [u8; 32] = hex::decode(commit_id.as_ref())
         .expect("Hex decoding of commit id string should succeed")
@@ -408,8 +420,20 @@ fn commit_from_dir<P: AsRef<Path>>(
     let leaf_dir = main_dir.join(LEAF_DIR);
     tracing::trace!("before index_merkle_from_path");
 
-    let index =
-        index_from_path(main_dir, leaf_dir, &maybe_hash, commit_store.clone())?;
+    let tree_pos = if let Some(ref hash_hex) = commit_id {
+        let tree_pos_path = main_dir.join(hash_hex).join(TREE_POS_FILE);
+        Some(tree_pos_from_path(tree_pos_path)?.tree_pos)
+    } else {
+        None
+    };
+
+    let (index, contracts_merkle) = index_merkle_from_path(
+        main_dir,
+        leaf_dir,
+        &maybe_hash,
+        commit_store.clone(),
+        tree_pos.as_ref(),
+    )?;
     tracing::trace!("after index_merkle_from_path");
 
     let bytecode_dir = main_dir.join(BYTECODE_DIR);
@@ -465,12 +489,11 @@ fn commit_from_dir<P: AsRef<Path>>(
         }
     }
 
-    let (base, contracts_merkle) = if let Some(hash_hex) = commit_id {
+    let base = if let Some(ref hash_hex) = commit_id {
         let base_info_path = main_dir.join(hash_hex).join(BASE_FILE);
-        let base_info = base_from_path(base_info_path)?;
-        (base_info.maybe_base, base_info.contracts_merkle)
+        base_from_path(base_info_path)?.maybe_base
     } else {
-        (None, ContractsMerkle::default())
+        None
     };
 
     Ok(Commit {
@@ -482,15 +505,17 @@ fn commit_from_dir<P: AsRef<Path>>(
     })
 }
 
-fn index_from_path(
+fn index_merkle_from_path(
     main_path: impl AsRef<Path>,
     leaf_dir: impl AsRef<Path>,
     maybe_commit_id: &Option<Hash>,
     commit_store: Arc<Mutex<CommitStore>>,
-) -> io::Result<NewContractIndex> {
+    maybe_tree_pos: Option<&HashMap<ContractId, u64>>,
+) -> io::Result<(NewContractIndex, ContractsMerkle)> {
     let leaf_dir = leaf_dir.as_ref();
 
     let mut index: NewContractIndex = NewContractIndex::new();
+    let mut merkle: ContractsMerkle = ContractsMerkle::default();
 
     for entry in fs::read_dir(leaf_dir)? {
         let entry = entry?;
@@ -521,6 +546,30 @@ fn index_from_path(
                         ),
                         )
                     })?;
+                if let Some(h) = element.hash() {
+                    match maybe_tree_pos {
+                        Some(tree_pos) => {
+                            if let Some(pos) = tree_pos.get(&contract_id) {
+                                merkle.insert_with_int_pos(
+                                    position_from_contract(&contract_id),
+                                    *pos,
+                                    h,
+                                    &contract_id,
+                                );
+                            }
+                        }
+                        None => {
+                            merkle.insert_with_int_pos(
+                                position_from_contract(&contract_id),
+                                element
+                                    .int_pos()
+                                    .expect("int pos should be present"),
+                                h,
+                                &contract_id,
+                            );
+                        }
+                    }
+                }
                 if element_depth != u32::MAX {
                     index.insert_contract_index(&contract_id, element);
                 } else {
@@ -533,7 +582,7 @@ fn index_from_path(
         }
     }
 
-    Ok(index)
+    Ok((index, merkle))
 }
 
 fn base_from_path<P: AsRef<Path>>(path: P) -> io::Result<BaseInfo> {
@@ -550,6 +599,20 @@ fn base_from_path<P: AsRef<Path>>(path: P) -> io::Result<BaseInfo> {
     Ok(base_info)
 }
 
+fn tree_pos_from_path<P: AsRef<Path>>(path: P) -> io::Result<TreePos> {
+    let path = path.as_ref();
+
+    let tree_pos_bytes = fs::read(path)?;
+    let tree_pos = rkyv::from_bytes(&tree_pos_bytes).map_err(|err| {
+        io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!("Invalid tree positions file \"{path:?}\": {err}"),
+        )
+    })?;
+
+    Ok(tree_pos)
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct Commit {
     index: NewContractIndex,
@@ -647,7 +710,7 @@ impl Commit {
 
         let root = *element.tree().root();
         let pos = position_from_contract(&contract_id);
-        let internal_pos = contracts_merkle.insert(pos, root);
+        let internal_pos = contracts_merkle.insert(pos, root, &contract_id);
         element.set_hash(Some(root));
         element.set_int_pos(Some(internal_pos));
     }
@@ -1048,10 +1111,8 @@ fn write_commit_inner<P: AsRef<Path>, S: AsRef<str>>(
     }
     tracing::trace!("persisting index finished");
 
-    base_info.contracts_merkle = commit.contracts_merkle.clone(); //todo: clone
-
     let base_main_path =
-        base_path_main(directories.main_dir, commit_id.as_ref())?;
+        base_path_main(&directories.main_dir, commit_id.as_ref())?;
     let base_info_bytes =
         rkyv::to_bytes::<_, 128>(&base_info).map_err(|err| {
             io::Error::new(
@@ -1061,6 +1122,19 @@ fn write_commit_inner<P: AsRef<Path>, S: AsRef<str>>(
         })?;
     fs::write(base_main_path, base_info_bytes)?;
 
+    let tree_pos_main_path =
+        tree_pos_path_main(&directories.main_dir, commit_id.as_ref())?;
+    let tree_pos_bytes = rkyv::to_bytes::<_, 128>(
+        commit.contracts_merkle.tree_pos(),
+    )
+    .map_err(|err| {
+        io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!("Failed serializing tree positions file: {err}"),
+        )
+    })?;
+    fs::write(tree_pos_main_path, tree_pos_bytes)?;
+
     Ok(())
 }
 
@@ -1101,6 +1175,7 @@ fn finalize_commit<P: AsRef<Path>>(
     let root = hex::encode(root);
     let commit_path = main_dir.join(&root);
     let base_info_path = commit_path.join(BASE_FILE);
+    let tree_pos_path = commit_path.join(TREE_POS_FILE);
     let base_info = base_from_path(&base_info_path)?;
     for contract_hint in base_info.contract_hints {
         let contract_hex = hex::encode(contract_hint);
@@ -1130,6 +1205,7 @@ fn finalize_commit<P: AsRef<Path>>(
     }
 
     fs::remove_file(base_info_path)?;
+    fs::remove_file(tree_pos_path)?;
     fs::remove_dir(commit_path)?;
 
     Ok(())

piecrust/src/store/tree.rs

@@ -4,6 +4,7 @@
 //
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
+use std::collections::HashMap;
 use std::{
     cell::Ref,
     collections::{BTreeMap, BTreeSet},
@@ -97,19 +98,26 @@ impl NewContractIndex {
 pub struct ContractsMerkle {
     inner_tree: Tree,
     dict: BTreeMap<u64, u64>,
+    tree_pos: HashMap<ContractId, u64>,
 }
 
 impl Default for ContractsMerkle {
     fn default() -> Self {
         Self {
             inner_tree: Tree::new(),
             dict: BTreeMap::new(),
+            tree_pos: HashMap::new(),
         }
     }
 }
 
 impl ContractsMerkle {
-    pub fn insert(&mut self, pos: u64, hash: Hash) -> u64 {
+    pub fn insert(
+        &mut self,
+        pos: u64,
+        hash: Hash,
+        contract_id: &ContractId,
+    ) -> u64 {
         let new_pos = match self.dict.get(&pos) {
             None => {
                 let new_pos = (self.dict.len() + 1) as u64;
@@ -119,12 +127,20 @@ impl ContractsMerkle {
             Some(p) => *p,
         };
         self.inner_tree.insert(new_pos, hash);
+        self.tree_pos.insert(*contract_id, new_pos);
         new_pos
     }
 
-    pub fn insert_with_int_pos(&mut self, pos: u64, int_pos: u64, hash: Hash) {
+    pub fn insert_with_int_pos(
+        &mut self,
+        pos: u64,
+        int_pos: u64,
+        hash: Hash,
+        contract_id: &ContractId,
+    ) {
         self.dict.insert(pos, int_pos);
         self.inner_tree.insert(int_pos, hash);
+        self.tree_pos.insert(*contract_id, int_pos);
     }
 
     pub fn opening(&self, pos: u64) -> Option<TreeOpening> {
@@ -135,6 +151,10 @@ impl ContractsMerkle {
     pub fn root(&self) -> Ref<Hash> {
         self.inner_tree.root()
     }
+
+    pub fn tree_pos(&self) -> &HashMap<ContractId, u64> {
+        &self.tree_pos
+    }
 }
 
 #[derive(Debug, Clone, Archive, Deserialize, Serialize)]
@@ -151,7 +171,12 @@ pub struct ContractIndex {
 pub struct BaseInfo {
     pub contract_hints: Vec<ContractId>,
     pub maybe_base: Option<Hash>,
-    pub contracts_merkle: ContractsMerkle,
+}
+
+#[derive(Debug, Clone, Default, Archive, Deserialize, Serialize)]
+#[archive_attr(derive(CheckBytes))]
+pub struct TreePos {
+    pub tree_pos: HashMap<ContractId, u64>,
 }
 
 #[derive(Debug, Clone, Archive, Deserialize, Serialize)]