Processing storage tree return metadata (#169)

* define HashTrace to the smt trace * process_block return tree_metadata
Sin7Y · Nov 27, 2023 · 9f0263f · 9f0263f
1 parent 1906d93
commit 9f0263f
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 40 deletions.
diff --git a/core/src/merkle_tree/patch.rs b/core/src/merkle_tree/patch.rs
@@ -2,7 +2,7 @@ use super::iter_ext::IteratorExt;
 use super::TreeError;
 use crate::crypto::hash::Hasher;
 
-use crate::trace::trace::PoseidonRow;
+use crate::trace::trace::{PoseidonRow, HashTrace};
 use crate::types::merkle_tree::constant::ROOT_TREE_DEPTH;
 use crate::types::merkle_tree::{
     tree_key_to_u256, u256_to_tree_key, NodeEntry, TreeKey, TreeKeyU256, TreeValue,
@@ -117,14 +117,7 @@ impl UpdatesBatch {
                 Mutex<
                     Vec<(
                         usize,
-                        (
-                            PoseidonRow,
-                            TreeValue,
-                            TreeValue,
-                            TreeValue,
-                            TreeValue,
-                            PoseidonRow,
-                        ),
+                        HashTrace,
                     )>,
                 >,
             >,
@@ -250,7 +243,7 @@ impl UpdatesBatch {
 
                                 hash_trace.lock().unwrap().push((
                                     *index,
-                                    (
+                                    HashTrace(
                                         hash.1,
                                         current_hash.clone(),
                                         nei_hash.clone(),

diff --git a/core/src/merkle_tree/tree.rs b/core/src/merkle_tree/tree.rs
@@ -11,6 +11,7 @@ use crate::types::merkle_tree::{
     LevelIndex, NodeEntry, TreeKey, TreeMetadata, TreeOperation, TreeValue, ZkHash,
 };
 use crate::types::proof::StorageLogMetadata;
+use crate::trace::trace::HashTrace;
 use itertools::Itertools;
 use log::{debug, info};
 use std::borrow::{Borrow, BorrowMut};
@@ -102,32 +103,19 @@ impl AccountTree {
     pub fn process_block<I>(
         &mut self,
         storage_logs: I,
-    ) -> Vec<(
-        PoseidonRow,
-        TreeValue,
-        TreeValue,
-        TreeValue,
-        TreeValue,
-        PoseidonRow,
-    )>
+    ) -> (Vec<HashTrace>, Option<TreeMetadata>)
     where
         I: IntoIterator,
         I::Item: Borrow<WitnessStorageLog>,
     {
-        self.process_blocks(once(storage_logs))
+        let (hash_traces, tree_metadata) = self.process_blocks(once(storage_logs));
+        (hash_traces, tree_metadata.last().cloned())
     }
 
     pub fn process_blocks<I>(
         &mut self,
         blocks: I,
-    ) -> Vec<(
-        PoseidonRow,
-        TreeValue,
-        TreeValue,
-        TreeValue,
-        TreeValue,
-        PoseidonRow,
-    )>
+    ) -> (Vec<HashTrace>, Vec<TreeMetadata>)
     where
         I: IntoIterator,
         I::Item: IntoIterator,
@@ -178,14 +166,7 @@ impl AccountTree {
         &mut self,
         updates_batch: Vec<Vec<(TreeKey, TreeOperation)>>,
     ) -> Result<
-        Vec<(
-            PoseidonRow,
-            TreeValue,
-            TreeValue,
-            TreeValue,
-            TreeValue,
-            PoseidonRow,
-        )>,
+        (Vec<HashTrace>, Vec<TreeMetadata>),
         TreeError,
     > {
         let total_blocks = updates_batch.len();
@@ -237,7 +218,7 @@ impl AccountTree {
             .flat_map(|e| e)
             .collect();
 
-        let _tree_metadata: Vec<_> = {
+        let tree_metadata: Vec<_> = {
             let patch_metadata =
                 self.apply_patch(updates.0, &storage_logs_with_blocks, &leaf_indices);
 
@@ -278,7 +259,7 @@ impl AccountTree {
         };
 
         self.block_number += total_blocks as u32;
-        Ok(hash_trace)
+        Ok((hash_trace, tree_metadata))
     }
 
     /// Prepares all the data which will be needed to calculate new Merkle Trees

diff --git a/core/src/trace/trace.rs b/core/src/trace/trace.rs
@@ -1,6 +1,5 @@
 use crate::program::REGISTER_NUM;
-use crate::types::account::Address;
-use crate::types::merkle_tree::TreeValue;
+use crate::types::{account::Address, merkle_tree::TreeValue};
 use crate::utils::split_limbs_from_field;
 use crate::utils::split_u16_limbs_from_field;
 use plonky2::field::goldilocks_field::GoldilocksField;
@@ -259,6 +258,9 @@ impl Display for PoseidonRow {
     }
 }
 
+#[derive(Debug, Copy, Clone, Default, Serialize, Deserialize)]
+pub struct HashTrace(pub PoseidonRow, pub TreeValue, pub TreeValue, pub TreeValue, pub TreeValue, pub PoseidonRow);
+
 #[derive(Debug, Copy, Clone, Serialize, Deserialize)]
 pub struct StorageRow {
     pub tx_idx: GoldilocksField,

diff --git a/executor/src/trace.rs b/executor/src/trace.rs
@@ -204,7 +204,7 @@ pub fn gen_storage_hash_table(
     let mut trace = std::mem::replace(&mut process.storage_log, Vec::new());
     trace.extend(std::mem::replace(&mut process.program_log, Vec::new()));
     let mut pre_root = account_tree.root_hash();
-    let hash_traces = account_tree.process_block(trace.iter());
+    let (hash_traces, _) = account_tree.process_block(trace.iter());
     let _ = account_tree.save();
 
     let mut root_hashes = Vec::new();