diff --git a/ethexe/db/src/rocks.rs b/ethexe/db/src/rocks.rs index 7d2f48a8b89..e61523f35e8 100644 --- a/ethexe/db/src/rocks.rs +++ b/ethexe/db/src/rocks.rs @@ -19,7 +19,7 @@ use crate::{CASDatabase, KVDatabase}; use anyhow::Result; use gprimitives::H256; -use rocksdb::{Options, DB}; +use rocksdb::{DBIteratorWithThreadMode, Options, DB}; use std::{path::PathBuf, sync::Arc}; /// Database for storing states and codes in memory. @@ -85,13 +85,40 @@ impl KVDatabase for RocksDatabase { .expect("Failed to write data, database is not in valid state"); } - fn iter_prefix(&self, prefix: &[u8]) -> Box, Vec)> + '_> { - Box::new( - self.inner - .prefix_iterator(prefix) - .map(|kv| kv.expect("unexpected error during iteration")) - .map(|(k, v)| (<[u8]>::into_vec(k), <[u8]>::into_vec(v))), - ) + fn iter_prefix<'a>( + &'a self, + prefix: &'a [u8], + ) -> Box, Vec)> + '_> { + Box::new(PrefixIterator { + prefix, + prefix_iter: self.inner.prefix_iterator(prefix), + done: false, + }) + } +} + +pub struct PrefixIterator<'a> { + prefix: &'a [u8], + prefix_iter: DBIteratorWithThreadMode<'a, DB>, + done: bool, +} + +impl<'a> Iterator for PrefixIterator<'a> { + type Item = (Vec, Vec); + + fn next(&mut self) -> Option { + if self.done { + return None; + } + + match self.prefix_iter.next() { + Some(Ok((k, v))) if k.starts_with(self.prefix) => Some((k.to_vec(), v.to_vec())), + Some(Err(_)) => panic!("Failed to read data, database is not in valid state"), + _ => { + self.done = true; + None + } + } } }