Rewrite proof decoding code to not use a separate map (#1462)

* Rewrite proof decoding code to not use a separate map * Work in progress * Adjust some APIs * Use strong typing for stack entries * Don't decode trie nodes twice when decoding * Work on next_key * `next_key` is expected to work now * Fix benches * Give me a break * Fix tests compilation * Uncomment now-working test * Implement trie_node_info * Reimplement `iter_ordered` * Don't require lifetime on `trie_root_merkle_value` * Update light client * Inline and fix `closest_ancestor_in_proof` * Remove empty proof special case * Improve fmt::Debug for EntryKeyIter * Bugfix in `next_key` * Uncomment failing test * CHANGELOG * Fix warnings * Typo * Fix children with nibble `e` being skipped * Stupidest typo * Add some TODOs * Small revert in runtime_service
smol-dot · Dec 21, 2023 · 3186bb4 · 3186bb4
1 parent 4f3fa6e
commit 3186bb4
Show file tree

Hide file tree

Showing 14 changed files with 1,156 additions and 699 deletions.
diff --git a/lib/benches/proof-decode.rs b/lib/benches/proof-decode.rs
@@ -15,6 +15,8 @@
 // You should have received a copy of the GNU General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+use std::iter;
+
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput};
 use rand::Rng as _;
 
@@ -95,7 +97,15 @@ fn benchmark_proof_decode(c: &mut Criterion) {
                             .map(|_| rand::thread_rng().gen_range(0..16).try_into().unwrap())
                             .collect::<Vec<_>>()
                     },
-                    |key| i.next_key(&proof.trie_root, &key, false, &[], true),
+                    |key| {
+                        i.next_key(
+                            &proof.trie_root,
+                            key.iter().copied(),
+                            false,
+                            iter::empty(),
+                            true,
+                        )
+                    },
                     BatchSize::SmallInput,
                 )
             },
@@ -114,7 +124,7 @@ fn benchmark_proof_decode(c: &mut Criterion) {
                             .map(|_| rand::thread_rng().gen_range(0..16).try_into().unwrap())
                             .collect::<Vec<_>>()
                     },
-                    |key| i.closest_descendant_merkle_value(&proof.trie_root, &key),
+                    |key| i.closest_descendant_merkle_value(&proof.trie_root, key.iter().copied()),
                     BatchSize::SmallInput,
                 )
             },
@@ -133,7 +143,7 @@ fn benchmark_proof_decode(c: &mut Criterion) {
                             .map(|_| rand::thread_rng().gen_range(0..16).try_into().unwrap())
                             .collect::<Vec<_>>()
                     },
-                    |key| i.closest_ancestor_in_proof(&proof.trie_root, &key),
+                    |key| i.closest_ancestor_in_proof(&proof.trie_root, key.iter().copied()),
                     BatchSize::SmallInput,
                 )
             },

diff --git a/lib/src/sync/warp_sync.rs b/lib/src/sync/warp_sync.rs
@@ -1715,19 +1715,23 @@ impl<TSrc, TRq> BuildRuntime<TSrc, TRq> {
             let code_nibbles = trie::bytes_to_nibbles(b":code".iter().copied()).collect::<Vec<_>>();
             match decoded_downloaded_runtime.closest_ancestor_in_proof(
                 &self.inner.warped_header_state_root,
-                &code_nibbles[..code_nibbles.len() - 1],
+                code_nibbles.iter().take(code_nibbles.len() - 1).copied(),
             ) {
                 Ok(Some(closest_ancestor_key)) => {
+                    let closest_ancestor_key = closest_ancestor_key.collect::<Vec<_>>();
                     let next_nibble = code_nibbles[closest_ancestor_key.len()];
                     let merkle_value = decoded_downloaded_runtime
-                        .trie_node_info(&self.inner.warped_header_state_root, closest_ancestor_key)
+                        .trie_node_info(
+                            &self.inner.warped_header_state_root,
+                            closest_ancestor_key.iter().copied(),
+                        )
                         .unwrap()
                         .children
                         .child(next_nibble)
                         .merkle_value();
 
                     match merkle_value {
-                        Some(mv) => (mv.to_owned(), closest_ancestor_key.to_vec()),
+                        Some(mv) => (mv.to_owned(), closest_ancestor_key),
                         None => {
                             self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                             self.inner.runtime_download = RuntimeDownload::NotStarted {
@@ -2034,9 +2038,9 @@ impl<TSrc, TRq> BuildChainInformation<TSrc, TRq> {
                     let (proof, downloaded_source) = calls.get(&nk.call_in_progress()).unwrap();
                     let value = match proof.next_key(
                         &self.inner.warped_header_state_root,
-                        &nk.key().collect::<Vec<_>>(), // TODO: overhead
+                        nk.key(),
                         nk.or_equal(),
-                        &nk.prefix().collect::<Vec<_>>(), // TODO: overhead
+                        nk.prefix(),
                         nk.branch_nodes(),
                     ) {
                         Ok(v) => v,
@@ -2052,14 +2056,14 @@ impl<TSrc, TRq> BuildChainInformation<TSrc, TRq> {
                             );
                         }
                     };
-                    nk.inject_key(value.map(|v| v.iter().copied()))
+                    nk.inject_key(value)
                 }
                 chain_information::build::InProgress::ClosestDescendantMerkleValue(mv) => {
                     // TODO: child tries not supported
                     let (proof, downloaded_source) = calls.get(&mv.call_in_progress()).unwrap();
                     let value = match proof.closest_descendant_merkle_value(
                         &self.inner.warped_header_state_root,
-                        &mv.key().collect::<Vec<_>>(), // TODO: overhead
+                        mv.key(),
                     ) {
                         Ok(v) => v,
                         Err(proof_decode::IncompleteProofError { .. }) => {

diff --git a/lib/src/transactions/validate/tests.rs b/lib/src/transactions/validate/tests.rs
@@ -67,17 +67,17 @@ fn validate_from_proof() {
                 let next_key = call_proof
                     .next_key(
                         main_trie_root,
-                        &nk.key().collect::<Vec<_>>(),
+                        nk.key(),
                         nk.or_equal(),
-                        &nk.prefix().collect::<Vec<_>>(),
+                        nk.prefix(),
                         nk.branch_nodes(),
                     )
                     .unwrap();
-                validation_in_progress = nk.inject_key(next_key.map(|k| k.iter().copied()));
+                validation_in_progress = nk.inject_key(next_key);
             }
             super::Query::ClosestDescendantMerkleValue(mv) => {
                 let value = call_proof
-                    .closest_descendant_merkle_value(main_trie_root, &mv.key().collect::<Vec<_>>())
+                    .closest_descendant_merkle_value(main_trie_root, mv.key())
                     .unwrap();
                 validation_in_progress = mv.inject_merkle_value(value);
             }

diff --git a/lib/src/trie/nibble.rs b/lib/src/trie/nibble.rs
@@ -76,6 +76,12 @@ impl From<Nibble> for u8 {
     }
 }
 
+impl From<Nibble> for usize {
+    fn from(nibble: Nibble) -> usize {
+        usize::from(nibble.0)
+    }
+}
+
 impl fmt::LowerHex for Nibble {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::LowerHex::fmt(&self.0, f)

diff --git 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::{borrow::ToOwned as _, vec, vec::Vec};
+use alloc::{vec, vec::Vec};
 use core::{fmt, iter, mem};
 
 mod tests;
@@ -169,7 +169,8 @@ impl PrefixScan {
                     };
 
                     match (
-                        decoded_proof.trie_node_info(&self.trie_root_hash, info_of_node),
+                        decoded_proof
+                            .trie_node_info(&self.trie_root_hash, info_of_node.iter().copied()),
                         query_ty,
                     ) {
                         (Ok(info), QueryTy::Exact) => info,
@@ -179,7 +180,7 @@ impl PrefixScan {
                                     // 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));
+                                    next.push((child_key.collect::<Vec<_>>(), QueryTy::Exact));
                                     continue;
                                 }
                                 proof_decode::Child::AbsentFromProof { .. } => {
@@ -258,7 +259,7 @@ impl PrefixScan {
                             next.push((direction, QueryTy::Direction));
                         }
                         proof_decode::Child::InProof { child_key, .. } => {
-                            next.push((child_key.to_owned(), QueryTy::Exact))
+                            next.push((child_key.collect::<Vec<_>>(), QueryTy::Exact))
                         }
                     }
                 }