From 062720f62ef6d3126cf3a12299e8c6e3de2d0e16 Mon Sep 17 00:00:00 2001 From: Michael Sproul Date: Tue, 15 Feb 2022 17:45:53 +1100 Subject: [PATCH] Use SmallVec in Bitfield --- consensus/ssz_types/Cargo.toml | 1 + consensus/ssz_types/src/bitfield.rs | 35 +++++++++++-------- consensus/types/Cargo.toml | 1 + .../src/test_utils/test_random/bitfield.rs | 5 +-- 4 files changed, 26 insertions(+), 16 deletions(-) diff --git a/consensus/ssz_types/Cargo.toml b/consensus/ssz_types/Cargo.toml index b71de4ccdb..9c23ce92b5 100644 --- a/consensus/ssz_types/Cargo.toml +++ b/consensus/ssz_types/Cargo.toml @@ -18,6 +18,7 @@ eth2_ssz = "0.4.1" typenum = "1.12.0" arbitrary = { version = "1.0", features = ["derive"], optional = true } derivative = "2.1.1" +smallvec = "1.8.0" [dev-dependencies] serde_json = "1.0.58" diff --git a/consensus/ssz_types/src/bitfield.rs b/consensus/ssz_types/src/bitfield.rs index 57b262bea1..d6d972677e 100644 --- a/consensus/ssz_types/src/bitfield.rs +++ b/consensus/ssz_types/src/bitfield.rs @@ -5,10 +5,17 @@ use derivative::Derivative; use eth2_serde_utils::hex::{encode as hex_encode, PrefixedHexVisitor}; use serde::de::{Deserialize, Deserializer}; use serde::ser::{Serialize, Serializer}; +use smallvec::{smallvec, SmallVec, ToSmallVec}; use ssz::{Decode, Encode}; use tree_hash::Hash256; use typenum::Unsigned; +/// Maximum number of bytes to store on the stack in a bitfield's `SmallVec`. +/// +/// The default of 32 bytes is enough to take us through to ~500K validators, as the byte length of +/// attestation bitfields is roughly `N // 32 slots // 64 committes // 8 bits`. +pub const SMALLVEC_LEN: usize = 32; + /// A marker trait applied to `Variable` and `Fixed` that defines the behaviour of a `Bitfield`. pub trait BitfieldBehaviour: Clone {} @@ -91,7 +98,7 @@ pub type BitVector = Bitfield>; #[derive(Clone, Debug, Derivative)] #[derivative(PartialEq, Hash(bound = ""))] pub struct Bitfield { - bytes: Vec, + bytes: SmallVec<[u8; SMALLVEC_LEN]>, len: usize, _phantom: PhantomData, } @@ -106,7 +113,7 @@ impl Bitfield> { pub fn with_capacity(num_bits: usize) -> Result { if num_bits <= N::to_usize() { Ok(Self { - bytes: vec![0; bytes_for_bit_len(num_bits)], + bytes: smallvec![0; bytes_for_bit_len(num_bits)], len: num_bits, _phantom: PhantomData, }) @@ -138,7 +145,7 @@ impl Bitfield> { /// /// assert_eq!(b.into_bytes(), vec![0b0001_0000]); /// ``` - pub fn into_bytes(self) -> Vec { + pub fn into_bytes(self) -> SmallVec<[u8; SMALLVEC_LEN]> { let len = self.len(); let mut bytes = self.bytes; @@ -163,7 +170,7 @@ impl Bitfield> { /// produces (SSZ). /// /// Returns `None` if `bytes` are not a valid encoding. - pub fn from_bytes(bytes: Vec) -> Result { + pub fn from_bytes(bytes: SmallVec<[u8; SMALLVEC_LEN]>) -> Result { let bytes_len = bytes.len(); let mut initial_bitfield: Bitfield> = { let num_bits = bytes.len() * 8; @@ -235,7 +242,7 @@ impl Bitfield> { /// All bits are initialized to `false`. pub fn new() -> Self { Self { - bytes: vec![0; bytes_for_bit_len(Self::capacity())], + bytes: smallvec![0; bytes_for_bit_len(Self::capacity())], len: Self::capacity(), _phantom: PhantomData, } @@ -258,7 +265,7 @@ impl Bitfield> { /// /// assert_eq!(BitVector4::new().into_bytes(), vec![0b0000_0000]); /// ``` - pub fn into_bytes(self) -> Vec { + pub fn into_bytes(self) -> SmallVec<[u8; SMALLVEC_LEN]> { self.into_raw_bytes() } @@ -266,7 +273,7 @@ impl Bitfield> { /// produces (SSZ). /// /// Returns `None` if `bytes` are not a valid encoding. - pub fn from_bytes(bytes: Vec) -> Result { + pub fn from_bytes(bytes: SmallVec<[u8; SMALLVEC_LEN]>) -> Result { Self::from_raw_bytes(bytes, Self::capacity()) } @@ -355,7 +362,7 @@ impl Bitfield { } /// Returns the underlying bytes representation of the bitfield. - pub fn into_raw_bytes(self) -> Vec { + pub fn into_raw_bytes(self) -> SmallVec<[u8; SMALLVEC_LEN]> { self.bytes } @@ -372,9 +379,9 @@ impl Bitfield { /// - `bytes` is not the minimal required bytes to represent a bitfield of `bit_len` bits. /// - `bit_len` is not a multiple of 8 and `bytes` contains set bits that are higher than, or /// equal to `bit_len`. - fn from_raw_bytes(bytes: Vec, bit_len: usize) -> Result { + fn from_raw_bytes(bytes: SmallVec<[u8; SMALLVEC_LEN]>, bit_len: usize) -> Result { if bit_len == 0 { - if bytes.len() == 1 && bytes == [0] { + if bytes.len() == 1 && bytes[0] == 0 { // A bitfield with `bit_len` 0 can only be represented by a single zero byte. Ok(Self { bytes, @@ -512,7 +519,7 @@ impl Encode for Bitfield> { } fn ssz_append(&self, buf: &mut Vec) { - buf.append(&mut self.clone().into_bytes()) + buf.extend_from_slice(&self.clone().into_bytes()) } } @@ -522,7 +529,7 @@ impl Decode for Bitfield> { } fn from_ssz_bytes(bytes: &[u8]) -> Result { - Self::from_bytes(bytes.to_vec()).map_err(|e| { + Self::from_bytes(bytes.to_smallvec()).map_err(|e| { ssz::DecodeError::BytesInvalid(format!("BitList failed to decode: {:?}", e)) }) } @@ -542,7 +549,7 @@ impl Encode for Bitfield> { } fn ssz_append(&self, buf: &mut Vec) { - buf.append(&mut self.clone().into_bytes()) + buf.extend_from_slice(&self.clone().into_bytes()) } } @@ -556,7 +563,7 @@ impl Decode for Bitfield> { } fn from_ssz_bytes(bytes: &[u8]) -> Result { - Self::from_bytes(bytes.to_vec()).map_err(|e| { + Self::from_bytes(bytes.to_smallvec()).map_err(|e| { ssz::DecodeError::BytesInvalid(format!("BitVector failed to decode: {:?}", e)) }) } diff --git a/consensus/types/Cargo.toml b/consensus/types/Cargo.toml index 260368e917..07362e280e 100644 --- a/consensus/types/Cargo.toml +++ b/consensus/types/Cargo.toml @@ -47,6 +47,7 @@ superstruct = "0.4.0" serde_json = "1.0.74" milhouse = { path = "../../../milhouse", optional = true } rpds = "0.11.0" +smallvec = "1.8.0" [dev-dependencies] criterion = "0.3.3" diff --git a/consensus/types/src/test_utils/test_random/bitfield.rs b/consensus/types/src/test_utils/test_random/bitfield.rs index 2ba3576b77..5cb4e7d521 100644 --- a/consensus/types/src/test_utils/test_random/bitfield.rs +++ b/consensus/types/src/test_utils/test_random/bitfield.rs @@ -1,9 +1,10 @@ use super::*; use crate::{BitList, BitVector, Unsigned}; +use smallvec::smallvec; impl TestRandom for BitList { fn random_for_test(rng: &mut impl RngCore) -> Self { - let mut raw_bytes = vec![0; std::cmp::max(1, (N::to_usize() + 7) / 8)]; + let mut raw_bytes = smallvec![0; std::cmp::max(1, (N::to_usize() + 7) / 8)]; rng.fill_bytes(&mut raw_bytes); Self::from_bytes(raw_bytes).expect("we generate a valid BitList") } @@ -11,7 +12,7 @@ impl TestRandom for BitList { impl TestRandom for BitVector { fn random_for_test(rng: &mut impl RngCore) -> Self { - let mut raw_bytes = vec![0; std::cmp::max(1, (N::to_usize() + 7) / 8)]; + let mut raw_bytes = smallvec![0; std::cmp::max(1, (N::to_usize() + 7) / 8)]; rng.fill_bytes(&mut raw_bytes); Self::from_bytes(raw_bytes).expect("we generate a valid BitVector") }