From 157a1a93a43e2045cd698037259b6c433e5ed04a Mon Sep 17 00:00:00 2001
From: Yulong Wu <yulong.wu@rdx.works>
Date: Fri, 13 Oct 2023 12:58:22 +0100
Subject: [PATCH] Add pruning support to RocksDBWithMerkleTreeSubstateStore

---
 .../src/rocks_db_with_merkle_tree/mod.rs      | 68 ++++++++++++++++---
 1 file changed, 59 insertions(+), 9 deletions(-)

diff --git a/radix-engine-stores/src/rocks_db_with_merkle_tree/mod.rs b/radix-engine-stores/src/rocks_db_with_merkle_tree/mod.rs
index 327993fd268..3b00d11f1aa 100644
--- a/radix-engine-stores/src/rocks_db_with_merkle_tree/mod.rs
+++ b/radix-engine-stores/src/rocks_db_with_merkle_tree/mod.rs
@@ -1,5 +1,5 @@
 use crate::hash_tree::tree_store::{
-    encode_key, NodeKey, ReadableTreeStore, TreeNode, VersionedTreeNode,
+    encode_key, NodeKey, ReadableTreeStore, StaleTreePart, TreeNode, TreeNodeV1, VersionedTreeNode,
 };
 use itertools::Itertools;
 use radix_engine_common::data::scrypto::{scrypto_decode, scrypto_encode};
@@ -25,6 +25,7 @@ const STALE_MERKLE_TREE_PARTS_CF: &str = "stale_merkle_tree_parts";
 
 pub struct RocksDBWithMerkleTreeSubstateStore {
     db: DBWithThreadMode<SingleThreaded>,
+    pruning_enabled: bool,
 }
 
 impl RocksDBWithMerkleTreeSubstateStore {
@@ -41,10 +42,10 @@ impl RocksDBWithMerkleTreeSubstateStore {
         let mut options = Options::default();
         options.create_if_missing(true);
         options.create_missing_column_families(true);
-        Self::with_options(&options, root)
+        Self::with_options(&options, root, true)
     }
 
-    pub fn with_options(options: &Options, root: PathBuf) -> Self {
+    pub fn with_options(options: &Options, root: PathBuf, pruning_enabled: bool) -> Self {
         let db = DB::open_cf_descriptors(
             options,
             root.as_path(),
@@ -59,7 +60,10 @@ impl RocksDBWithMerkleTreeSubstateStore {
             .collect::<Vec<_>>(),
         )
         .unwrap();
-        Self { db }
+        Self {
+            db,
+            pruning_enabled,
+        }
     }
 
     fn cf(&self, cf: &str) -> &ColumnFamily {
@@ -201,11 +205,14 @@ impl CommittableSubstateDatabase for RocksDBWithMerkleTreeSubstateStore {
                 scrypto_encode(&VersionedTreeNode::new_latest(node)).unwrap(),
             );
         }
-        batch.put_cf(
-            self.cf(STALE_MERKLE_TREE_PARTS_CF),
-            next_state_version.to_be_bytes(),
-            scrypto_encode(&state_hash_tree_update.stale_tree_parts).unwrap(),
-        );
+        if !self.pruning_enabled {
+            // If pruning is not enabled, we store the stale nodes in DB.
+            batch.put_cf(
+                self.cf(STALE_MERKLE_TREE_PARTS_CF),
+                next_state_version.to_be_bytes(),
+                scrypto_encode(&state_hash_tree_update.stale_tree_parts).unwrap(),
+            );
+        }
 
         // update the metadata
         batch.put_cf(
@@ -220,6 +227,49 @@ impl CommittableSubstateDatabase for RocksDBWithMerkleTreeSubstateStore {
 
         // flush the batch
         self.db.write(batch).unwrap();
+
+        if self.pruning_enabled {
+            for part in state_hash_tree_update.stale_tree_parts {
+                match part {
+                    StaleTreePart::Node(node_key) => {
+                        self.db
+                            .delete_cf(self.cf(MERKLE_NODES_CF), encode_key(&node_key))
+                            .unwrap();
+                    }
+                    StaleTreePart::Subtree(node_key) => {
+                        let mut queue = VecDeque::new();
+                        queue.push_back(node_key);
+
+                        while let Some(node_key) = queue.pop_front() {
+                            if let Some(bytes) = self
+                                .db
+                                .get_cf(self.cf(MERKLE_NODES_CF), encode_key(&node_key))
+                                .unwrap()
+                            {
+                                self.db
+                                    .delete_cf(self.cf(MERKLE_NODES_CF), encode_key(&node_key))
+                                    .unwrap();
+                                let value: VersionedTreeNode = scrypto_decode(&bytes).unwrap();
+                                match value.into_latest() {
+                                    TreeNodeV1::Internal(x) => {
+                                        for child in x.children {
+                                            queue.push_back(
+                                                node_key.gen_child_node_key(
+                                                    child.version,
+                                                    child.nibble,
+                                                ),
+                                            )
+                                        }
+                                    }
+                                    TreeNodeV1::Leaf(_) => {}
+                                    TreeNodeV1::Null => {}
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
     }
 }