Make JMT optional

radixdlt · Oct 12, 2023 · 038793e · 038793e
1 parent 2d24dba
commit 038793e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/radix-engine-stores/src/hash_tree_support.rs b/radix-engine-stores/src/hash_tree_support.rs
@@ -12,15 +12,17 @@ pub struct HashTreeUpdatingDatabase<D> {
     tree_store: TypedInMemoryTreeStore,
     current_version: Version,
     current_hash: Hash,
+    jmt_enabled: bool,
 }
 
 impl<D> HashTreeUpdatingDatabase<D> {
-    pub fn new(underlying: D) -> Self {
+    pub fn new(underlying: D, jmt_enabled: bool) -> Self {
         HashTreeUpdatingDatabase {
             underlying,
             tree_store: TypedInMemoryTreeStore::with_pruning(),
             current_version: 0,
             current_hash: Hash([0; Hash::LENGTH]),
+            jmt_enabled,
         }
     }
 
@@ -37,11 +39,13 @@ impl<D> HashTreeUpdatingDatabase<D> {
     }
 
     fn update_with(&mut self, db_updates: &DatabaseUpdates) {
-        self.current_hash = put_at_next_version(
-            &mut self.tree_store,
-            Some(self.current_version).filter(|version| *version > 0),
-            db_updates,
-        );
+        if self.jmt_enabled {
+            self.current_hash = put_at_next_version(
+                &mut self.tree_store,
+                Some(self.current_version).filter(|version| *version > 0),
+                db_updates,
+            );
+        }
         self.current_version += 1;
     }
 }

diff --git a/scrypto-unit/src/test_runner.rs b/scrypto-unit/src/test_runner.rs
@@ -329,7 +329,7 @@ impl<E: NativeVmExtension, D: TestDatabase> TestRunnerBuilder<E, D> {
         TestRunnerBuilder {
             custom_genesis: self.custom_genesis,
             custom_extension: self.custom_extension,
-            custom_database: HashTreeUpdatingDatabase::new(self.custom_database),
+            custom_database: HashTreeUpdatingDatabase::new(self.custom_database, true),
             trace: self.trace,
             skip_receipt_check: false,
         }

diff --git a/simulator/src/replay/cmd_run_in_memory.rs b/simulator/src/replay/cmd_run_in_memory.rs
@@ -28,6 +28,8 @@ pub struct RunInMemory {
     /// The max version to execute
     #[clap(short, long)]
     pub max_version: Option<u64>,
+    #[clap(long)]
+    pub enable_jmt: bool,
 }
 
 impl RunInMemory {
@@ -42,7 +44,7 @@ impl RunInMemory {
         let mut archive = Archive::new(tar);
 
         let in_memory = InMemorySubstateDatabase::standard();
-        let mut database = HashTreeUpdatingDatabase::new(in_memory);
+        let mut database = HashTreeUpdatingDatabase::new(in_memory, self.enable_jmt);
         let scrypto_vm = ScryptoVm::<DefaultWasmEngine>::default();
         let start = std::time::Instant::now();
         for entry in archive.entries().map_err(Error::IOError)? {