diff --git a/src/db.rs b/src/db.rs index b7754f2..b16183c 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,7 +1,6 @@ use anyhow::{Context, Result}; use electrs_rocksdb as rocksdb; -use std::iter; use std::path::Path; use std::sync::atomic::{AtomicBool, Ordering}; @@ -240,48 +239,13 @@ impl DBStore { self.iter_prefix_cf(self.txid_cf(), prefix) } - fn iter_cf<'a, const N: usize, F: FnMut(&[u8]) -> bool + 'a>( + fn iter_cf<'a, const N: usize>( &'a self, cf: &rocksdb::ColumnFamily, readopts: rocksdb::ReadOptions, - start: Option<&[u8]>, - mut filter_fn: F, + prefix: Option, ) -> impl Iterator + '_ { - let mut raw_iter = self.db.raw_iterator_cf_opt(cf, readopts); - let mut done = false; - - if let Some(key) = start { - raw_iter.seek(key); - } else { - raw_iter.seek_to_first(); - }; - - iter::from_fn(move || loop { - // based on ::next - break if done { - None - } else if let Some((key, _)) = raw_iter.item() { - let ret = if filter_fn(key) { - Some( - key.try_into() - .with_context(|| format!("expected {N} bytes, got {}", key.len())) - .expect("database key has wrong length"), - ) - } else { - None - }; - raw_iter.next(); - if ret.is_some() { - ret - } else { - continue; - } - } else { - done = true; - raw_iter.status().expect("cf iterator failed"); - None - }; - }) + DBIterator::new(self.db.raw_iterator_cf_opt(cf, readopts), prefix) } fn iter_prefix_cf( @@ -291,18 +255,13 @@ impl DBStore { ) -> impl Iterator + '_ { let mut opts = rocksdb::ReadOptions::default(); opts.set_prefix_same_as_start(true); // requires .set_prefix_extractor() above. - self.iter_cf(cf, opts, Some(&prefix), |_| true) + self.iter_cf(cf, opts, Some(prefix)) } pub(crate) fn iter_headers(&self) -> impl Iterator + '_ { let mut opts = rocksdb::ReadOptions::default(); opts.fill_cache(false); - self.iter_cf( - self.headers_cf(), - opts, - None, - |key| key != TIP_KEY, // headers' rows are longer than TIP_KEY - ) + self.iter_cf(self.headers_cf(), opts, None) } pub(crate) fn get_tip(&self) -> Option> { @@ -406,6 +365,55 @@ impl DBStore { } } +struct DBIterator<'a, const N: usize> { + raw: rocksdb::DBRawIterator<'a>, + prefix: Option, + done: bool, +} + +impl<'a, const N: usize> DBIterator<'a, N> { + fn new(mut raw: rocksdb::DBRawIterator<'a>, prefix: Option) -> Self { + match prefix { + Some(key) => raw.seek(key), + None => raw.seek_to_first(), + }; + Self { + raw, + prefix, + done: false, + } + } +} + +impl<'a, const N: usize> Iterator for DBIterator<'a, N> { + type Item = [u8; N]; + + fn next(&mut self) -> Option { + while !self.done { + let key = match self.raw.key() { + Some(key) => key, + None => break, // end of scan + }; + let prefix_match = match self.prefix { + Some(key_prefix) => key.starts_with(&key_prefix), + None => true, + }; + if !prefix_match { + break; // prefix mismatch + } + let result: Option<[u8; N]> = key.try_into().ok(); + self.raw.next(); + match result { + Some(value) => return Some(value), + None => continue, // skip keys with size != N + } + } + self.raw.status().expect("DB scan failed"); + self.done = true; + None + } +} + impl Drop for DBStore { fn drop(&mut self) { info!("closing DB at {}", self.db.path().display());