Merge remote-tracking branch 'origin/freezer-tools' into tree-states

beacon_node/store/src/lib.rs

@@ -48,7 +48,7 @@ use strum::{EnumString, IntoStaticStr};
 pub use types::*;
 pub use validator_pubkey_cache::ValidatorPubkeyCache;
 
-pub type ColumnIter<'a> = Box<dyn Iterator<Item = Result<(Hash256, Vec<u8>), Error>> + 'a>;
+pub type ColumnIter<'a, K> = Box<dyn Iterator<Item = Result<(K, Vec<u8>), Error>> + 'a>;
 pub type ColumnKeyIter<'a> = Box<dyn Iterator<Item = Result<Hash256, Error>> + 'a>;
 
 pub trait KeyValueStore<E: EthSpec>: Sync + Send + Sized + 'static {
@@ -85,7 +85,7 @@ pub trait KeyValueStore<E: EthSpec>: Sync + Send + Sized + 'static {
     fn compact(&self) -> Result<(), Error>;
 
     /// Iterate through all keys and values in a particular column.
-    fn iter_column(&self, _column: DBColumn) -> ColumnIter {
+    fn iter_column<K: Key>(&self, _column: DBColumn) -> ColumnIter<K> {
         // Default impl for non LevelDB databases
         Box::new(std::iter::empty())
     }
@@ -97,6 +97,26 @@ pub trait KeyValueStore<E: EthSpec>: Sync + Send + Sized + 'static {
     }
 }
 
+pub trait Key: Sized + 'static {
+    fn from_bytes(key: &[u8]) -> Result<Self, Error>;
+}
+
+impl Key for Hash256 {
+    fn from_bytes(key: &[u8]) -> Result<Self, Error> {
+        if key.len() == 32 {
+            Ok(Hash256::from_slice(key))
+        } else {
+            Err(Error::InvalidKey)
+        }
+    }
+}
+
+impl Key for Vec<u8> {
+    fn from_bytes(key: &[u8]) -> Result<Self, Error> {
+        Ok(key.to_vec())
+    }
+}
+
 pub fn get_key_for_col(column: &str, key: &[u8]) -> Vec<u8> {
     let mut result = column.as_bytes().to_vec();
     result.extend_from_slice(key);
@@ -246,6 +266,34 @@ impl DBColumn {
     pub fn as_bytes(self) -> &'static [u8] {
         self.as_str().as_bytes()
     }
+
+    /// Most database keys are 32 bytes, but some freezer DB keys are 8 bytes.
+    ///
+    /// This function returns the number of bytes used by keys in a given column.
+    pub fn key_size(self) -> usize {
+        match self {
+            Self::BeaconMeta
+            | Self::BeaconBlock
+            | Self::BeaconState
+            | Self::BeaconStateSummary
+            | Self::BeaconStateTemporary
+            | Self::ExecPayload
+            | Self::BeaconChain
+            | Self::OpPool
+            | Self::Eth1Cache
+            | Self::ForkChoice
+            | Self::PubkeyCache
+            | Self::BeaconRestorePoint
+            | Self::DhtEnrs
+            | Self::OptimisticTransitionBlock
+            | Self::BeaconStateDiff => 32,
+            Self::BeaconBlockRoots
+            | Self::BeaconStateRoots
+            | Self::BeaconHistoricalRoots
+            | Self::BeaconRandaoMixes
+            | Self::BeaconBlockFrozen => 8,
+        }
+    }
 }
 
 /// An item that may stored in a `Store` by serializing and deserializing from bytes.