drk/cache: overlay added along with its smt impl

bin/drk/src/cache.rs

@@ -16,23 +16,35 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use darkfi::Result;
-use sled_overlay::sled;
+use darkfi::{blockchain::HeaderHash, Result};
+use darkfi_sdk::{
+    crypto::{
+        pasta_prelude::PrimeField,
+        smt::{PoseidonFp, SparseMerkleTree, StorageAdapter, SMT_FP_DEPTH},
+        MerkleTree,
+    },
+    error::{ContractError, ContractResult},
+    pasta::pallas,
+};
+use darkfi_serial::serialize;
+use log::error;
+use num_bigint::BigUint;
+use sled_overlay::{sled, SledDbOverlay, SledDbOverlayStateDiff};
 
-pub const SLED_ORDER_TREE: &[u8] = b"_order";
+pub const SLED_SCANNED_BLOCKS_TREE: &[u8] = b"_scanned_blocks";
 pub const SLED_STATE_INVERSE_DIFF_TREE: &[u8] = b"_state_inverse_diff";
 pub const SLED_MERKLE_TREES_TREE: &[u8] = b"_merkle_trees";
 pub const SLED_MONEY_SMT_TREE: &[u8] = b"_money_smt";
 
-/// Structure holding all sled trees that define the Blockchain cache.
+/// Structure holding all sled trees that define the blockchain cache.
 #[derive(Clone)]
 pub struct Cache {
     /// Main pointer to the sled db connection
     pub sled_db: sled::Db,
-    /// The `sled` tree storing the order of the blockchain's blocks,
+    /// The `sled` tree storing the scanned blocks from the blockchain,
     /// where the key is the height number, and the value is the blocks'
     /// hash.
-    pub order: sled::Tree,
+    pub scanned_blocks: sled::Tree,
     /// The `sled` tree storing each blocks' full database state inverse
     /// changes, where the key is the block height number, and the value
     /// is the serialized database inverse diff.
@@ -41,7 +53,7 @@ pub struct Cache {
     /// where the key is the tree name, and the value is the serialized
     /// merkle tree itself.
     pub merkle_trees: sled::Tree,
-    /// The `sled` tree storing the Sparse Merkle Tree of the money
+    /// The `sled` tree storing the Sparse Merkle Tree of the Money
     /// contract.
     pub money_smt: sled::Tree,
 }
@@ -49,11 +61,213 @@ pub struct Cache {
 impl Cache {
     /// Instantiate a new `Cache` with the given `sled` database.
     pub fn new(db: &sled::Db) -> Result<Self> {
-        let order = db.open_tree(SLED_ORDER_TREE)?;
+        let scanned_blocks = db.open_tree(SLED_SCANNED_BLOCKS_TREE)?;
         let state_inverse_diff = db.open_tree(SLED_STATE_INVERSE_DIFF_TREE)?;
         let merkle_trees = db.open_tree(SLED_MERKLE_TREES_TREE)?;
         let money_smt = db.open_tree(SLED_MONEY_SMT_TREE)?;
 
-        Ok(Self { sled_db: db.clone(), order, state_inverse_diff, merkle_trees, money_smt })
+        Ok(Self {
+            sled_db: db.clone(),
+            scanned_blocks,
+            state_inverse_diff,
+            merkle_trees,
+            money_smt,
+        })
+    }
+}
+
+/// Overlay structure over a [`Cache`] instance.
+pub struct CacheOverlay(SledDbOverlay);
+
+impl CacheOverlay {
+    /// Instantiate a new `CacheOverlay` over the given [`Cache`] instance.
+    pub fn new(cache: &Cache) -> Result<CacheOverlay> {
+        // Here we configure all our cache sled trees to be protected in the overlay
+        let protected_trees = vec![
+            SLED_SCANNED_BLOCKS_TREE,
+            SLED_STATE_INVERSE_DIFF_TREE,
+            SLED_MERKLE_TREES_TREE,
+            SLED_MONEY_SMT_TREE,
+        ];
+        let mut overlay = SledDbOverlay::new(&cache.sled_db, protected_trees);
+
+        // Open all our cache sled trees in the overlay
+        overlay.open_tree(SLED_SCANNED_BLOCKS_TREE, true)?;
+        overlay.open_tree(SLED_STATE_INVERSE_DIFF_TREE, true)?;
+        overlay.open_tree(SLED_MERKLE_TREES_TREE, true)?;
+        overlay.open_tree(SLED_MONEY_SMT_TREE, true)?;
+
+        Ok(Self(overlay))
+    }
+
+    /// Insert a `u32` and a block hash into overlay's scanned blocks
+    /// tree. The block height is used as the key, and the blockhash is
+    /// used as value.
+    pub fn insert_scanned_block(&mut self, height: &u32, hash: &HeaderHash) -> Result<()> {
+        self.0.insert(SLED_SCANNED_BLOCKS_TREE, &height.to_be_bytes(), hash.inner())?;
+        Ok(())
+    }
+
+    /// Insert a `u32` and a block inverse diff into overlay's inverse
+    /// diffs tree. The block height is used as the key, and the
+    /// serialized database inverse diff is used as value.
+    pub fn insert_state_inverse_diff(
+        &mut self,
+        height: &u32,
+        diff: &SledDbOverlayStateDiff,
+    ) -> Result<()> {
+        self.0.insert(SLED_STATE_INVERSE_DIFF_TREE, &height.to_be_bytes(), &serialize(diff))?;
+        Ok(())
+    }
+
+    /// Insert a bytes slice and a merkle tree into overlay's merkle
+    /// trees tree. The provided bytes slice is used as the key, and
+    /// the serialized merkle tree is used as value.
+    pub fn insert_merkle_tree(&mut self, key: &[u8], tree: &MerkleTree) -> Result<()> {
+        self.0.insert(SLED_MERKLE_TREES_TREE, key, &serialize(tree))?;
+        Ok(())
+    }
+}
+
+pub type CacheSmt = SparseMerkleTree<
+    'static,
+    SMT_FP_DEPTH,
+    { SMT_FP_DEPTH + 1 },
+    pallas::Base,
+    PoseidonFp,
+    CacheOverlay,
+>;
+
+pub struct CacheSmtStorage<'a> {
+    overlay: &'a mut CacheOverlay,
+    tree: &'a [u8],
+}
+
+impl<'a> CacheSmtStorage<'a> {
+    pub fn new(overlay: &'a mut CacheOverlay, tree: &'a [u8]) -> Self {
+        Self { overlay, tree }
+    }
+}
+
+impl StorageAdapter for CacheSmtStorage<'_> {
+    type Value = pallas::Base;
+
+    fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
+        if let Err(e) = self.overlay.0.insert(self.tree, &key.to_bytes_le(), &value.to_repr()) {
+            error!(target: "cache::StorageAdapter::put", "Inserting key {key:?}, value {value:?} into DB failed: {e:?}");
+            return Err(ContractError::SmtPutFailed)
+        }
+        Ok(())
+    }
+
+    fn get(&self, key: &BigUint) -> Option<pallas::Base> {
+        let value = match self.overlay.0.get(self.tree, &key.to_bytes_le()) {
+            Ok(v) => v,
+            Err(e) => {
+                error!(target: "cache::StorageAdapter::get", "Fetching key {key:?} from DB failed: {e:?}");
+                return None
+            }
+        };
+
+        let value = value?;
+
+        let mut repr = [0; 32];
+        repr.copy_from_slice(&value);
+
+        pallas::Base::from_repr(repr).into()
+    }
+
+    fn del(&mut self, key: &BigUint) -> ContractResult {
+        if let Err(e) = self.overlay.0.remove(self.tree, &key.to_bytes_le()) {
+            error!(target: "cache::StorageAdapter::del", "Removing key {key:?} from DB failed: {e:?}");
+            return Err(ContractError::SmtDelFailed)
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use darkfi::{zk::halo2::Field, Result};
+    use darkfi_sdk::{
+        crypto::smt::{gen_empty_nodes, util::FieldHasher, PoseidonFp, SparseMerkleTree},
+        pasta::pallas,
+    };
+    use rand::rngs::OsRng;
+    use sled_overlay::sled;
+
+    use crate::cache::{Cache, CacheOverlay, CacheSmtStorage, SLED_MONEY_SMT_TREE};
+
+    #[test]
+    fn test_cache_smt() -> Result<()> {
+        // Setup cache and its overlay
+        let sled_db = sled::Config::new().temporary(true).open()?;
+        let cache = Cache::new(&sled_db)?;
+        let mut overlay = CacheOverlay::new(&cache)?;
+
+        // Setup SMT
+        const HEIGHT: usize = 3;
+        let hasher = PoseidonFp::new();
+        let empty_leaf = pallas::Base::ZERO;
+        let empty_nodes = gen_empty_nodes::<{ HEIGHT + 1 }, _, _>(&hasher, empty_leaf);
+        let store = CacheSmtStorage::new(&mut overlay, SLED_MONEY_SMT_TREE);
+        let mut smt = SparseMerkleTree::<HEIGHT, { HEIGHT + 1 }, _, _, _>::new(
+            store,
+            hasher.clone(),
+            &empty_nodes,
+        );
+
+        // Verify database is empty
+        assert!(cache.money_smt.is_empty());
+
+        let leaves = vec![
+            (pallas::Base::from(1), pallas::Base::random(&mut OsRng)),
+            (pallas::Base::from(2), pallas::Base::random(&mut OsRng)),
+            (pallas::Base::from(3), pallas::Base::random(&mut OsRng)),
+        ];
+        smt.insert_batch(leaves.clone()).unwrap();
+
+        let hash1 = leaves[0].1;
+        let hash2 = leaves[1].1;
+        let hash3 = leaves[2].1;
+
+        let hash = |l, r| hasher.hash([l, r]);
+
+        let hash01 = hash(empty_nodes[3], hash1);
+        let hash23 = hash(hash2, hash3);
+
+        let hash0123 = hash(hash01, hash23);
+        let root = hash(hash0123, empty_nodes[1]);
+        assert_eq!(root, smt.root());
+
+        // Now try to construct a membership proof for leaf 3
+        let pos = leaves[2].0;
+        let path = smt.prove_membership(&pos);
+        assert_eq!(path.path[0], empty_nodes[1]);
+        assert_eq!(path.path[1], hash01);
+        assert_eq!(path.path[2], hash2);
+
+        assert_eq!(hash23, hash(path.path[2], hash3));
+        assert_eq!(hash0123, hash(path.path[1], hash(path.path[2], hash3)));
+        assert_eq!(root, hash(hash(path.path[1], hash(path.path[2], hash3)), path.path[0]));
+
+        assert!(path.verify(&root, &hash3, &pos));
+
+        // Grab the overlay diff
+        let diff = overlay.0.diff(&[])?;
+
+        // Apply the overlay
+        overlay.0.apply_diff(&diff)?;
+
+        // Verify database contains keys
+        assert!(!cache.money_smt.is_empty());
+
+        // We are now going to rollback the changes
+        overlay.0.apply_diff(&diff.inverse())?;
+
+        // Verify database is empty again
+        assert!(cache.money_smt.is_empty());
+
+        Ok(())
     }
 }