Optimize prefix_proof by reducing the number of iterations

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<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`].
@@ -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::<Vec<_>>();
@@ -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.

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:?}"),

lib/src/trie/proof_decode.rs

@@ -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,
                         },
                     },
@@ -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.
@@ -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 => {
@@ -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) => {
@@ -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
@@ -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.
@@ -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(())
     }
 }