Skip to content

Commit

Permalink
Make Nibbles Copy
Browse files Browse the repository at this point in the history
  • Loading branch information
@richardpringle
richardpringle committed Sep 21, 2023
1 parent 3dc47ba commit 41ea07c
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 73 deletions.
26 changes: 17 additions & 9 deletions firewood/src/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -370,7 +370,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
let mut val = Some(val);

// walk down the merkle tree starting from next_node, currently the root
for (key_nib_offset, key_nib) in key_nibbles.iter().enumerate() {
for (key_nib_offset, key_nib) in key_nibbles.into_iter().enumerate() {
// special handling for extension nodes
if nskip > 0 {
nskip -= 1;
Expand All @@ -393,7 +393,9 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
// create a new leaf
let leaf_ptr = self
.new_node(Node::new(NodeType::Leaf(LeafNode(
PartialPath(key_nibbles.iter().skip(key_nib_offset + 1).collect()),
PartialPath(
key_nibbles.into_iter().skip(key_nib_offset + 1).collect(),
),
Data(val.take().unwrap()),
))))?
.as_ptr();
Expand All @@ -412,7 +414,10 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
// of the stored key to pass into split
let n_path = n.0.to_vec();
let n_value = Some(n.1.clone());
let rem_path = key_nibbles.iter().skip(key_nib_offset).collect::<Vec<_>>();
let rem_path = key_nibbles
.into_iter()
.skip(key_nib_offset)
.collect::<Vec<_>>();
self.split(
node,
&mut parents,
Expand All @@ -428,7 +433,10 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
let n_path = n.0.to_vec();
let n_ptr = n.1;
nskip = n_path.len() - 1;
let rem_path = key_nibbles.iter().skip(key_nib_offset).collect::<Vec<_>>();
let rem_path = key_nibbles
.into_iter()
.skip(key_nib_offset)
.collect::<Vec<_>>();

if let Some(v) = self.split(
node,
Expand Down Expand Up @@ -1044,7 +1052,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {

let mut nskip = 0;
let mut nodes: Vec<DiskAddress> = Vec::new();
for (i, nib) in key_nibbles.iter().enumerate() {
for (i, nib) in key_nibbles.into_iter().enumerate() {
if nskip > 0 {
nskip -= 1;
continue;
Expand All @@ -1060,7 +1068,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
// the key passed in must match the entire remainder of this
// extension node, otherwise we break out
let n_path = &*n.0;
let remaining_path = key_nibbles.iter().skip(i);
let remaining_path = key_nibbles.into_iter().skip(i);
if remaining_path.size_hint().0 < n_path.len() {
// all bytes aren't there
break;
Expand Down Expand Up @@ -1115,7 +1123,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
let mut u_ref = self.get_node(root)?;
let mut nskip = 0;

for (i, nib) in key_nibbles.iter().enumerate() {
for (i, nib) in key_nibbles.into_iter().enumerate() {
if nskip > 0 {
nskip -= 1;
continue;
Expand All @@ -1126,14 +1134,14 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
None => return Ok(None),
},
NodeType::Leaf(n) => {
if !key_nibbles.iter().skip(i).eq(n.0.iter().cloned()) {
if !key_nibbles.into_iter().skip(i).eq(n.0.iter().cloned()) {
return Ok(None);
}
return Ok(Some(Ref(u_ref)));
}
NodeType::Extension(n) => {
let n_path = &*n.0;
let rem_path = key_nibbles.iter().skip(i);
let rem_path = key_nibbles.into_iter().skip(i);
if rem_path.size_hint().0 < n_path.len() {
return Ok(None);
}
Expand Down
11 changes: 6 additions & 5 deletions firewood/src/merkle/partial_path.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE.md for licensing terms.

use crate::nibbles::NibblesIterator;
use std::fmt::{self, Debug};

use crate::nibbles::Nibbles;

/// PartialPath keeps a list of nibbles to represent a path on the Trie.
#[derive(PartialEq, Eq, Clone)]
pub struct PartialPath(pub Vec<u8>);
Expand Down Expand Up @@ -43,6 +42,8 @@ impl PartialPath {
}

// TODO: remove all non `Nibbles` usages and delete this function.
// I also think `PartialPath` could probably borrow instead of own data.
//
/// returns a tuple of the decoded partial path and whether the path is terminal
pub fn decode(raw: &[u8]) -> (Self, bool) {
let prefix = raw[0];
Expand All @@ -53,10 +54,10 @@ impl PartialPath {
}

/// returns a tuple of the decoded partial path and whether the path is terminal
pub fn from_nibbles<'a, const N: usize>(nibbles: Nibbles<'a, N>) -> (Self, bool) {
let prefix = nibbles[0];
pub fn from_nibbles<const N: usize>(mut nibbles: NibblesIterator<'_, N>) -> (Self, bool) {
let prefix = nibbles.next().unwrap();
let is_odd = (prefix & 1) as usize;
let decoded = nibbles.iter().skip(1).skip(1 - is_odd).collect();
let decoded = nibbles.skip(1 - is_odd).collect();

(Self(decoded), prefix > 1)
}
Expand Down
33 changes: 19 additions & 14 deletions firewood/src/nibbles.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ static NIBBLES: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15

/// Nibbles is a newtype that contains only a reference to a [u8], and produces
/// nibbles. Nibbles can be indexed using nib\[x\] or you can get an iterator
/// with iter()
/// with `into_iter()`
///
/// Nibbles can be constructed with a number of leading zeroes. This is used
/// in firewood because there is a sentinel node, so we always want the first
Expand All @@ -23,21 +23,21 @@ static NIBBLES: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
/// # use firewood::nibbles;
/// # fn main() {
/// let nib = nibbles::Nibbles::<0>::new(&[0x56, 0x78]);
/// assert_eq!(nib.iter().collect::<Vec<_>>(), [0x5, 0x6, 0x7, 0x8]);
/// assert_eq!(nib.into_iter().collect::<Vec<_>>(), [0x5, 0x6, 0x7, 0x8]);
///
/// // nibbles can be efficiently advanced without rendering the
/// // intermediate values
/// assert_eq!(nib.iter().skip(3).collect::<Vec<_>>(), [0x8]);
/// assert_eq!(nib.into_iter().skip(3).collect::<Vec<_>>(), [0x8]);
///
/// // nibbles can also be indexed
///
/// assert_eq!(nib[1], 0x6);
///
/// // or reversed
/// assert_eq!(nib.iter().rev().next(), Some(0x8));
/// assert_eq!(nib.into_iter().rev().next(), Some(0x8));
/// # }
/// ```
#[derive(Debug)]
#[derive(Debug, Copy, Clone)]
pub struct Nibbles<'a, const LEADING_ZEROES: usize>(&'a [u8]);

impl<'a, const LEADING_ZEROES: usize> Index<usize> for Nibbles<'a, LEADING_ZEROES> {
Expand All @@ -54,16 +54,21 @@ impl<'a, const LEADING_ZEROES: usize> Index<usize> for Nibbles<'a, LEADING_ZEROE
}
}

impl<'a, const LEADING_ZEROES: usize> Nibbles<'a, LEADING_ZEROES> {
impl<'a, const LEADING_ZEROES: usize> IntoIterator for Nibbles<'a, LEADING_ZEROES> {
type Item = u8;
type IntoIter = NibblesIterator<'a, LEADING_ZEROES>;

#[must_use]
pub fn iter(&self) -> NibblesIterator<'_, LEADING_ZEROES> {
fn into_iter(self) -> Self::IntoIter {
NibblesIterator {
data: self,
head: 0,
tail: self.len(),
}
}
}

impl<'a, const LEADING_ZEROES: usize> Nibbles<'a, LEADING_ZEROES> {
#[must_use]
pub fn len(&self) -> usize {
LEADING_ZEROES + 2 * self.0.len()
Expand All @@ -79,11 +84,11 @@ impl<'a, const LEADING_ZEROES: usize> Nibbles<'a, LEADING_ZEROES> {
}
}

/// An interator returned by [Nibbles::iter]
/// An interator returned by [Nibbles::into_iter]
/// See their documentation for details.
#[derive(Clone, Debug)]
pub struct NibblesIterator<'a, const LEADING_ZEROES: usize> {
data: &'a Nibbles<'a, LEADING_ZEROES>,
data: Nibbles<'a, LEADING_ZEROES>,
head: usize,
tail: usize,
}
Expand Down Expand Up @@ -161,14 +166,14 @@ mod test {
fn leading_zero_nibbles_iter() {
let nib = Nibbles::<1>(&TEST_BYTES);
let expected: [u8; 9] = [0u8, 0xd, 0xe, 0xa, 0xd, 0xb, 0xe, 0xe, 0xf];
expected.into_iter().eq(nib.iter());
expected.into_iter().eq(nib.into_iter());
}

#[test]
fn skip_skips_zeroes() {
let nib1 = Nibbles::<1>(&TEST_BYTES);
let nib0 = Nibbles::<0>(&TEST_BYTES);
assert!(nib1.iter().skip(1).eq(nib0.iter()));
assert!(nib1.into_iter().skip(1).eq(nib0.into_iter()));
}

#[test]
Expand All @@ -187,7 +192,7 @@ mod test {
#[test]
fn size_hint_0() {
let nib = Nibbles::<0>(&TEST_BYTES);
let mut nib_iter = nib.iter();
let mut nib_iter = nib.into_iter();
assert_eq!((8, Some(8)), nib_iter.size_hint());
let _ = nib_iter.next();
assert_eq!((7, Some(7)), nib_iter.size_hint());
Expand All @@ -196,7 +201,7 @@ mod test {
#[test]
fn size_hint_1() {
let nib = Nibbles::<1>(&TEST_BYTES);
let mut nib_iter = nib.iter();
let mut nib_iter = nib.into_iter();
assert_eq!((9, Some(9)), nib_iter.size_hint());
let _ = nib_iter.next();
assert_eq!((8, Some(8)), nib_iter.size_hint());
Expand All @@ -205,7 +210,7 @@ mod test {
#[test]
fn backwards() {
let nib = Nibbles::<1>(&TEST_BYTES);
let nib_iter = nib.iter().rev();
let nib_iter = nib.into_iter().rev();
let expected = [0xf, 0xe, 0xe, 0xb, 0xd, 0xa, 0xe, 0xd, 0x0];

assert!(nib_iter.eq(expected));
Expand Down
86 changes: 41 additions & 45 deletions firewood/src/proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ use shale::ShaleError;
use shale::ShaleStore;
use thiserror::Error;

use crate::nibbles::Nibbles;
use crate::nibbles::NibblesIterator;
use crate::{
db::DbError,
merkle::{
Expand Down Expand Up @@ -114,53 +116,45 @@ impl<N: AsRef<[u8]> + Send> Proof<N> {
key: K,
root_hash: [u8; 32],
) -> Result<Option<Vec<u8>>, ProofError> {
let mut key_nibbles = Vec::new();
key_nibbles.extend(key.as_ref().iter().copied().flat_map(to_nibble_array));
let mut key_nibbles = Nibbles::<0>::new(key.as_ref()).into_iter();

let mut remaining_key_nibbles: &[u8] = &key_nibbles;
let mut cur_hash = root_hash;
let proofs_map = &self.0;
let mut index = 0;

loop {
let cur_proof = proofs_map
.get(&cur_hash)
.ok_or(ProofError::ProofNodeMissing)?;
let (sub_proof, size) =
self.locate_subproof(remaining_key_nibbles, cur_proof.as_ref())?;
index += size;
let (sub_proof, traversed_nibbles) =
self.locate_subproof(key_nibbles, cur_proof.as_ref())?;
key_nibbles = traversed_nibbles;

match sub_proof {
cur_hash = match sub_proof {
// Return when reaching the end of the key.
Some(p) if index == key_nibbles.len() => return Ok(Some(p.rlp)),
Some(p) if key_nibbles.size_hint().0 == 0 => return Ok(Some(p.rlp)),
// The trie doesn't contain the key.
Some(SubProof { hash: None, .. }) | None => return Ok(None),
Some(p) => {
cur_hash = p.hash.unwrap();
remaining_key_nibbles = &key_nibbles[index..];
}
}
Some(SubProof {
hash: Some(hash), ..
}) => hash,
_ => return Ok(None),
};
}
}

fn locate_subproof(
fn locate_subproof<'a>(
&self,
key_nibbles: &[u8],
mut key_nibbles: NibblesIterator<'a, 0>,
rlp_encoded_node: &[u8],
) -> Result<(Option<SubProof>, usize), ProofError> {
) -> Result<(Option<SubProof>, NibblesIterator<'a, 0>), ProofError> {
let rlp = rlp::Rlp::new(rlp_encoded_node);

match rlp.item_count() {
Ok(EXT_NODE_SIZE) => {
let decoded_key_nibbles: Vec<_> = rlp
.at(0)
.unwrap()
.as_val::<Vec<u8>>()
.unwrap()
.into_iter()
.flat_map(to_nibble_array)
.collect();
let (cur_key_path, term) = PartialPath::decode(&decoded_key_nibbles);
let decoded_key = rlp.at(0).unwrap().as_val::<Vec<u8>>().unwrap();
let decoded_key_nibbles = Nibbles::<0>::new(&decoded_key);

let (cur_key_path, term) =
PartialPath::from_nibbles(decoded_key_nibbles.into_iter());
let cur_key = cur_key_path.into_inner();

let rlp = rlp.at(1).unwrap();
Expand All @@ -171,30 +165,32 @@ impl<N: AsRef<[u8]> + Send> Proof<N> {
rlp.as_raw().to_vec()
};

// Check if the key of current node match with the given key.
if key_nibbles.len() < cur_key.len() || key_nibbles[..cur_key.len()] != cur_key {
return Ok((None, 0));
// Check if the key of current node match with the given key
// and consume the current-key portion of the nibbles-iterator
let does_not_match = key_nibbles.size_hint().0 < cur_key.len()
|| !cur_key.iter().all(|val| key_nibbles.next() == Some(*val));

if does_not_match {
return Ok((None, Nibbles::<0>::new(&[]).into_iter()));
}

if term {
Ok((
Some(SubProof {
rlp: data,
hash: None,
}),
cur_key.len(),
))
let sub_proof = if term {
SubProof {
rlp: data,
hash: None,
}
} else {
self.generate_subproof(data)
.map(|subproof| (Some(subproof), cur_key.len()))
}
self.generate_subproof(data)?
};

Ok((sub_proof.into(), key_nibbles))
}

Ok(BRANCH_NODE_SIZE) if key_nibbles.is_empty() => Err(ProofError::NoSuchNode),
Ok(BRANCH_NODE_SIZE) if key_nibbles.size_hint().0 == 0 => Err(ProofError::NoSuchNode),

Ok(BRANCH_NODE_SIZE) => {
let index = key_nibbles[0];
let rlp = rlp.at(index as usize).unwrap();
let index = key_nibbles.next().unwrap() as usize;
let rlp = rlp.at(index).unwrap();

let data = if rlp.is_data() {
rlp.as_val::<Vec<u8>>().unwrap()
Expand All @@ -203,7 +199,7 @@ impl<N: AsRef<[u8]> + Send> Proof<N> {
};

self.generate_subproof(data)
.map(|subproof| (Some(subproof), 1))
.map(|subproof| (Some(subproof), key_nibbles))
}

Ok(_) => Err(ProofError::DecodeError(rlp::DecoderError::RlpInvalidLength)),
Expand Down

0 comments on commit 41ea07c

Please sign in to comment.