Modify runtime_host.rs to support child tries (#684)

* Prepare for addition of child tries diffs

* Add support for child tries in runtime_host.rs

* Add support in read_only_runtime_host.rs

* Add trie() methods to everything relevant

* Intercept child trie requests in full node

* Replace trie() -> Trie with child_trie() -> Option

* Plug in child trie reading from proof

* Be more relaxed in inject_merkle_value, actually

* Docfix

* Rustfmt

* CHANGELOG

* Update full node for API changes

full-node/src/run/consensus_service.rs

@@ -824,6 +824,14 @@ impl SyncBackground {
 
                     // Access to the best block storage.
                     author::build::BuilderAuthoring::StorageGet(req) => {
+                        // TODO: child tries not supported
+                        if req.child_trie().is_some() {
+                            log::warn!("child-tries-not-supported");
+                            block_authoring =
+                                req.inject_value(None::<(iter::Empty<&'static [u8]>, _)>);
+                            continue;
+                        }
+
                         let key = req.key().as_ref().to_vec();
                         let value = self
                             .database
@@ -844,6 +852,13 @@ impl SyncBackground {
                         block_authoring = req.resume_unknown();
                     }
                     author::build::BuilderAuthoring::NextKey(req) => {
+                        // TODO: child tries not supported
+                        if req.child_trie().is_some() {
+                            log::warn!("child-tries-not-supported");
+                            block_authoring = req.inject_key(None::<iter::Empty<_>>);
+                            continue;
+                        }
+
                         let search_params = trie::branch_search::Config {
                             key_before: req.key().collect::<Vec<_>>().into_iter(),
                             or_equal: req.or_equal(),

lib/src/author/build.rs

@@ -317,6 +317,11 @@ impl StorageGet {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// Injects the corresponding storage value.
     pub fn inject_value(
         self,
@@ -338,6 +343,11 @@ impl ClosestDescendantMerkleValue {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// Indicate that the value is unknown and resume the calculation.
     ///
     /// This function be used if you are unaware of the Merkle value. The algorithm will perform
@@ -347,7 +357,10 @@ impl ClosestDescendantMerkleValue {
     }
 
     /// Injects the corresponding Merkle value.
-    pub fn inject_merkle_value(self, merkle_value: &[u8]) -> BuilderAuthoring {
+    ///
+    /// `None` can be passed if there is no descendant or, in the case of a child trie read, in
+    /// order to indicate that the child trie does not exist.
+    pub fn inject_merkle_value(self, merkle_value: Option<&[u8]>) -> BuilderAuthoring {
         self.1
             .with_runtime_inner(self.0.inject_merkle_value(merkle_value))
     }
@@ -364,6 +377,11 @@ impl NextKey {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// If `true`, then the provided value must the one superior or equal to the requested key.
     /// If `false`, then the provided value must be strictly superior to the requested key.
     pub fn or_equal(&self) -> bool {

lib/src/author/runtime.rs

@@ -620,6 +620,11 @@ impl StorageGet {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// Injects the corresponding storage value.
     pub fn inject_value(
         self,
@@ -641,6 +646,11 @@ impl ClosestDescendantMerkleValue {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// Indicate that the value is unknown and resume the calculation.
     ///
     /// This function be used if you are unaware of the Merkle value. The algorithm will perform
@@ -650,7 +660,10 @@ impl ClosestDescendantMerkleValue {
     }
 
     /// Injects the corresponding Merkle value.
-    pub fn inject_merkle_value(self, merkle_value: &[u8]) -> BlockBuild {
+    ///
+    /// `None` can be passed if there is no descendant or, in the case of a child trie read, in
+    /// order to indicate that the child trie does not exist.
+    pub fn inject_merkle_value(self, merkle_value: Option<&[u8]>) -> BlockBuild {
         BlockBuild::from_inner(self.0.inject_merkle_value(merkle_value), self.1)
     }
 }
@@ -666,6 +679,11 @@ impl NextKey {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// If `true`, then the provided value must the one superior or equal to the requested key.
     /// If `false`, then the provided value must be strictly superior to the requested key.
     pub fn or_equal(&self) -> bool {

lib/src/chain/blocks_tree/verify.rs

@@ -924,6 +924,11 @@ impl<T> StorageGet<T> {
         self.inner.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.inner.child_trie()
+    }
+
     /// Access to the Nth ancestor's information and hierarchy. Returns `None` if `n` is too
     /// large. A value of `0` for `n` corresponds to the parent block. A value of `1` corresponds
     /// to the parent's parent. And so on.
@@ -984,6 +989,11 @@ impl<T> StorageClosestDescendantMerkleValue<T> {
         self.inner.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.inner.child_trie()
+    }
+
     /// Indicate that the value is unknown and resume the calculation.
     ///
     /// This function be used if you are unaware of the Merkle value. The algorithm will perform
@@ -994,7 +1004,10 @@ impl<T> StorageClosestDescendantMerkleValue<T> {
     }
 
     /// Injects the corresponding Merkle value.
-    pub fn inject_merkle_value(self, merkle_value: &[u8]) -> BodyVerifyStep2<T> {
+    ///
+    /// `None` can be passed if there is no descendant or, in the case of a child trie read, in
+    /// order to indicate that the child trie does not exist.
+    pub fn inject_merkle_value(self, merkle_value: Option<&[u8]>) -> BodyVerifyStep2<T> {
         let inner = self.inner.inject_merkle_value(merkle_value);
         self.context.with_body_verify(inner)
     }
@@ -1013,6 +1026,11 @@ impl<T> StorageNextKey<T> {
         self.inner.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.inner.child_trie()
+    }
+
     /// If `true`, then the provided value must the one superior or equal to the requested key.
     /// If `false`, then the provided value must be strictly superior to the requested key.
     pub fn or_equal(&self) -> bool {

lib/src/chain/chain_information/build.rs

@@ -280,6 +280,11 @@ impl StorageGet {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// Injects the corresponding storage value.
     pub fn inject_value(
         self,
@@ -304,6 +309,11 @@ impl NextKey {
         self.0.key()
     }
 
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
+        self.0.child_trie()
+    }
+
     /// If `true`, then the provided value must the one superior or equal to the requested key.
     /// If `false`, then the provided value must be strictly superior to the requested key.
     pub fn or_equal(&self) -> bool {

lib/src/executor/host.rs

@@ -2157,17 +2157,17 @@ impl ExternalStorageGet {
             .unwrap()
     }
 
-    /// Returns the trie that must be read from.
-    pub fn trie(&'_ self) -> Trie<impl AsRef<[u8]> + '_> {
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
         if let Some((child_trie_ptr, child_trie_size)) = self.child_trie_ptr_size {
             let child_trie = self
                 .inner
                 .vm
                 .read_memory(child_trie_ptr, child_trie_size)
                 .unwrap();
-            Trie::ChildTrieDefault { child_trie }
+            Some(child_trie)
         } else {
-            Trie::MainTrie
+            None
         }
     }
 
@@ -2384,21 +2384,21 @@ impl ExternalStorageSet {
         }
     }
 
-    /// Returns the trie that must be written to.
-    pub fn trie(&'_ self) -> Trie<impl AsRef<[u8]> + '_> {
+    /// If `Some`, write to the given child trie. If `None`, write to the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
         match &self.write_inner {
             ExternalStorageSetInner::Regular {
                 child_trie_ptr_size: Some((ptr, size)),
                 ..
             } => {
                 let child_trie = self.inner.vm.read_memory(*ptr, *size).unwrap();
-                Trie::ChildTrieDefault { child_trie }
+                Some(child_trie)
             }
             ExternalStorageSetInner::Regular {
                 child_trie_ptr_size: None,
                 ..
-            } => Trie::MainTrie,
-            ExternalStorageSetInner::ChildTrieRootCommit { .. } => Trie::MainTrie,
+            } => None,
+            ExternalStorageSetInner::ChildTrieRootCommit { .. } => None,
         }
     }
 
@@ -2497,13 +2497,13 @@ impl ExternalStorageAppend {
             .unwrap()
     }
 
-    /// Returns the trie that must be written to.
+    /// If `Some`, write to the given child trie. If `None`, write to the main trie.
     ///
-    /// > **Note**: At the moment, this function always returns [`Trie::MainTrie`], as there is
-    /// >           no host function that appends to a child trie storage.
-    pub fn trie(&'_ self) -> Trie<impl AsRef<[u8]> + '_> {
+    /// > **Note**: At the moment, this function always returns None, as there is no host function
+    /// >           that appends to a child trie storage.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
         // Note that there is no equivalent of this host function for child tries.
-        Trie::<&'static [u8]>::MainTrie
+        None::<&'static [u8]>
     }
 
     /// Returns the value to append.
@@ -2558,17 +2558,17 @@ impl ExternalStorageClearPrefix {
         }
     }
 
-    /// Returns the trie that must be written to.
-    pub fn trie(&'_ self) -> Trie<impl AsRef<[u8]> + '_> {
+    /// If `Some`, write to the given child trie. If `None`, write to the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
         if let Some((child_trie_ptr, child_trie_size)) = self.child_trie_ptr_size {
             let child_trie = self
                 .inner
                 .vm
                 .read_memory(child_trie_ptr, child_trie_size)
                 .unwrap();
-            Trie::ChildTrieDefault { child_trie }
+            Some(child_trie)
         } else {
-            Trie::MainTrie
+            None
         }
     }
 
@@ -2637,23 +2637,23 @@ pub struct ExternalStorageRoot {
 
 impl ExternalStorageRoot {
     /// Returns the trie whose root hash must be provided.
-    pub fn trie(&'_ self) -> Trie<impl AsRef<[u8]> + '_> {
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
         if let Some((ptr, size)) = self.child_trie_ptr_size {
             let child_trie = self.inner.vm.read_memory(ptr, size).unwrap();
-            Trie::ChildTrieDefault { child_trie }
+            Some(child_trie)
         } else {
-            Trie::MainTrie
+            None
         }
     }
 
     /// Writes the trie root hash to the Wasm VM and prepares it for resume.
     ///
-    /// Must be passed `None` if [`ExternalStorageRoot::trie`] returned [`Trie::ChildTrieDefault`]
-    /// and the trie doesn't exist.
+    /// Must be passed `None` if [`ExternalStorageRoot::child_trie`] returned `Some` and the trie
+    /// doesn't exist.
     ///
     /// # Panic
     ///
-    /// Panics if `None` is passed and [`ExternalStorageRoot::trie`] returned [`Trie::MainTrie`].
+    /// Panics if `None` is passed and [`ExternalStorageRoot::child_trie`] returned `None`.
     ///
     pub fn resume(self, hash: Option<&[u8; 32]>) -> HostVm {
         let host_fn = match self.inner.registered_functions[self.calling] {
@@ -2690,12 +2690,6 @@ impl fmt::Debug for ExternalStorageRoot {
     }
 }
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
-pub enum Trie<T> {
-    MainTrie,
-    ChildTrieDefault { child_trie: T },
-}
-
 /// Must provide the storage key that follows, in lexicographic order, a specific one.
 pub struct ExternalStorageNextKey {
     inner: Inner,
@@ -2721,17 +2715,17 @@ impl ExternalStorageNextKey {
             .unwrap()
     }
 
-    /// Returns the trie that must be read from.
-    pub fn trie(&'_ self) -> Trie<impl AsRef<[u8]> + '_> {
+    /// If `Some`, read from the given child trie. If `None`, read from the main trie.
+    pub fn child_trie(&'_ self) -> Option<impl AsRef<[u8]> + '_> {
         if let Some((child_trie_ptr, child_trie_size)) = self.child_trie_ptr_size {
             let child_trie = self
                 .inner
                 .vm
                 .read_memory(child_trie_ptr, child_trie_size)
                 .unwrap();
-            Trie::ChildTrieDefault { child_trie }
+            Some(child_trie)
         } else {
-            Trie::MainTrie
+            None
         }
     }