Skip to content

Commit

Permalink
Optimize prefix_proof by reducing the number of iterations (#363)
Browse files Browse the repository at this point in the history 
* Optimize prefix_proof by reducing the number of iterations

* Sort keys in the test

* Update storage_prefix_keys_ordered

* CHANGELOG

* Rustfmt
  • Loading branch information
@tomaka
tomaka authored Mar 29, 2023
1 parent f09b506 commit 29c6ccc
Show file tree
Hide file tree
Showing 5 changed files with 203 additions and 94 deletions.
118 changes: 95 additions & 23 deletions lib/src/trie/prefix_proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@

use super::{nibble, proof_decode};

use alloc::{vec, vec::Vec};
use alloc::{borrow::ToOwned as _, vec, vec::Vec};
use core::{fmt, iter, mem};

mod tests;
Expand All @@ -49,7 +49,10 @@ pub struct Config<'a> {
pub fn prefix_scan(config: Config<'_>) -> PrefixScan {
PrefixScan {
trie_root_hash: config.trie_root_hash,
next_queries: vec![nibble::bytes_to_nibbles(config.prefix.iter().copied()).collect()],
next_queries: vec![(
nibble::bytes_to_nibbles(config.prefix.iter().copied()).collect(),
QueryTy::Exact,
)],
final_result: Vec::with_capacity(32),
}
}
Expand All @@ -58,17 +61,29 @@ pub fn prefix_scan(config: Config<'_>) -> PrefixScan {
pub struct PrefixScan {
trie_root_hash: [u8; 32],
// TODO: we have lots of Vecs here; maybe find a way to optimize
next_queries: Vec<Vec<nibble::Nibble>>,
next_queries: Vec<(Vec<nibble::Nibble>, QueryTy)>,
// TODO: we have lots of Vecs here; maybe find a way to optimize
final_result: Vec<Vec<u8>>,
}

#[derive(Copy, Clone, Debug)]
enum QueryTy {
/// Expect to find a trie node with this exact key.
Exact,
/// The last nibble of the key is a dummy to force the remote to prove to us either that this
/// node exists or that this node doesn't exist, and if it doesn't exist prove it by including
/// the "actual" child that we're looking for in the proof.
/// It is guaranteed that the trie contains a node whose key is the requested key without its
/// last nibble.
Direction,
}

impl PrefixScan {
/// Returns the list of keys whose storage proof must be queried.
pub fn requested_keys(
&'_ self,
) -> impl Iterator<Item = impl Iterator<Item = nibble::Nibble> + '_> + '_ {
self.next_queries.iter().map(|l| l.iter().copied())
self.next_queries.iter().map(|(l, _)| l.iter().copied())
}

/// Injects the proof presumably containing the keys returned by [`PrefixScan::requested_keys`].
Expand All @@ -93,22 +108,65 @@ impl PrefixScan {
let mut next = Vec::with_capacity(non_terminal_queries.len() * 2);

debug_assert!(!non_terminal_queries.is_empty());
while let Some(query) = non_terminal_queries.pop() {
let info = match decoded_proof.trie_node_info(&query) {
Some(info) => info,
None if !is_first_iteration => {
// Node not in the proof. There's no point in adding this node to `next`
// as we will fail again if we try to verify the proof again.
// If `is_first_iteration`, it means that the proof is incorrect.
self.next_queries.push(query);
continue;
}
None => {
// Push all the non-processed queries back to `next_queries` before
// returning the error, so that we can try again.
self.next_queries.push(query);
self.next_queries.extend(non_terminal_queries);
return Err((self, Error::MissingProofEntry));
while let Some((query_key, query_ty)) = non_terminal_queries.pop() {
// Get the information from the proof about this key.
// If the query type is "direction", then instead we look up the parent (that we
// know for sure exists in the trie) then find the child.
let info = {
let info_of_node = match query_ty {
QueryTy::Exact => &query_key[..],
QueryTy::Direction => &query_key[..query_key.len() - 1],
};

match (decoded_proof.trie_node_info(info_of_node), query_ty) {
(Some(info), QueryTy::Exact) => info,
(Some(info), QueryTy::Direction) => {
match info.children.child(query_key[query_key.len() - 1]) {
proof_decode::Child::InProof { child_key } => {
// Rather than complicate this code, we just add the child to
// `next` (this time an `Exact` query) and process it during
// the next iteration.
next.push((child_key.to_owned(), QueryTy::Exact));
continue;
}
proof_decode::Child::AbsentFromProof if !is_first_iteration => {
// Node not in the proof. There's no point in adding this node
// to `next` as we will fail again if we try to verify the
// proof again.
// If `is_first_iteration`, it means that the proof is
// incorrect.
self.next_queries.push((query_key, QueryTy::Direction));
continue;
}
proof_decode::Child::AbsentFromProof => {
// Push all the non-processed queries back to `next_queries`
// before returning the error, so that we can try again.
self.next_queries.push((query_key, QueryTy::Direction));
self.next_queries.extend(non_terminal_queries);
return Err((self, Error::MissingProofEntry));
}
proof_decode::Child::NoChild => {
// We know for sure that there is a child in this direction,
// otherwise the query wouldn't have been added to this
// state machine.
unreachable!()
}
}
}
(None, _) if !is_first_iteration => {
// Node not in the proof. There's no point in adding this node to `next`
// as we will fail again if we try to verify the proof again.
// If `is_first_iteration`, it means that the proof is incorrect.
self.next_queries.push((query_key, query_ty));
continue;
}
(None, _) => {
// Push all the non-processed queries back to `next_queries` before
// returning the error, so that we can try again.
self.next_queries.push((query_key, query_ty));
self.next_queries.extend(non_terminal_queries);
return Err((self, Error::MissingProofEntry));
}
}
};

Expand All @@ -119,8 +177,8 @@ impl PrefixScan {
) {
// Trie nodes with a value are always aligned to "bytes-keys". In other words,
// the number of nibbles is always even.
debug_assert_eq!(query.len() % 2, 0);
let key = query
debug_assert_eq!(query_key.len() % 2, 0);
let key = query_key
.chunks(2)
.map(|n| (u8::from(n[0]) << 4) | u8::from(n[1]))
.collect::<Vec<_>>();
Expand All @@ -132,7 +190,21 @@ impl PrefixScan {

// For each child of the node, put into `next` the key that goes towards this
// child.
next.extend(info.children.unfold_append_to_key(query));
for (nibble, child) in info.children.children().enumerate() {
match child {
proof_decode::Child::NoChild => continue,
proof_decode::Child::AbsentFromProof => {
let mut direction = query_key.clone();
direction.push(
nibble::Nibble::try_from(u8::try_from(nibble).unwrap()).unwrap(),
);
next.push((direction, QueryTy::Direction));
}
proof_decode::Child::InProof { child_key } => {
next.push((child_key.to_owned(), QueryTy::Exact))
}
}
}
}

// Finished when nothing more to request.
Expand Down
7 changes: 5 additions & 2 deletions lib/src/trie/prefix_proof/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14499,8 +14499,11 @@ fn regression_test_174() {
prefix_scan = scan;
continue;
}
Ok(ResumeOutcome::Success { keys }) => {
assert_eq!(keys, EXPECTED);
Ok(ResumeOutcome::Success { mut keys }) => {
let mut expected = EXPECTED.to_owned();
expected.sort();
keys.sort();
assert_eq!(keys, expected);
return;
}
Err((_, err)) => panic!("{err:?}"),
Expand Down
135 changes: 82 additions & 53 deletions lib/src/trie/proof_decode.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -468,9 +468,8 @@ impl<T: AsRef<[u8]>> DecodedTrieProof<T> {
_ => None,
},
trie_node_info: TrieNodeInfo {
children: Children {
children_bitmap: *children_bitmap,
},
children: self.children_from_key(key, *children_bitmap),

storage_value,
},
},
Expand All @@ -479,6 +478,36 @@ impl<T: AsRef<[u8]>> DecodedTrieProof<T> {
)
}

fn children_from_key(&self, key: &Vec<nibble::Nibble>, children_bitmap: u16) -> Children {
debug_assert_eq!(self.entries.get(key).unwrap().2, children_bitmap);

let mut children = [Child::NoChild; 16];
let mut child_search = key.clone();

for nibble in nibble::all_nibbles().filter(|n| (children_bitmap & (1 << u8::from(*n))) != 0)
{
child_search.push(nibble);

children[usize::from(u8::from(nibble))] = if let Some((child, _)) = self
.entries
.range::<[nibble::Nibble], _>((
ops::Bound::Included(&child_search[..]),
ops::Bound::Unbounded,
))
.next()
.filter(|(maybe_child, _)| maybe_child.starts_with(&child_search))
{
Child::InProof { child_key: child }
} else {
Child::AbsentFromProof
};

child_search.pop();
}

Children { children }
}

/// Returns information about a trie node.
///
/// Returns `None` if the proof doesn't contain enough information about this trie node.
Expand Down Expand Up @@ -517,7 +546,9 @@ impl<T: AsRef<[u8]>> DecodedTrieProof<T> {
// that it doesn't exist.
return Some(TrieNodeInfo {
storage_value: StorageValue::None,
children: Children { children_bitmap: 0 },
children: Children {
children: [Child::NoChild; 16],
},
});
}
Some((found_key, (storage_value, _, children_bitmap))) if *found_key == key => {
Expand All @@ -541,9 +572,7 @@ impl<T: AsRef<[u8]>> DecodedTrieProof<T> {
)
}
},
children: Children {
children_bitmap: *children_bitmap,
},
children: self.children_from_key(found_key, *children_bitmap),
});
}
Some((found_key, (_, _, children_bitmap))) if key.starts_with(found_key) => {
Expand All @@ -553,26 +582,10 @@ impl<T: AsRef<[u8]>> DecodedTrieProof<T> {
if children_bitmap & (1 << u8::from(key[found_key.len()])) == 0 {
// Child absent.
// It has been proven that the requested key doesn't exist in the trie.
return Some(TrieNodeInfo {
storage_value: StorageValue::None,
children: Children { children_bitmap: 0 },
});
} else if let Some((descendant, _)) = self
.entries
.range::<[nibble::Nibble], _>((
ops::Bound::Included(key),
ops::Bound::Unbounded,
))
.next()
.filter(|(k, _)| k.starts_with(key))
{
// There exists a node in the proof that starts with `key`. Consequently,
// we know that `key` doesn't correspond to a node that exists.
let nibble = descendant[key.len()];
return Some(TrieNodeInfo {
storage_value: StorageValue::None,
children: Children {
children_bitmap: 1 << u8::from(nibble),
children: [Child::NoChild; 16],
},
});
} else if self
Expand All @@ -593,7 +606,9 @@ impl<T: AsRef<[u8]>> DecodedTrieProof<T> {
// Thus, the requested key doesn't exist in the trie.
return Some(TrieNodeInfo {
storage_value: StorageValue::None,
children: Children { children_bitmap: 0 },
children: Children {
children: [Child::NoChild; 16],
},
});
} else {
// Child present.
Expand Down Expand Up @@ -746,55 +761,69 @@ pub struct TrieNodeInfo<'a> {
/// Storage value of the node, if any.
pub storage_value: StorageValue<'a>,
/// Which children the node has.
pub children: Children,
pub children: Children<'a>,
}

/// See [`TrieNodeInfo::children`].
#[derive(Copy, Clone)]
pub struct Children {
/// If `(children_bitmap & (1 << n)) == 1` (where `n is in 0..16`), then this node has a
/// child whose key starts with the key of the parent, followed with
/// `Nibble::try_from(n).unwrap()`, followed with 0 or more extra nibbles unknown here.
children_bitmap: u16,
pub struct Children<'a> {
children: [Child<'a>; 16],
}

/// Information about a specific child in the list of children.
#[derive(Copy, Clone)]
pub enum Child<'a> {
/// Child exists and can be found in the proof.
InProof {
/// Key of the child. Always starts with the key of its parent.
child_key: &'a [nibble::Nibble],
},
/// Child exists but isn't present in the proof.
AbsentFromProof,
/// Child doesn't exist.
NoChild,
}

impl Children {
impl<'a> Children<'a> {
/// Returns `true` if a child in the direction of the given nibble is present.
pub fn has_child(&self, nibble: nibble::Nibble) -> bool {
self.children_bitmap & (1 << u8::from(nibble)) != 0
match self.children[usize::from(u8::from(nibble))] {
Child::InProof { .. } | Child::AbsentFromProof => true,
Child::NoChild => false,
}
}

/// Iterates over all the children of the node. For each child, contains the nibble that must
/// be appended to the key of the node in order to find the child.
pub fn next_nibbles(&'_ self) -> impl Iterator<Item = nibble::Nibble> + '_ {
nibble::all_nibbles().filter(move |n| (self.children_bitmap & (1 << u8::from(*n)) != 0))
/// Returns the information about the child in the given direction.
pub fn child(&self, direction: nibble::Nibble) -> Child<'a> {
self.children[usize::from(u8::from(direction))]
}

/// Iterators over all the children of the node. Returns an iterator producing one element per
/// child, where the element is `key` plus the nibble of this child.
pub fn unfold_append_to_key(
/// Returns an iterator of 16 items, one for each child.
pub fn children(
&'_ self,
key: Vec<nibble::Nibble>,
) -> impl Iterator<Item = Vec<nibble::Nibble>> + '_ {
nibble::all_nibbles()
.filter(move |n| (self.children_bitmap & (1 << u8::from(*n)) != 0))
.map(move |nibble| {
let mut k = key.clone();
k.push(nibble);
k
})
) -> impl DoubleEndedIterator + ExactSizeIterator<Item = Child<'a>> + '_ {
self.children.iter().copied()
}
}

impl fmt::Debug for Children {
impl<'a> fmt::Debug for Children<'a> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:016b}", self.children_bitmap)
fmt::Binary::fmt(&self, f)
}
}

impl fmt::Binary for Children {
impl<'a> fmt::Binary for Children<'a> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:016b}", self.children_bitmap)
for child in &self.children {
let chr = match child {
Child::InProof { .. } | Child::AbsentFromProof => '1',
Child::NoChild => '0',
};

fmt::Write::write_char(f, chr)?
}

Ok(())
}
}

Expand Down
Loading

0 comments on commit 29c6ccc

Please sign in to comment.