From d2b455fdc2670b9f2b723b740d938d8475dfbae3 Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Tue, 28 Mar 2023 17:03:12 +0200 Subject: [PATCH 1/5] Optimize prefix_proof by reducing the number of iterations --- lib/src/trie/prefix_proof.rs | 118 ++++++++++++++++++++++++------ lib/src/trie/proof_decode.rs | 135 +++++++++++++++++++++-------------- 2 files changed, 176 insertions(+), 77 deletions(-) diff --git a/lib/src/trie/prefix_proof.rs b/lib/src/trie/prefix_proof.rs index e8375bb0bc..079e742a83 100644 --- a/lib/src/trie/prefix_proof.rs +++ b/lib/src/trie/prefix_proof.rs @@ -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; @@ -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), } } @@ -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>, + next_queries: Vec<(Vec, QueryTy)>, // TODO: we have lots of Vecs here; maybe find a way to optimize final_result: Vec>, } +#[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 + '_> + '_ { - 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`]. @@ -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)); + } } }; @@ -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::>(); @@ -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. diff --git a/lib/src/trie/proof_decode.rs b/lib/src/trie/proof_decode.rs index d6b2f69af4..04d393d375 100644 --- a/lib/src/trie/proof_decode.rs +++ b/lib/src/trie/proof_decode.rs @@ -468,9 +468,8 @@ impl> DecodedTrieProof { _ => None, }, trie_node_info: TrieNodeInfo { - children: Children { - children_bitmap: *children_bitmap, - }, + children: self.children_from_key(key, *children_bitmap), + storage_value, }, }, @@ -479,6 +478,36 @@ impl> DecodedTrieProof { ) } + fn children_from_key(&self, key: &Vec, 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. @@ -517,7 +546,9 @@ impl> DecodedTrieProof { // 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 => { @@ -541,9 +572,7 @@ impl> DecodedTrieProof { ) } }, - 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) => { @@ -553,26 +582,10 @@ impl> DecodedTrieProof { 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 @@ -593,7 +606,9 @@ impl> DecodedTrieProof { // 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. @@ -746,55 +761,67 @@ 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], } -impl Children { +/// 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<'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 + '_ { - 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( - &'_ self, - key: Vec, - ) -> impl Iterator> + '_ { - 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 - }) + /// Returns an iterator of 16 items, one for each child. + pub fn children(&'_ self) -> impl ExactSizeIterator> + '_ { + 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(()) } } From 263ca92ab3e02c8afccacd59dc2ffee2856e2330 Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Tue, 28 Mar 2023 18:19:07 +0200 Subject: [PATCH 2/5] Sort keys in the test --- lib/src/trie/prefix_proof/tests.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/lib/src/trie/prefix_proof/tests.rs b/lib/src/trie/prefix_proof/tests.rs index 6de2b768d8..baea7a5058 100644 --- a/lib/src/trie/prefix_proof/tests.rs +++ b/lib/src/trie/prefix_proof/tests.rs @@ -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:?}"), From b498b8972feac95b96d871fc646159034a535586 Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Tue, 28 Mar 2023 19:36:49 +0200 Subject: [PATCH 3/5] Update storage_prefix_keys_ordered --- lib/src/trie/proof_decode.rs | 2 +- light-base/src/runtime_service.rs | 36 +++++++++++++++++-------------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/lib/src/trie/proof_decode.rs b/lib/src/trie/proof_decode.rs index 04d393d375..0b2bb36f92 100644 --- a/lib/src/trie/proof_decode.rs +++ b/lib/src/trie/proof_decode.rs @@ -799,7 +799,7 @@ impl<'a> Children<'a> { } /// Returns an iterator of 16 items, one for each child. - pub fn children(&'_ self) -> impl ExactSizeIterator> + '_ { + pub fn children(&'_ self) -> impl DoubleEndedIterator + ExactSizeIterator> + '_ { self.children.iter().copied() } } diff --git a/light-base/src/runtime_service.rs b/light-base/src/runtime_service.rs index 883c96f744..e0a30a670c 100644 --- a/light-base/src/runtime_service.rs +++ b/light-base/src/runtime_service.rs @@ -59,11 +59,10 @@ use crate::{platform::Platform, sync_service}; use alloc::{ borrow::ToOwned as _, boxed::Box, - collections::BTreeMap, + collections::{BTreeMap, VecDeque}, format, string::{String, ToString as _}, sync::{Arc, Weak}, - vec, vec::Vec, }; use core::{ @@ -875,7 +874,9 @@ impl<'a> RuntimeCallLock<'a> { prefix: &[u8], ) -> Result + '_>, RuntimeCallError> { // TODO: this could be a function in the proof_decode module - let mut to_find = vec![trie::bytes_to_nibbles(prefix.iter().copied()).collect::>()]; + let mut to_find = VecDeque::>::new(); + to_find.push_back(trie::bytes_to_nibbles(prefix.iter().copied()).collect::>()); + let mut output = Vec::new(); let call_proof = match &self.call_proof { @@ -883,7 +884,7 @@ impl<'a> RuntimeCallLock<'a> { Err(err) => return Err(err.clone()), }; - for key in mem::take(&mut to_find) { + while let Some(key) = to_find.pop_front() { let node_info = call_proof .trie_node_info(&key) .ok_or(RuntimeCallError::MissingProofEntry)?; @@ -894,24 +895,27 @@ impl<'a> RuntimeCallLock<'a> { | proof_decode::StorageValue::HashKnownValueMissing(_) ) { assert_eq!(key.len() % 2, 0); - output.push( - trie::nibbles_to_bytes_suffix_extend(key.iter().copied()).collect::>(), - ); + let key_as_bytes = + trie::nibbles_to_bytes_suffix_extend(key.iter().copied()).collect::>(); + debug_assert!(output.last().map_or(true, |last| *last < key_as_bytes)); + output.push(key_as_bytes); } - for nibble in trie::all_nibbles() { - if !node_info.children.has_child(nibble) { - continue; + for child in node_info.children.children().rev() { + match child { + proof_decode::Child::NoChild => continue, + proof_decode::Child::AbsentFromProof => { + return Err(RuntimeCallError::MissingProofEntry); + } + proof_decode::Child::InProof { child_key } => { + debug_assert!(to_find.front().map_or(true, |f| child_key < f)); + to_find.push_front(child_key.to_owned()); + } } - - let mut child = key.clone(); - child.push(nibble); - to_find.push(child); } } - // TODO: maybe we could iterate over the proof in an ordered way rather than sorting at the end - output.sort(); + // TODO: debug_assert!(output.is_sorted()); // TODO: https://github.com/rust-lang/rust/issues/53485 Ok(output.into_iter()) } From 847a60ab79688cbab56c77b3314be1d8c8656132 Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Tue, 28 Mar 2023 19:37:55 +0200 Subject: [PATCH 4/5] CHANGELOG --- wasm-node/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/wasm-node/CHANGELOG.md b/wasm-node/CHANGELOG.md index 9da4648a11..72d60f8705 100644 --- a/wasm-node/CHANGELOG.md +++ b/wasm-node/CHANGELOG.md @@ -10,6 +10,7 @@ - Improved the ganularity of the tasks that handle JSON-RPC requests and libp2p connections. Smoldot now yields more often to the browser, reducing the chances and the severity of freezes during the rendering of the web page. ([#349](https://github.com/smol-dot/smoldot/pull/349)) - Smoldot is now compiled with the `bulk-memory-operations` and `sign-extensions-ops` WebAssembly features enabled. This is expected to considerably speed up its execution. The minimum version required to run smoldot is now Chrome 75, Firefox 79, NodeJS v12.5, and Deno v0.4. ([#356](https://github.com/smol-dot/smoldot/pull/356)) - When the `state_getKeysPaged` JSON-RPC function is called, and the list of keys returned in the response is truncated (due to the `count` and `startKey` parameters), the rest of the keys are now put in a cache with the expectation that `state_getKeysPaged` is called again in order to obtain the rest of the keys. The `state_getKeysPaged` JSON-RPC function is unfortunately very often used by PolkadotJS despite being completely unsuitable for light clients. ([#361](https://github.com/smol-dot/smoldot/pull/361)) +- Significantly optimize the performance of the proof verification in `state_getKeys` and `state_getKeysPaged`. ([#363](https://github.com/smol-dot/smoldot/pull/363)) ### Fixed From ed2fa223d5a4a82ec991bf690ee0cbdad428151c Mon Sep 17 00:00:00 2001 From: Pierre Krieger Date: Tue, 28 Mar 2023 19:58:09 +0200 Subject: [PATCH 5/5] Rustfmt --- lib/src/trie/proof_decode.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lib/src/trie/proof_decode.rs b/lib/src/trie/proof_decode.rs index 0b2bb36f92..264a977d2b 100644 --- a/lib/src/trie/proof_decode.rs +++ b/lib/src/trie/proof_decode.rs @@ -799,7 +799,9 @@ impl<'a> Children<'a> { } /// Returns an iterator of 16 items, one for each child. - pub fn children(&'_ self) -> impl DoubleEndedIterator + ExactSizeIterator> + '_ { + pub fn children( + &'_ self, + ) -> impl DoubleEndedIterator + ExactSizeIterator> + '_ { self.children.iter().copied() } }