Skip to content

Commit

Permalink
perf: improve IntegerList API to avoid allocations
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Sep 27, 2024
1 parent fbb0b11 commit 45f7a77
Show file tree
Hide file tree
Showing 13 changed files with 39 additions and 63 deletions.
1 change: 0 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

46 changes: 13 additions & 33 deletions crates/primitives-traits/src/integer_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,34 +16,30 @@ pub struct IntegerList(pub RoaringTreemap);

impl fmt::Debug for IntegerList {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let vec: Vec<u64> = self.0.iter().collect();
write!(f, "IntegerList {vec:?}")
f.write_str("IntegerList")?;
f.debug_list().entries(self.0.iter()).finish()
}
}

impl IntegerList {
/// Creates an `IntegerList` from a list of integers.
///
/// # Returns
///
/// Returns an error if the list is empty or not pre-sorted.
pub fn new<T: AsRef<[u64]>>(list: T) -> Result<Self, RoaringBitmapError> {
Ok(Self(
RoaringTreemap::from_sorted_iter(list.as_ref().iter().copied())
.map_err(|_| RoaringBitmapError::InvalidInput)?,
))
pub fn new(list: impl IntoIterator<Item = u64>) -> Result<Self, RoaringBitmapError> {
RoaringTreemap::from_sorted_iter(list)
.map(Self)
.map_err(|_| RoaringBitmapError::InvalidInput)
}

// Creates an IntegerList from a pre-sorted list of integers.
///
/// # Panics
///
/// Panics if the list is empty or not pre-sorted.
pub fn new_pre_sorted<T: AsRef<[u64]>>(list: T) -> Self {
Self(
RoaringTreemap::from_sorted_iter(list.as_ref().iter().copied())
.expect("IntegerList must be pre-sorted and non-empty"),
)
#[inline]
#[track_caller]
pub fn new_pre_sorted(list: impl IntoIterator<Item = u64>) -> Self {
Self::new(list).expect("IntegerList must be pre-sorted and non-empty")
}

/// Serializes a [`IntegerList`] into a sequence of bytes.
Expand All @@ -67,28 +63,13 @@ impl IntegerList {
}
}

macro_rules! impl_uint {
($($w:tt),+) => {
$(
impl From<Vec<$w>> for IntegerList {
fn from(v: Vec<$w>) -> Self {
Self::new_pre_sorted(v.iter().map(|v| *v as u64).collect::<Vec<_>>())
}
}
)+
};
}

impl_uint!(usize, u64, u32, u8, u16);

impl Serialize for IntegerList {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let vec = self.0.iter().collect::<Vec<u64>>();
let mut seq = serializer.serialize_seq(Some(self.len() as usize))?;
for e in vec {
for e in &self.0 {
seq.serialize_element(&e)?;
}
seq.end()
Expand All @@ -107,11 +88,10 @@ impl<'de> Visitor<'de> for IntegerListVisitor {
where
E: SeqAccess<'de>,
{
let mut list = Vec::new();
let mut list = Vec::with_capacity(seq.size_hint().unwrap_or(0).min(1024));
while let Some(item) = seq.next_element()? {
list.push(item);
}

IntegerList::new(list).map_err(|_| serde::de::Error::invalid_value(Unexpected::Seq, &self))
}
}
Expand All @@ -132,7 +112,7 @@ use arbitrary::{Arbitrary, Unstructured};
impl<'a> Arbitrary<'a> for IntegerList {
fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self, arbitrary::Error> {
let mut nums: Vec<u64> = Vec::arbitrary(u)?;
nums.sort();
nums.sort_unstable();
Self::new(nums).map_err(|_| arbitrary::Error::IncorrectFormat)
}
}
Expand Down
6 changes: 2 additions & 4 deletions crates/prune/prune/src/segments/user/account_history.rs
Original file line number Diff line number Diff line change
Expand Up @@ -275,10 +275,8 @@ mod tests {
.iter()
.filter(|(key, _)| key.highest_block_number > last_pruned_block_number)
.map(|(key, blocks)| {
let new_blocks = blocks
.iter()
.skip_while(|block| *block <= last_pruned_block_number)
.collect::<Vec<_>>();
let new_blocks =
blocks.iter().skip_while(|block| *block <= last_pruned_block_number);
(key.clone(), BlockNumberList::new_pre_sorted(new_blocks))
})
.collect::<Vec<_>>();
Expand Down
10 changes: 5 additions & 5 deletions crates/prune/prune/src/segments/user/history.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,18 +101,18 @@ where
// contain the target block number, as it's in this shard.
else {
let blocks = raw_blocks.value()?;
let higher_blocks =
blocks.iter().skip_while(|block| *block <= to_block).collect::<Vec<_>>();
let higher_blocks = || blocks.iter().skip_while(|block| *block <= to_block);
let n_higher_blocks = higher_blocks().count();

// If there were blocks less than or equal to the target one
// (so the shard has changed), update the shard.
if blocks.len() as usize == higher_blocks.len() {
if blocks.len() as usize == n_higher_blocks {
return Ok(PruneShardOutcome::Unchanged);
}

// If there will be no more blocks in the shard after pruning blocks below target
// block, we need to remove it, as empty shards are not allowed.
if higher_blocks.is_empty() {
if n_higher_blocks == 0 {
if key.as_ref().highest_block_number == u64::MAX {
let prev_row = cursor
.prev()?
Expand Down Expand Up @@ -151,7 +151,7 @@ where
} else {
cursor.upsert(
RawKey::new(key),
RawValue::new(BlockNumberList::new_pre_sorted(higher_blocks)),
RawValue::new(BlockNumberList::new_pre_sorted(higher_blocks())),
)?;
Ok(PruneShardOutcome::Updated)
}
Expand Down
6 changes: 2 additions & 4 deletions crates/prune/prune/src/segments/user/storage_history.rs
Original file line number Diff line number Diff line change
Expand Up @@ -281,10 +281,8 @@ mod tests {
.iter()
.filter(|(key, _)| key.sharded_key.highest_block_number > last_pruned_block_number)
.map(|(key, blocks)| {
let new_blocks = blocks
.iter()
.skip_while(|block| *block <= last_pruned_block_number)
.collect::<Vec<_>>();
let new_blocks =
blocks.iter().skip_while(|block| *block <= last_pruned_block_number);
(key.clone(), BlockNumberList::new_pre_sorted(new_blocks))
})
.collect::<Vec<_>>();
Expand Down
2 changes: 1 addition & 1 deletion crates/stages/stages/src/stages/index_account_history.rs
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ mod tests {
}

fn list(list: &[u64]) -> BlockNumberList {
BlockNumberList::new(list).unwrap()
BlockNumberList::new(list.iter().copied()).unwrap()
}

fn cast(
Expand Down
2 changes: 1 addition & 1 deletion crates/stages/stages/src/stages/index_storage_history.rs
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,7 @@ mod tests {
}

fn list(list: &[u64]) -> BlockNumberList {
BlockNumberList::new(list).unwrap()
BlockNumberList::new(list.iter().copied()).unwrap()
}

fn cast(
Expand Down
6 changes: 3 additions & 3 deletions crates/stages/stages/src/stages/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,11 @@ where
let mut cache: HashMap<P, Vec<u64>> = HashMap::default();

let mut collect = |cache: &HashMap<P, Vec<u64>>| {
for (key, indice_list) in cache {
let last = indice_list.last().expect("qed");
for (key, indices) in cache {
let last = indices.last().expect("qed");
collector.insert(
sharded_key_factory(*key, *last),
BlockNumberList::new_pre_sorted(indice_list),
BlockNumberList::new_pre_sorted(indices.iter().copied()),
)?;
}
Ok::<(), StageError>(())
Expand Down
1 change: 1 addition & 0 deletions crates/storage/db-api/src/models/integer_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ impl Compress for IntegerList {
fn compress(self) -> Self::Compressed {
self.to_bytes()
}

fn compress_to_buf<B: bytes::BufMut + AsMut<[u8]>>(self, buf: &mut B) {
self.to_mut_bytes(buf)
}
Expand Down
6 changes: 3 additions & 3 deletions crates/storage/db/src/implementation/mdbx/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1319,7 +1319,7 @@ mod tests {

for i in 1..5 {
let key = ShardedKey::new(real_key, i * 100);
let list: IntegerList = vec![i * 100u64].into();
let list = IntegerList::new_pre_sorted([i * 100u64]);

db.update(|tx| tx.put::<AccountsHistory>(key.clone(), list.clone()).expect(""))
.unwrap();
Expand All @@ -1340,7 +1340,7 @@ mod tests {
.expect("should be able to retrieve it.");

assert_eq!(ShardedKey::new(real_key, 200), key);
let list200: IntegerList = vec![200u64].into();
let list200 = IntegerList::new_pre_sorted([200u64]);
assert_eq!(list200, list);
}
// Seek greatest index
Expand All @@ -1357,7 +1357,7 @@ mod tests {
.expect("should be able to retrieve it.");

assert_eq!(ShardedKey::new(real_key, 400), key);
let list400: IntegerList = vec![400u64].into();
let list400 = IntegerList::new_pre_sorted([400u64]);
assert_eq!(list400, list);
}
}
Expand Down
13 changes: 7 additions & 6 deletions crates/storage/db/src/tables/codecs/fuzz/inputs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,13 @@ pub struct IntegerListInput(pub Vec<u64>);
impl From<IntegerListInput> for IntegerList {
fn from(list: IntegerListInput) -> Self {
let mut v = list.0;

// Empty lists are not supported by `IntegerList`, so we want to skip these cases.
if v.is_empty() {
return vec![1u64].into()
}
v.sort();
v.into()
let list = if v.is_empty() {
&[1u64][..]
} else {
v.sort_unstable();
&v
};
IntegerList::new_pre_sorted(list.iter().copied())
}
}
2 changes: 1 addition & 1 deletion crates/storage/provider/src/providers/database/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1356,7 +1356,7 @@ impl<TX: DbTxMut + DbTx, Spec: Send + Sync> DatabaseProvider<TX, Spec> {
};
self.tx.put::<T>(
sharded_key_factory(partial_key, highest_block_number),
BlockNumberList::new_pre_sorted(list),
BlockNumberList::new_pre_sorted(list.iter().copied()),
)?;
}
}
Expand Down
1 change: 0 additions & 1 deletion examples/custom-inspector/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ license.workspace = true
[dependencies]
reth.workspace = true
reth-node-ethereum.workspace = true
reth-rpc-types.workspace = true
alloy-rpc-types.workspace = true
clap = { workspace = true, features = ["derive"] }
futures-util.workspace = true
Expand Down

0 comments on commit 45f7a77

Please sign in to comment.