diff --git a/Cargo.lock b/Cargo.lock index 940ae824..8cec655a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2397,6 +2397,7 @@ dependencies = [ "kona-preimage", "kona-primitives", "os_pipe", + "proptest", "reqwest", "revm", "rocksdb", diff --git a/bin/host/Cargo.toml b/bin/host/Cargo.toml index fe4bf6a4..5dae44e1 100644 --- a/bin/host/Cargo.toml +++ b/bin/host/Cargo.toml @@ -42,3 +42,6 @@ tracing-subscriber.workspace = true command-fds.workspace = true os_pipe.workspace = true rocksdb.workspace = true + +[dev-dependencies] +proptest.workspace = true diff --git a/bin/host/src/kv/disk.rs b/bin/host/src/kv/disk.rs index ee76e0ab..48aa03b1 100644 --- a/bin/host/src/kv/disk.rs +++ b/bin/host/src/kv/disk.rs @@ -1,7 +1,8 @@ //! Contains a concrete implementation of the [KeyValueStore] trait that stores data on disk //! using [rocksdb]. -use super::KeyValueStore; +use super::{KeyValueStore, MemoryKeyValueStore}; +use alloy_primitives::B256; use anyhow::{anyhow, Result}; use rocksdb::{Options, DB}; use std::path::PathBuf; @@ -46,3 +47,51 @@ impl Drop for DiskKeyValueStore { let _ = DB::destroy(&Self::get_db_options(), self.data_directory.as_path()); } } + +impl TryFrom for MemoryKeyValueStore { + type Error = anyhow::Error; + + fn try_from(disk_store: DiskKeyValueStore) -> Result { + let mut memory_store = MemoryKeyValueStore::new(); + let mut db_iter = disk_store.db.full_iterator(rocksdb::IteratorMode::Start); + + while let Some(Ok((key, value))) = db_iter.next() { + memory_store.set( + B256::try_from(key.as_ref()) + .map_err(|e| anyhow!("Failed to convert slice to B256: {e}"))?, + value.to_vec(), + )?; + } + + Ok(memory_store) + } +} + +#[cfg(test)] +mod test { + use super::DiskKeyValueStore; + use crate::kv::{KeyValueStore, MemoryKeyValueStore}; + use proptest::{ + arbitrary::any, + collection::{hash_map, vec}, + proptest, + }; + use std::env::temp_dir; + + proptest! { + /// Test that converting from a [DiskKeyValueStore] to a [MemoryKeyValueStore] is lossless. + #[test] + fn convert_disk_kv_to_mem_kv(k_v in hash_map(any::<[u8; 32]>(), vec(any::(), 0..128), 1..128)) { + let tempdir = temp_dir(); + let mut disk_kv = DiskKeyValueStore::new(tempdir); + k_v.iter().for_each(|(k, v)| { + disk_kv.set(k.into(), v.to_vec()).unwrap(); + }); + + let mem_kv = MemoryKeyValueStore::try_from(disk_kv).unwrap(); + for (k, v) in k_v { + assert_eq!(mem_kv.get(k.into()).unwrap(), v.to_vec()); + } + } + } +}