reth_trie_common/hash_builder/

state.rs

use crate::TrieMask;
use alloc::vec::Vec;
use alloy_trie::{hash_builder::HashBuilderValue, nodes::RlpNode, HashBuilder};
use nybbles::Nibbles;

/// The hash builder state for storing in the database.
/// Check the `reth-trie` crate for more info on hash builder.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
pub struct HashBuilderState {
    /// The current key.
    pub key: Vec<u8>,
    /// The current node value.
    pub value: HashBuilderValue,
    /// The builder stack.
    pub stack: Vec<RlpNode>,

    /// Group masks.
    pub groups: Vec<TrieMask>,
    /// Tree masks.
    pub tree_masks: Vec<TrieMask>,
    /// Hash masks.
    pub hash_masks: Vec<TrieMask>,

    /// Flag indicating if the current node is stored in the database.
    pub stored_in_database: bool,
}

impl From<HashBuilderState> for HashBuilder {
    fn from(state: HashBuilderState) -> Self {
        Self {
            key: Nibbles::from_nibbles_unchecked(state.key),
            stack: state.stack,
            value: state.value,
            state_masks: state.groups,
            tree_masks: state.tree_masks,
            hash_masks: state.hash_masks,
            stored_in_database: state.stored_in_database,
            updated_branch_nodes: None,
            proof_retainer: None,
            rlp_buf: Vec::with_capacity(32),
            all_branch_nodes: false,
        }
    }
}

impl From<HashBuilder> for HashBuilderState {
    fn from(state: HashBuilder) -> Self {
        Self {
            key: state.key.to_vec(),
            stack: state.stack,
            value: state.value,
            groups: state.state_masks,
            tree_masks: state.tree_masks,
            hash_masks: state.hash_masks,
            stored_in_database: state.stored_in_database,
        }
    }
}

#[cfg(any(test, feature = "reth-codec"))]
impl reth_codecs::Compact for HashBuilderState {
    fn to_compact<B>(&self, buf: &mut B) -> usize
    where
        B: bytes::BufMut + AsMut<[u8]>,
    {
        let mut len = 0;

        // Collection codecs return flag bits, not their encoded byte length.
        let start = buf.remaining_mut();
        self.key.to_compact(buf);
        len += start - buf.remaining_mut();

        buf.put_u16(self.stack.len() as u16);
        len += 2;
        for item in &self.stack {
            buf.put_u16(item.len() as u16);
            buf.put_slice(&item[..]);
            len += 2 + item.len();
        }

        let start = buf.remaining_mut();
        self.value.to_compact(buf);
        len += start - buf.remaining_mut();

        buf.put_u16(self.groups.len() as u16);
        len += 2;
        for item in &self.groups {
            len += (*item).to_compact(buf);
        }

        buf.put_u16(self.tree_masks.len() as u16);
        len += 2;
        for item in &self.tree_masks {
            len += (*item).to_compact(buf);
        }

        buf.put_u16(self.hash_masks.len() as u16);
        len += 2;
        for item in &self.hash_masks {
            len += (*item).to_compact(buf);
        }

        buf.put_u8(self.stored_in_database as u8);
        len += 1;
        len
    }

    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) {
        use bytes::Buf;

        let (key, mut buf) = Vec::from_compact(buf, 0);

        let stack_len = buf.get_u16() as usize;
        let mut stack = Vec::with_capacity(stack_len);
        for _ in 0..stack_len {
            let item_len = buf.get_u16() as usize;
            stack.push(RlpNode::from_raw(&buf[..item_len]).unwrap());
            buf.advance(item_len);
        }

        let (value, mut buf) = HashBuilderValue::from_compact(buf, 0);

        let groups_len = buf.get_u16() as usize;
        let mut groups = Vec::with_capacity(groups_len);
        for _ in 0..groups_len {
            let (item, rest) = TrieMask::from_compact(buf, 0);
            groups.push(item);
            buf = rest;
        }

        let tree_masks_len = buf.get_u16() as usize;
        let mut tree_masks = Vec::with_capacity(tree_masks_len);
        for _ in 0..tree_masks_len {
            let (item, rest) = TrieMask::from_compact(buf, 0);
            tree_masks.push(item);
            buf = rest;
        }

        let hash_masks_len = buf.get_u16() as usize;
        let mut hash_masks = Vec::with_capacity(hash_masks_len);
        for _ in 0..hash_masks_len {
            let (item, rest) = TrieMask::from_compact(buf, 0);
            hash_masks.push(item);
            buf = rest;
        }

        let stored_in_database = buf.get_u8() != 0;
        (Self { key, stack, value, groups, tree_masks, hash_masks, stored_in_database }, buf)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(feature = "arbitrary")]
    use proptest::{collection::vec, prelude::*};
    use reth_codecs::Compact;

    #[test]
    fn hash_builder_state_regression() {
        let mut state = HashBuilderState::default();
        state.stack.push(Default::default());
        let mut buf = vec![];
        let len = state.clone().to_compact(&mut buf);
        let (decoded, _) = HashBuilderState::from_compact(&buf, len);
        assert_eq!(state, decoded);
    }

    #[test]
    fn hash_builder_state_fixed_slice() {
        let mut state = HashBuilderState { key: vec![1, 2, 3], ..Default::default() };
        state.value.set_bytes_owned(vec![42; 128]);
        let mut expected = vec![];
        assert_eq!(state.to_compact(&mut expected), expected.len());

        let mut storage = vec![0; expected.len() + 1];
        let mut remaining = storage.as_mut_slice();
        let len = state.to_compact(&mut remaining);
        assert_eq!(len, expected.len());
        assert_eq!(remaining.len(), 1);
        assert_eq!(&storage[..len], expected);
        let (decoded, rest) = HashBuilderState::from_compact(&storage[..len], len);
        assert_eq!(decoded, state);
        assert!(rest.is_empty());
    }

    // Generate fields directly: the arbitrary interop's fixed byte buffer can run out while
    // generating variable-length fields, aborting the test before checking the codec.
    #[cfg(feature = "arbitrary")]
    proptest::proptest! {
        #[test]
        fn hash_builder_state_roundtrip(
            key in vec(0u8..16, 0..=64),
            value in any::<HashBuilderValue>(),
            stack in vec(any::<RlpNode>(), 0..=64),
            groups in vec(any::<TrieMask>(), 0..=64),
            tree_masks in vec(any::<TrieMask>(), 0..=64),
            hash_masks in vec(any::<TrieMask>(), 0..=64),
            stored_in_database in any::<bool>(),
        ) {
            let state = HashBuilderState {
                key, value, stack, groups, tree_masks, hash_masks, stored_in_database,
            };
            let mut buf = vec![];
            let len = state.to_compact(&mut buf);
            prop_assert_eq!(len, buf.len());
            let (decoded, rest) = HashBuilderState::from_compact(&buf, len);
            prop_assert_eq!(state, decoded);
            prop_assert!(rest.is_empty());
        }
    }
}