perf: improve IntegerList API to avoid allocations

paradigmxyz · Sep 27, 2024 · 90a4fd1 · 90a4fd1
1 parent fbb0b11
commit 90a4fd1
Show file tree

Hide file tree

Showing 12 changed files with 34 additions and 58 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/primitives-traits/src/integer_list.rs b/crates/primitives-traits/src/integer_list.rs
@@ -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.
@@ -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()
@@ -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))
     }
 }
@@ -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)
     }
 }

diff --git a/crates/prune/prune/src/segments/user/account_history.rs b/crates/prune/prune/src/segments/user/account_history.rs
@@ -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<_>>();

diff --git a/crates/prune/prune/src/segments/user/storage_history.rs b/crates/prune/prune/src/segments/user/storage_history.rs
@@ -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<_>>();

diff --git a/crates/stages/stages/src/stages/index_account_history.rs b/crates/stages/stages/src/stages/index_account_history.rs
@@ -182,7 +182,7 @@ mod tests {
     }
 
     fn list(list: &[u64]) -> BlockNumberList {
-        BlockNumberList::new(list).unwrap()
+        BlockNumberList::new(list.iter().copied()).unwrap()
     }
 
     fn cast(

diff --git a/crates/stages/stages/src/stages/index_storage_history.rs b/crates/stages/stages/src/stages/index_storage_history.rs
@@ -197,7 +197,7 @@ mod tests {
     }
 
     fn list(list: &[u64]) -> BlockNumberList {
-        BlockNumberList::new(list).unwrap()
+        BlockNumberList::new(list.iter().copied()).unwrap()
     }
 
     fn cast(

diff --git a/crates/stages/stages/src/stages/utils.rs b/crates/stages/stages/src/stages/utils.rs
@@ -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>(())

diff --git a/crates/storage/db-api/src/models/integer_list.rs b/crates/storage/db-api/src/models/integer_list.rs
@@ -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)
     }

diff --git a/crates/storage/db/src/implementation/mdbx/mod.rs b/crates/storage/db/src/implementation/mdbx/mod.rs
@@ -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();
@@ -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
@@ -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);
         }
     }

diff --git a/crates/storage/db/src/tables/codecs/fuzz/inputs.rs b/crates/storage/db/src/tables/codecs/fuzz/inputs.rs
@@ -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())
     }
 }
diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs
@@ -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()),
                 )?;
             }
         }

diff --git a/examples/custom-inspector/Cargo.toml b/examples/custom-inspector/Cargo.toml
@@ -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