Skip to content

Commit

Permalink
fix: arbitrary impls
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Oct 3, 2024
1 parent f4fb65d commit 3f22168
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 9 deletions.
2 changes: 1 addition & 1 deletion src/hash_builder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down
44 changes: 36 additions & 8 deletions src/hash_builder/value.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ 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.
#[cfg_attr(feature = "serde", serde(with = "hex"))]
buf: Vec<u8>,
/// The kind of value that is stored in `buf`.
kind: HashBuilderValueKind,
Expand All @@ -27,6 +27,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<Self> {
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<Self>;

fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
use proptest::prelude::*;

proptest::arbitrary::any::<HashBuilderValueKind>()
.prop_flat_map(|kind| {
let range = match kind {
HashBuilderValueKind::Bytes => 0..=128,
HashBuilderValueKind::Hash => 32..=32,
};
proptest::collection::vec(any::<u8>(), range)
.prop_map(move |buf| Self { buf, kind })
})
.boxed()
}
}

impl HashBuilderValue {
/// Creates a new empty value.
pub fn new() -> Self {
Expand All @@ -50,11 +83,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<u8>) {
self.buf = bytes;
Expand All @@ -81,7 +109,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,
Expand All @@ -107,7 +135,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,
Expand Down

0 comments on commit 3f22168

Please sign in to comment.