diff --git a/src/hash_builder/mod.rs b/src/hash_builder/mod.rs index 821e348..1ca8679 100644 --- a/src/hash_builder/mod.rs +++ b/src/hash_builder/mod.rs @@ -13,7 +13,7 @@ use core::cmp; use tracing::trace; mod value; -pub use value::{HashBuilderValue, HashBuilderValueKind, HashBuilderValueRef}; +pub use value::{HashBuilderValue, HashBuilderValueRef}; /// A component used to construct the root hash of the trie. /// diff --git a/src/hash_builder/value.rs b/src/hash_builder/value.rs index 5dab229..93a8bad 100644 --- a/src/hash_builder/value.rs +++ b/src/hash_builder/value.rs @@ -7,7 +7,6 @@ use core::fmt; /// Stores [`HashBuilderValueRef`] efficiently by reusing resources. #[derive(Clone, PartialEq, Eq)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[cfg_attr(feature = "arbitrary", derive(derive_arbitrary::Arbitrary, proptest_derive::Arbitrary))] pub struct HashBuilderValue { /// Stores the bytes of either the leaf node value or the hash of adjacent nodes. buf: Vec, @@ -27,6 +26,39 @@ impl fmt::Debug for HashBuilderValue { } } +#[cfg(feature = "arbitrary")] +impl<'u> arbitrary::Arbitrary<'u> for HashBuilderValue { + fn arbitrary(g: &mut arbitrary::Unstructured<'u>) -> arbitrary::Result { + let kind = HashBuilderValueKind::arbitrary(g)?; + let buf = match kind { + HashBuilderValueKind::Bytes => Vec::arbitrary(g)?, + HashBuilderValueKind::Hash => B256::arbitrary(g)?.to_vec(), + }; + Ok(Self { buf, kind }) + } +} + +#[cfg(feature = "arbitrary")] +impl proptest::arbitrary::Arbitrary for HashBuilderValue { + type Parameters = (); + type Strategy = proptest::strategy::BoxedStrategy; + + fn arbitrary_with((): Self::Parameters) -> Self::Strategy { + use proptest::prelude::*; + + proptest::arbitrary::any::() + .prop_flat_map(|kind| { + let range = match kind { + HashBuilderValueKind::Bytes => 0..=128, + HashBuilderValueKind::Hash => 32..=32, + }; + proptest::collection::vec(any::(), range) + .prop_map(move |buf| Self { buf, kind }) + }) + .boxed() + } +} + impl HashBuilderValue { /// Creates a new empty value. pub fn new() -> Self { @@ -50,11 +82,6 @@ impl HashBuilderValue { &self.buf } - /// Returns the kind of the value. - pub const fn kind(&self) -> HashBuilderValueKind { - self.kind - } - /// Like `set_from_ref`, but takes ownership of the bytes. pub fn set_bytes_owned(&mut self, bytes: Vec) { self.buf = bytes; @@ -81,7 +108,7 @@ impl HashBuilderValue { #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "arbitrary", derive(derive_arbitrary::Arbitrary, proptest_derive::Arbitrary))] -pub enum HashBuilderValueKind { +enum HashBuilderValueKind { /// Value of the leaf node. #[default] Bytes, @@ -107,7 +134,7 @@ impl<'a> HashBuilderValueRef<'a> { } /// Returns the kind of the value. - pub const fn kind(&self) -> HashBuilderValueKind { + const fn kind(&self) -> HashBuilderValueKind { match *self { HashBuilderValueRef::Bytes(_) => HashBuilderValueKind::Bytes, HashBuilderValueRef::Hash(_) => HashBuilderValueKind::Hash,