feat(interfaces): add ShardedKey (#142)

* add ShardedKey

* Update crates/interfaces/src/db/models/sharded_key.rs

Co-authored-by: rakita <rakita@users.noreply.github.com>

Co-authored-by: rakita <rakita@users.noreply.github.com>

crates/db/src/kv/mod.rs

@@ -135,11 +135,12 @@ pub mod test_utils {
 mod tests {
     use super::{test_utils, Env, EnvKind};
     use reth_interfaces::db::{
-        tables::{Headers, PlainAccountState, PlainStorageState},
+        models::ShardedKey,
+        tables::{AccountHistory, Headers, PlainAccountState, PlainStorageState},
         Database, DbCursorRO, DbDupCursorRO, DbTx, DbTxMut,
     };
     use reth_libmdbx::{NoWriteMap, WriteMap};
-    use reth_primitives::{Account, Address, Header, StorageEntry, H256, U256};
+    use reth_primitives::{Account, Address, Header, IntegerList, StorageEntry, H256, U256};
     use std::str::FromStr;
     use tempfile::TempDir;
 
@@ -264,7 +265,7 @@ mod tests {
             assert!(Some(value22) == cursor.next_dup_val().unwrap());
         }
 
-        // Seek value with subkey
+        // Seek value with exact subkey
         {
             let tx = env.tx().expect(ERROR_INIT_TX);
             let mut cursor = tx.cursor_dup::<PlainStorageState>().unwrap();
@@ -278,6 +279,38 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn db_sharded_key() {
+        let db: Env<WriteMap> = test_utils::create_test_db(EnvKind::RW);
+        let real_key = Address::from_str("0xa2c122be93b0074270ebee7f6b7292c7deb45047").unwrap();
+
+        for i in 1..5 {
+            let key = ShardedKey::new(real_key, i * 100);
+            let list: IntegerList = vec![i * 100u64].into();
+
+            db.update(|tx| tx.put::<AccountHistory>(key.clone(), list.clone()).expect("")).unwrap();
+        }
+
+        // Seek value with non existing key.
+        {
+            let tx = db.tx().expect(ERROR_INIT_TX);
+            let mut cursor = tx.cursor::<AccountHistory>().unwrap();
+
+            // It will seek the one greater or equal to the query. Since we have `Address | 100`,
+            // `Address | 200` in the database and we're querying `Address | 150` it will return us
+            // `Address | 200`.
+            let mut walker = cursor.walk(ShardedKey::new(real_key, 150)).unwrap();
+            let (key, list) = walker
+                .next()
+                .expect("element should exist.")
+                .expect("should be able to retrieve it.");
+
+            assert_eq!(ShardedKey::new(real_key, 200), key);
+            let list200: IntegerList = vec![200u64].into();
+            assert_eq!(list200, list);
+        }
+    }
 }
 
 #[cfg(test)]

crates/interfaces/src/db/models/mod.rs

@@ -3,6 +3,8 @@
 pub mod accounts;
 pub mod blocks;
 pub mod integer_list;
+pub mod sharded_key;
 
 pub use accounts::*;
 pub use blocks::*;
+pub use sharded_key::ShardedKey;

crates/interfaces/src/db/models/sharded_key.rs

@@ -0,0 +1,62 @@
+//! Sharded key
+
+use crate::db::{
+    table::{Decode, Encode},
+    Error,
+};
+use eyre::eyre;
+use reth_primitives::TxNumber;
+
+/// Sometimes data can be too big to be saved for a single key. This helps out by dividing the data
+/// into different shards. Example:
+///
+/// `Address | 200` -> data is from transaction 0 to 200.
+///
+/// `Address | 300` -> data is from transaction 201 to 300.
+#[derive(Debug, Default, Clone, PartialEq)]
+pub struct ShardedKey<T> {
+    /// The key for this type.
+    pub key: T,
+    /// Highest tx number to which `value` is related to.
+    pub highest_tx_number: TxNumber,
+}
+
+impl<T> ShardedKey<T> {
+    /// Creates a new `ShardedKey<T>`.
+    pub fn new(key: T, highest_tx_number: TxNumber) -> Self {
+        ShardedKey { key, highest_tx_number }
+    }
+}
+
+impl<T> Encode for ShardedKey<T>
+where
+    T: Encode,
+    Vec<u8>: From<<T as Encode>::Encoded>,
+{
+    type Encoded = Vec<u8>;
+
+    fn encode(self) -> Self::Encoded {
+        let mut buf: Vec<u8> = Encode::encode(self.key).into();
+        buf.extend_from_slice(&self.highest_tx_number.to_be_bytes());
+        buf
+    }
+}
+
+impl<T> Decode for ShardedKey<T>
+where
+    T: Decode,
+{
+    fn decode<B: Into<bytes::Bytes>>(value: B) -> Result<Self, Error> {
+        let value: bytes::Bytes = value.into();
+        let tx_num_index = value.len() - 8;
+
+        let highest_tx_number = u64::from_be_bytes(
+            value.as_ref()[tx_num_index..]
+                .try_into()
+                .map_err(|_| Error::Decode(eyre!("Into bytes error.")))?,
+        );
+        let key = T::decode(value.slice(..tx_num_index))?;
+
+        Ok(ShardedKey::new(key, highest_tx_number))
+    }
+}

crates/interfaces/src/db/tables.rs

@@ -4,6 +4,7 @@ use crate::db::{
     models::{
         accounts::{AccountBeforeTx, TxNumberAddress},
         blocks::{BlockNumHash, HeaderHash, NumTransactions, NumTxesInBlock},
+        ShardedKey,
     },
     DupSort,
 };
@@ -105,7 +106,7 @@ table!(Logs => TxNumber => Receipt); // Canonical only
 table!(PlainAccountState => Address => Account);
 dupsort!(PlainStorageState => Address => [H256] StorageEntry);
 
-table!(AccountHistory => Address => TxNumberList);
+table!(AccountHistory => ShardedKey<Address> => TxNumberList);
 table!(StorageHistory => Address_StorageKey => TxNumberList);
 
 dupsort!(AccountChangeSet => TxNumber => [Address] AccountBeforeTx);