Add child trie unit tests for runtime_host (#722)

* Add a unit test for runtime_host

* Remove TODO

* Use the block number from the test

* Update test

* Properly decode digest items

* Update to new format

* Add second test

* Don't return when first test succeeds

* Fix block number

* Add tests 3 and 4, now failing as expected

* Support child tries in the test

* Small code tweak

* Move serde structs

* Convert test to more simple format

* Remove `number_bytes` field as it's useless in the end

* CamelCase the field names

* Documentation and code comments tweaks

* Add TODO

* Add co-author for the test fixtures

* Language tweak

---------

Co-authored-by: pgherveou <[email protected]>

lib/src/executor/runtime_host.rs

@@ -47,6 +47,8 @@ use core::{fmt, iter};
 
 pub use trie::{Nibble, TrieEntryVersion};
 
+mod tests;
+
 /// Configuration for [`run`].
 pub struct Config<'a, TParams> {
     /// Virtual machine to be run.

lib/src/executor/runtime_host/tests.rs

@@ -0,0 +1,233 @@
+// Smoldot
+// Copyright (C) 2023  Pierre Krieger
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+#![cfg(test)]
+
+//! The test in this module reads various JSON files containing test fixtures and executes them.
+//!
+//! Each test fixture contains a block (header and body), plus the storage of its parent. The
+//! test consists in executing the block, to make sure that the state trie root matches the one
+//! calculated by smoldot.
+
+use core::{iter, ops};
+
+use super::{run, Config, RuntimeHostVm};
+use crate::{executor::host, trie};
+use alloc::collections::BTreeMap;
+
+#[test]
+fn execute_blocks() {
+    // Tests ordered alphabetically.
+    for (test_num, test_json) in [
+        include_str!("./child-trie-create-multiple.json"),
+        include_str!("./child-trie-create-one.json"),
+        include_str!("./child-trie-destroy.json"),
+        include_str!("./child-trie-read-basic.json"),
+        // TODO: more tests?
+    ]
+    .into_iter()
+    .enumerate()
+    {
+        // Decode the test JSON.
+        let test_data = serde_json::from_str::<Test>(test_json).unwrap();
+
+        // Turn the nice-looking data into something with better access times.
+        let storage = {
+            let mut storage = test_data
+                .parent_storage
+                .main_trie
+                .iter()
+                .map(|(key, value)| ((None, key.0.clone()), value.0.clone()))
+                .collect::<BTreeMap<_, _>>();
+            for (child_trie, child_trie_data) in &test_data.parent_storage.child_tries {
+                for (key, value) in child_trie_data {
+                    storage.insert((Some(child_trie.0.clone()), key.0.clone()), value.0.clone());
+                }
+            }
+            storage
+        };
+
+        // Build the runtime.
+        let virtual_machine = {
+            let code = storage
+                .get(&(None, b":code".to_vec()))
+                .expect("no runtime code found");
+            let heap_pages = crate::executor::storage_heap_pages_to_value(
+                storage.get(&(None, b":heappages".to_vec())).map(|v| &v[..]),
+            )
+            .unwrap();
+
+            host::HostVmPrototype::new(host::Config {
+                module: code,
+                heap_pages,
+                exec_hint: crate::executor::vm::ExecHint::Oneshot,
+                allow_unresolved_imports: false,
+            })
+            .unwrap()
+        };
+
+        // The runtime indicates the version of the trie items of the parent storage.
+        // While in principle each storage item could have a different version, in practice we
+        // just assume they're all the same.
+        let state_version = virtual_machine
+            .runtime_version()
+            .decode()
+            .state_version
+            .unwrap_or(host::TrieEntryVersion::V0);
+
+        // Start executing `Core_execute_block`. This runtime call will verify at the end whether
+        // the trie root hash of the block matches the one calculated by smoldot.
+        let mut execution = run(Config {
+            virtual_machine,
+            function_to_call: "Core_execute_block",
+            max_log_level: 3,
+            offchain_storage_changes: Default::default(),
+            storage_main_trie_changes: Default::default(),
+            parameter: {
+                // Block header + number of extrinsics + extrinsics
+                let encoded_body_len =
+                    crate::util::encode_scale_compact_usize(test_data.block.body.len());
+                iter::once(either::Right(either::Left(&test_data.block.header.0)))
+                    .chain(iter::once(either::Right(either::Right(encoded_body_len))))
+                    .chain(test_data.block.body.iter().map(|b| either::Left(&b.0)))
+            },
+        })
+        .unwrap();
+
+        loop {
+            match execution {
+                RuntimeHostVm::Finished(Ok(_)) => break, // Test successful!
+                RuntimeHostVm::Finished(Err(err)) => {
+                    panic!("Error during test #{}: {:?}", test_num, err)
+                }
+                RuntimeHostVm::SignatureVerification(sig) => execution = sig.verify_and_resume(),
+                RuntimeHostVm::ClosestDescendantMerkleValue(req) => {
+                    execution = req.resume_unknown()
+                }
+                RuntimeHostVm::StorageGet(get) => {
+                    let value = storage
+                        .get(&(
+                            get.child_trie().map(|c| c.as_ref().to_owned()),
+                            get.key().as_ref().to_owned(),
+                        ))
+                        .map(|v| (iter::once(&v[..]), state_version));
+                    execution = get.inject_value(value);
+                }
+                RuntimeHostVm::NextKey(req) => {
+                    // Because `NextKey` might ask for branch nodes, and that we don't build the
+                    // trie in its entirety, we have to use an algorithm that finds the branch
+                    // nodes for us.
+                    let next_key = {
+                        let mut search = trie::branch_search::BranchSearch::NextKey(
+                            trie::branch_search::start_branch_search(trie::branch_search::Config {
+                                key_before: req.key().collect::<Vec<_>>().into_iter(),
+                                or_equal: req.or_equal(),
+                                prefix: req.prefix().collect::<Vec<_>>().into_iter(),
+                                no_branch_search: !req.branch_nodes(),
+                            }),
+                        );
+
+                        loop {
+                            match search {
+                                trie::branch_search::BranchSearch::Found {
+                                    branch_trie_node_key,
+                                } => break branch_trie_node_key,
+                                trie::branch_search::BranchSearch::NextKey(bs_req) => {
+                                    let result = storage
+                                        .range((
+                                            if bs_req.or_equal() {
+                                                ops::Bound::Included((
+                                                    req.child_trie().map(|c| c.as_ref().to_owned()),
+                                                    bs_req.key_before().collect::<Vec<_>>(),
+                                                ))
+                                            } else {
+                                                ops::Bound::Excluded((
+                                                    req.child_trie().map(|c| c.as_ref().to_owned()),
+                                                    bs_req.key_before().collect::<Vec<_>>(),
+                                                ))
+                                            },
+                                            ops::Bound::Unbounded,
+                                        ))
+                                        .next()
+                                        .filter(|((trie, key), _)| {
+                                            *trie == req.child_trie().map(|c| c.as_ref().to_owned())
+                                                && key.starts_with(
+                                                    &bs_req.prefix().collect::<Vec<_>>(),
+                                                )
+                                        })
+                                        .map(|((_, k), _)| k);
+
+                                    search = bs_req.inject(result.map(|k| k.iter().copied()));
+                                }
+                            }
+                        }
+                    };
+
+                    execution = req.inject_key(next_key.map(|nk| nk.into_iter()));
+                }
+            }
+        }
+    }
+}
+
+// Serde structs used to decode the test fixtures.
+
+#[derive(serde::Deserialize)]
+struct Test {
+    block: Block,
+    #[serde(rename = "parentStorage")]
+    parent_storage: Storage,
+}
+
+#[derive(serde::Deserialize)]
+struct Block {
+    header: HexString,
+    body: Vec<HexString>,
+}
+
+#[derive(serde::Deserialize)]
+struct Storage {
+    #[serde(rename = "mainTrie")]
+    main_trie: hashbrown::HashMap<HexString, HexString>,
+    #[serde(rename = "childTries")]
+    child_tries: hashbrown::HashMap<HexString, hashbrown::HashMap<HexString, HexString>>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+struct HexString(Vec<u8>);
+
+impl<'a> serde::Deserialize<'a> for HexString {
+    fn deserialize<D>(deserializer: D) -> Result<HexString, D::Error>
+    where
+        D: serde::Deserializer<'a>,
+    {
+        let string = String::deserialize(deserializer)?;
+
+        if string.is_empty() {
+            return Ok(HexString(Vec::new()));
+        }
+
+        if !string.starts_with("0x") {
+            return Err(serde::de::Error::custom(
+                "hexadecimal string doesn't start with 0x",
+            ));
+        }
+
+        let bytes = hex::decode(&string[2..]).map_err(serde::de::Error::custom)?;
+        Ok(HexString(bytes))
+    }
+}

lib/src/header.rs

@@ -1080,6 +1080,15 @@ impl<'a> DigestItemRef<'a> {
         matches!(self, DigestItemRef::GrandpaConsensus(_))
     }
 
+    /// Decodes a SCALE-encoded digest item.
+    pub fn from_scale_encoded(bytes: &'a [u8], block_number_bytes: usize) -> Result<Self, Error> {
+        let (item, remain) = decode_item(bytes, block_number_bytes)?;
+        if !remain.is_empty() {
+            return Err(Error::TooLong);
+        }
+        Ok(item)
+    }
+
     /// Returns an iterator to list of buffers which, when concatenated, produces the SCALE
     /// encoding of that digest item.
     pub fn scale_encoding(