state: create SledSerialize, SledDeserialize traits

zebra-state/src/sled_state.rs

@@ -3,17 +3,18 @@
 use std::{collections::HashMap, convert::TryInto, sync::Arc};
 
 use tracing::trace;
+use zebra_chain::transparent;
 use zebra_chain::{
     block::{self, Block},
     parameters::{Network, GENESIS_PREVIOUS_BLOCK_HASH},
 };
-use zebra_chain::{
-    serialization::{ZcashDeserialize, ZcashSerialize},
-    transparent,
-};
 
 use crate::{BoxError, Config, HashOrHeight, QueuedBlock};
 
+mod sled_format;
+
+use sled_format::{SledDeserialize, SledSerialize};
+
 /// The finalized part of the chain state, stored in sled.
 ///
 /// This structure has two categories of methods:
@@ -38,7 +39,7 @@ pub struct FinalizedState {
     hash_by_height: sled::Tree,
     height_by_hash: sled::Tree,
     block_by_height: sled::Tree,
-    // tx_by_hash: sled::Tree,
+    tx_by_hash: sled::Tree,
     utxo_by_outpoint: sled::Tree,
     // sprout_nullifiers: sled::Tree,
     // sapling_nullifiers: sled::Tree,
@@ -48,76 +49,6 @@ pub struct FinalizedState {
     debug_stop_at_height: Option<block::Height>,
 }
 
-/// Helper trait for inserting (Key, Value) pairs into sled when both the key and
-/// value implement ZcashSerialize.
-trait SledSerialize {
-    /// Serialize and insert the given key and value into a sled tree.
-    fn zs_insert<K, V>(
-        &self,
-        key: &K,
-        value: &V,
-    ) -> Result<(), sled::transaction::UnabortableTransactionError>
-    where
-        K: ZcashSerialize,
-        V: ZcashSerialize;
-}
-
-/// Helper trait for retrieving values from sled trees when the key and value
-/// implement ZcashSerialize/ZcashDeserialize.
-trait SledDeserialize {
-    /// Serialize the given key and use that to get and deserialize the
-    /// corresponding value from a sled tree, if it is present.
-    fn zs_get<K, V>(&self, key: &K) -> Result<Option<V>, BoxError>
-    where
-        K: ZcashSerialize,
-        V: ZcashDeserialize;
-}
-
-impl SledSerialize for sled::transaction::TransactionalTree {
-    fn zs_insert<K, V>(
-        &self,
-        key: &K,
-        value: &V,
-    ) -> Result<(), sled::transaction::UnabortableTransactionError>
-    where
-        K: ZcashSerialize,
-        V: ZcashSerialize,
-    {
-        let key_bytes = key
-            .zcash_serialize_to_vec()
-            .expect("serializing into a vec won't fail");
-
-        let value_bytes = value
-            .zcash_serialize_to_vec()
-            .expect("serializing into a vec won't fail");
-
-        self.insert(key_bytes, value_bytes)?;
-
-        Ok(())
-    }
-}
-
-impl SledDeserialize for sled::Tree {
-    fn zs_get<K, V>(&self, key: &K) -> Result<Option<V>, BoxError>
-    where
-        K: ZcashSerialize,
-        V: ZcashDeserialize,
-    {
-        let key_bytes = key
-            .zcash_serialize_to_vec()
-            .expect("serializing into a vec won't fail");
-
-        let value_bytes = self.get(&key_bytes)?;
-
-        let value = value_bytes
-            .as_deref()
-            .map(ZcashDeserialize::zcash_deserialize)
-            .transpose()?;
-
-        Ok(value)
-    }
-}
-
 /// Where is the stop check being performed?
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 enum StopCheckContext {
@@ -136,7 +67,7 @@ impl FinalizedState {
             hash_by_height: db.open_tree(b"hash_by_height").unwrap(),
             height_by_hash: db.open_tree(b"height_by_hash").unwrap(),
             block_by_height: db.open_tree(b"block_by_height").unwrap(),
-            // tx_by_hash: db.open_tree(b"tx_by_hash").unwrap(),
+            tx_by_hash: db.open_tree(b"tx_by_hash").unwrap(),
             utxo_by_outpoint: db.open_tree(b"utxo_by_outpoint").unwrap(),
             // sprout_nullifiers: db.open_tree(b"sprout_nullifiers").unwrap(),
             // sapling_nullifiers: db.open_tree(b"sapling_nullifiers").unwrap(),
@@ -164,7 +95,7 @@ impl FinalizedState {
         total_flushed += self.hash_by_height.flush()?;
         total_flushed += self.height_by_hash.flush()?;
         total_flushed += self.block_by_height.flush()?;
-        // total_flushed += self.tx_by_hash.flush()?;
+        total_flushed += self.tx_by_hash.flush()?;
         total_flushed += self.utxo_by_outpoint.flush()?;
         // total_flushed += self.sprout_nullifiers.flush()?;
         // total_flushed += self.sapling_nullifiers.flush()?;
@@ -280,7 +211,6 @@ impl FinalizedState {
         let height = block
             .coinbase_height()
             .expect("finalized blocks are valid and have a coinbase height");
-        let height_bytes = height.0.to_be_bytes();
         let hash = block.hash();
 
         trace!(?height, "Finalized block");
@@ -290,31 +220,33 @@ impl FinalizedState {
             &self.height_by_hash,
             &self.block_by_height,
             &self.utxo_by_outpoint,
+            &self.tx_by_hash,
         )
             .transaction(
-                move |(hash_by_height, height_by_hash, block_by_height, utxo_by_outpoint)| {
+                move |(
-                    // TODO: do serialization above
+                    hash_by_height,
-                    // for some reason this wouldn't move into the closure (??)
+                    height_by_hash,
-                    let block_bytes = block
+                    block_by_height,
-                        .zcash_serialize_to_vec()
+                    utxo_by_outpoint,
-                        .expect("zcash_serialize_to_vec has wrong return type");
+                    tx_by_hash,
-
+                )| {
                     // TODO: check highest entry of hash_by_height as in RFC
 
-                    hash_by_height.insert(&height_bytes, &hash.0)?;
+                    hash_by_height.zs_insert(height, hash)?;
-                    height_by_hash.insert(&hash.0, &height_bytes)?;
+                    height_by_hash.zs_insert(hash, height)?;
-                    block_by_height.insert(&height_bytes, block_bytes)?;
+                    block_by_height.zs_insert(height, &*block)?;
-                    // tx_by_hash
 
                     for transaction in block.transactions.iter() {
                         let transaction_hash = transaction.hash();
+                        tx_by_hash.zs_insert(transaction_hash, transaction)?;
+
                         for (index, output) in transaction.outputs().iter().enumerate() {
                             let outpoint = transparent::OutPoint {
                                 hash: transaction_hash,
                                 index: index as _,
                             };
 
-                            utxo_by_outpoint.zs_insert(&outpoint, output)?;
+                            utxo_by_outpoint.zs_insert(outpoint, output)?;
                         }
                     }
                     // sprout_nullifiers
@@ -354,11 +286,11 @@ impl FinalizedState {
         let heights = crate::util::block_locator_heights(tip_height);
         let mut hashes = Vec::with_capacity(heights.len());
         for height in heights {
-            if let Some(bytes) = self.hash_by_height.get(&height.0.to_be_bytes())? {
+            if let Some(hash) = self.hash_by_height.zs_get(&height)? {
-                let hash = block::Hash(bytes.as_ref().try_into().unwrap());
                 hashes.push(hash)
             }
         }
+
         Ok(hashes)
     }
 
@@ -375,8 +307,8 @@ impl FinalizedState {
     }
 
     pub fn depth(&self, hash: block::Hash) -> Result<Option<u32>, BoxError> {
-        let height = match self.height_by_hash.get(&hash.0)? {
+        let height: block::Height = match self.height_by_hash.zs_get(&hash)? {
-            Some(bytes) => block::Height(u32::from_be_bytes(bytes.as_ref().try_into().unwrap())),
+            Some(height) => height,
             None => return Ok(None),
         };
 
@@ -388,16 +320,14 @@ impl FinalizedState {
     pub fn block(&self, hash_or_height: HashOrHeight) -> Result<Option<Arc<Block>>, BoxError> {
         let height = match hash_or_height {
             HashOrHeight::Height(height) => height,
-            HashOrHeight::Hash(hash) => match self.height_by_hash.get(&hash.0)? {
+            HashOrHeight::Hash(hash) => match self.height_by_hash.zs_get(&hash)? {
-                Some(bytes) => {
+                Some(height) => height,
-                    block::Height(u32::from_be_bytes(bytes.as_ref().try_into().unwrap()))
-                }
                 None => return Ok(None),
             },
         };
 
-        match self.block_by_height.get(&height.0.to_be_bytes())? {
+        match self.block_by_height.zs_get(&height)? {
-            Some(bytes) => Ok(Some(Arc::<Block>::zcash_deserialize(bytes.as_ref())?)),
+            Some(block) => Ok(Some(block)),
             None => Ok(None),
         }
     }
@@ -414,8 +344,7 @@ impl FinalizedState {
     /// Returns the finalized hash for a given `block::Height` if it is present.
     pub fn get_hash(&self, height: block::Height) -> Option<block::Hash> {
         self.hash_by_height
-            .get(&height.0.to_be_bytes())
+            .zs_get(&height)
             .expect("expected that sled errors would not occur")
-            .map(|bytes| block::Hash(bytes.as_ref().try_into().unwrap()))
     }
 }

zebra-state/src/sled_state/sled_format.rs

@@ -0,0 +1,205 @@
+//! Module defining exactly how to move types in and out of sled
+use std::{convert::TryInto, sync::Arc};
+
+use zebra_chain::{
+    block,
+    block::Block,
+    serialization::{ZcashDeserialize, ZcashSerialize},
+    transaction,
+    transaction::Transaction,
+    transparent,
+};
+
+use crate::BoxError;
+
+// Helper trait for defining the exact format used to interact with sled per
+// type.
+pub trait IntoSled {
+    // The type used to compare a value as a key to other keys stored in a
+    // sled::Tree
+    type Bytes: AsRef<[u8]>;
+
+    // function to convert the current type to its sled format in `zs_get()`
+    // without necessarily allocating a new IVec
+    fn as_bytes(&self) -> Self::Bytes;
+
+    // function to convert the current type into its sled format
+    fn into_ivec(self) -> sled::IVec;
+}
+
+// Helper type for retrieving types from sled with the correct format
+pub trait FromSled: Sized {
+    // function to convert the sled bytes back into the deserialized type
+    fn from_ivec(bytes: sled::IVec) -> Result<Self, BoxError>;
+}
+
+impl IntoSled for &Block {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.zcash_serialize_to_vec()
+            .expect("serialization to vec doesn't fail")
+    }
+
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().into()
+    }
+}
+
+impl FromSled for Arc<Block> {
+    fn from_ivec(bytes: sled::IVec) -> Result<Self, BoxError> {
+        let block = Arc::<Block>::zcash_deserialize(bytes.as_ref())?;
+        Ok(block)
+    }
+}
+
+impl IntoSled for &Arc<Transaction> {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.zcash_serialize_to_vec()
+            .expect("serialization to vec doesn't fail")
+    }
+
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().into()
+    }
+}
+
+impl IntoSled for transaction::Hash {
+    type Bytes = [u8; 32];
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.0
+    }
+
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().as_ref().into()
+    }
+}
+
+impl IntoSled for block::Hash {
+    type Bytes = [u8; 32];
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.0
+    }
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().as_ref().into()
+    }
+}
+
+impl FromSled for block::Hash {
+    fn from_ivec(bytes: sled::IVec) -> Result<Self, BoxError> {
+        let array = bytes.as_ref().try_into().unwrap();
+        Ok(Self(array))
+    }
+}
+
+impl IntoSled for block::Height {
+    type Bytes = [u8; 4];
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.0.to_be_bytes()
+    }
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().as_ref().into()
+    }
+}
+
+impl FromSled for block::Height {
+    fn from_ivec(bytes: sled::IVec) -> Result<Self, BoxError> {
+        let array = bytes.as_ref().try_into().unwrap();
+        Ok(block::Height(u32::from_be_bytes(array)))
+    }
+}
+
+impl IntoSled for &transparent::Output {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.zcash_serialize_to_vec()
+            .expect("serialization to vec doesn't fail")
+    }
+
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().into()
+    }
+}
+
+impl FromSled for transparent::Output {
+    fn from_ivec(bytes: sled::IVec) -> Result<Self, BoxError> {
+        Self::zcash_deserialize(&*bytes).map_err(Into::into)
+    }
+}
+
+impl IntoSled for transparent::OutPoint {
+    type Bytes = Vec<u8>;
+
+    fn as_bytes(&self) -> Self::Bytes {
+        self.zcash_serialize_to_vec()
+            .expect("serialization to vec doesn't fail")
+    }
+
+    fn into_ivec(self) -> sled::IVec {
+        self.as_bytes().into()
+    }
+}
+
+/// Helper trait for inserting (Key, Value) pairs into sled with a consistently
+/// defined format
+pub trait SledSerialize {
+    /// Serialize and insert the given key and value into a sled tree.
+    fn zs_insert<K, V>(
+        &self,
+        key: K,
+        value: V,
+    ) -> Result<(), sled::transaction::UnabortableTransactionError>
+    where
+        K: IntoSled,
+        V: IntoSled;
+}
+
+impl SledSerialize for sled::transaction::TransactionalTree {
+    fn zs_insert<K, V>(
+        &self,
+        key: K,
+        value: V,
+    ) -> Result<(), sled::transaction::UnabortableTransactionError>
+    where
+        K: IntoSled,
+        V: IntoSled,
+    {
+        let key_bytes = key.into_ivec();
+        let value_bytes = value.into_ivec();
+        self.insert(key_bytes, value_bytes)?;
+        Ok(())
+    }
+}
+
+/// Helper trait for retrieving values from sled trees with a consistently
+/// defined format
+pub trait SledDeserialize {
+    /// Serialize the given key and use that to get and deserialize the
+    /// corresponding value from a sled tree, if it is present.
+    fn zs_get<K, V>(&self, key: &K) -> Result<Option<V>, BoxError>
+    where
+        K: IntoSled,
+        V: FromSled;
+}
+
+impl SledDeserialize for sled::Tree {
+    fn zs_get<K, V>(&self, key: &K) -> Result<Option<V>, BoxError>
+    where
+        K: IntoSled,
+        V: FromSled,
+    {
+        let key_bytes = key.as_bytes();
+
+        let value_bytes = self.get(key_bytes)?;
+
+        let value = value_bytes.map(V::from_ivec).transpose()?;
+
+        Ok(value)
+    }
+}