From 9147e8653d538e484417d3b5045c263b7e3622c5 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 17 Nov 2022 16:33:58 +0200 Subject: [PATCH] cleanup dup code --- crates/stages/src/stages/bodies.rs | 53 ++++-------------------- crates/stages/src/stages/headers.rs | 40 ++---------------- crates/stages/src/test_utils/stage_db.rs | 32 ++++++++++++-- 3 files changed, 39 insertions(+), 86 deletions(-) diff --git a/crates/stages/src/stages/bodies.rs b/crates/stages/src/stages/bodies.rs index a43c702678..b62a872a24 100644 --- a/crates/stages/src/stages/bodies.rs +++ b/crates/stages/src/stages/bodies.rs @@ -475,10 +475,7 @@ mod tests { use assert_matches::assert_matches; use reth_eth_wire::BlockBody; use reth_interfaces::{ - db::{ - models::{BlockNumHash, StoredBlockBody}, - tables, DbCursorRO, DbTx, DbTxMut, - }, + db::{models::StoredBlockBody, tables, DbCursorRO, DbTx, DbTxMut}, p2p::bodies::{ client::BodiesClient, downloader::{BodiesStream, BodyDownloader}, @@ -486,10 +483,8 @@ mod tests { }, test_utils::{generators::random_block_range, TestConsensus}, }; - use reth_primitives::{ - BigEndianHash, BlockLocked, BlockNumber, Header, SealedHeader, H256, U256, - }; - use std::{collections::HashMap, ops::Deref, sync::Arc, time::Duration}; + use reth_primitives::{BlockLocked, BlockNumber, Header, SealedHeader, H256}; + use std::{collections::HashMap, sync::Arc, time::Duration}; /// The block hash of the genesis block. pub(crate) const GENESIS_HASH: H256 = H256::zero(); @@ -564,7 +559,7 @@ mod tests { let end = input.previous_stage.as_ref().map(|(_, num)| *num).unwrap_or_default(); let blocks = random_block_range(start..end, GENESIS_HASH); self.insert_genesis()?; - self.insert_headers(blocks.iter().map(|block| &block.header))?; + self.db.insert_headers(blocks.iter().map(|block| &block.header))?; self.set_responses(blocks.iter().map(body_by_hash).collect()); Ok(blocks) } @@ -602,7 +597,8 @@ mod tests { /// The genesis block always has no transactions and no ommers, and it always has the /// same hash. pub(crate) fn insert_genesis(&self) -> Result<(), TestRunnerError> { - self.insert_header(&SealedHeader::new(Header::default(), GENESIS_HASH))?; + let header = SealedHeader::new(Header::default(), GENESIS_HASH); + self.db.insert_headers(std::iter::once(&header))?; self.db.commit(|tx| { tx.put::( (0, GENESIS_HASH).into(), @@ -613,42 +609,7 @@ mod tests { Ok(()) } - /// Insert header into tables - pub(crate) fn insert_header( - &self, - header: &SealedHeader, - ) -> Result<(), TestRunnerError> { - self.insert_headers(std::iter::once(header))?; - Ok(()) - } - - /// Insert headers into tables - /// TODO: move to common inserter - pub(crate) fn insert_headers<'a, I>(&self, headers: I) -> Result<(), TestRunnerError> - where - I: Iterator, - { - let headers = headers.collect::>(); - self.db - .map_put::(&headers, |h| (h.hash(), h.number))?; - self.db.map_put::(&headers, |h| { - (BlockNumHash((h.number, h.hash())), h.deref().clone().unseal()) - })?; - self.db.map_put::(&headers, |h| { - (h.number, h.hash()) - })?; - - self.db.transform_append::(&headers, |prev, h| { - let prev_td = U256::from_big_endian(&prev.clone().unwrap_or_default()); - ( - BlockNumHash((h.number, h.hash())), - H256::from_uint(&(prev_td + h.difficulty)).as_bytes().to_vec(), - ) - })?; - - Ok(()) - } - + /// Retrieve the last body from the database pub(crate) fn last_body(&self) -> Option { self.db .query(|tx| Ok(tx.cursor::()?.last()?.map(|e| e.1))) diff --git a/crates/stages/src/stages/headers.rs b/crates/stages/src/stages/headers.rs index 49039a4c0c..ce6a2320b6 100644 --- a/crates/stages/src/stages/headers.rs +++ b/crates/stages/src/stages/headers.rs @@ -281,8 +281,8 @@ mod tests { TestConsensus, TestHeaderDownloader, TestHeadersClient, }, }; - use reth_primitives::{rpc::BigEndianHash, BlockNumber, SealedHeader, H256, U256}; - use std::{ops::Deref, sync::Arc}; + use reth_primitives::{BlockNumber, SealedHeader, H256, U256}; + use std::sync::Arc; pub(crate) struct HeadersTestRunner { pub(crate) consensus: Arc, @@ -327,7 +327,7 @@ mod tests { fn seed_execution(&mut self, input: ExecInput) -> Result { let start = input.stage_progress.unwrap_or_default(); let head = random_header(start, None); - self.insert_header(&head)?; + self.db.insert_headers(std::iter::once(&head))?; // use previous progress as seed size let end = input.previous_stage.map(|(_, num)| num).unwrap_or_default() + 1; @@ -425,40 +425,6 @@ mod tests { } impl HeadersTestRunner { - /// Insert header into tables - pub(crate) fn insert_header( - &self, - header: &SealedHeader, - ) -> Result<(), TestRunnerError> { - self.insert_headers(std::iter::once(header)) - } - - /// Insert headers into tables - pub(crate) fn insert_headers<'a, I>(&self, headers: I) -> Result<(), TestRunnerError> - where - I: Iterator, - { - let headers = headers.collect::>(); - self.db - .map_put::(&headers, |h| (h.hash(), h.number))?; - self.db.map_put::(&headers, |h| { - (BlockNumHash((h.number, h.hash())), h.deref().clone().unseal()) - })?; - self.db.map_put::(&headers, |h| { - (h.number, h.hash()) - })?; - - self.db.transform_append::(&headers, |prev, h| { - let prev_td = U256::from_big_endian(&prev.clone().unwrap_or_default()); - ( - BlockNumHash((h.number, h.hash())), - H256::from_uint(&(prev_td + h.difficulty)).as_bytes().to_vec(), - ) - })?; - - Ok(()) - } - pub(crate) fn check_no_header_entry_above( &self, block: BlockNumber, diff --git a/crates/stages/src/test_utils/stage_db.rs b/crates/stages/src/test_utils/stage_db.rs index f14b1cb4f2..9dd0c8dc97 100644 --- a/crates/stages/src/test_utils/stage_db.rs +++ b/crates/stages/src/test_utils/stage_db.rs @@ -2,9 +2,11 @@ use reth_db::{ kv::{test_utils::create_test_db, tx::Tx, Env, EnvKind}, mdbx::{WriteMap, RW}, }; -use reth_interfaces::db::{self, DBContainer, DbCursorRO, DbCursorRW, DbTx, DbTxMut, Table}; -use reth_primitives::BlockNumber; -use std::{borrow::Borrow, sync::Arc}; +use reth_interfaces::db::{ + self, models::BlockNumHash, tables, DBContainer, DbCursorRO, DbCursorRW, DbTx, DbTxMut, Table, +}; +use reth_primitives::{BigEndianHash, BlockNumber, SealedHeader, H256, U256}; +use std::{borrow::Borrow, ops::Deref, sync::Arc}; /// The [StageTestDB] is used as an internal /// database for testing stage implementation. @@ -146,4 +148,28 @@ impl StageTestDB { Ok(()) }) } + + /// Insert ordered collection of [SealedHeader] into the corresponding tables + /// that are supposed to be populated by the headers stage. + pub(crate) fn insert_headers<'a, I>(&self, headers: I) -> Result<(), db::Error> + where + I: Iterator, + { + let headers = headers.collect::>(); + self.map_put::(&headers, |h| (h.hash(), h.number))?; + self.map_put::(&headers, |h| { + (BlockNumHash((h.number, h.hash())), h.deref().clone().unseal()) + })?; + self.map_put::(&headers, |h| (h.number, h.hash()))?; + + self.transform_append::(&headers, |prev, h| { + let prev_td = U256::from_big_endian(&prev.clone().unwrap_or_default()); + ( + BlockNumHash((h.number, h.hash())), + H256::from_uint(&(prev_td + h.difficulty)).as_bytes().to_vec(), + ) + })?; + + Ok(()) + } }