Finish implementing state_version = 1 (#230)

* Implement ext_storage_root_version_2

* Add support for trie v1 in RuntimeHostVm

* Update chain_spec.rs after API changes

* Propagate changes to the various runtime state machines

* Add a generic parameter to StorageDiff

* Fix template parameters on StorageDiff

* Add StorageDiff::merge_map

* Update optimistic.rs

* Fix merge_map

* Update all.rs

* Give information about the trie node version in Merkle proofs

* Only provide a node version when a value is set

* Update light-base

* Fix tests

* Add BlockFull::state_trie_version

* Store the trie entry version in the storage cache

* Insert the trie entry version in the database

* WIP full node update

* Make the database return the trie entry version

* Finish update to full node

* Fix invalid write to non_finalized_changes

* Minor fixes

* Also implement ext_default_child_storage_root_version_2

* Fix doctest

* Fix other doctest

* Add comment about versions in trie.rs

* Spellcheck

full-node/src/run.rs

@@ -21,7 +21,7 @@ use futures::{channel::oneshot, prelude::*};
 use smoldot::{
     chain, chain_spec,
     database::full_sqlite,
-    header,
+    executor, header,
     identity::keystore,
     informant::HashDisplay,
     libp2p::{
@@ -742,18 +742,36 @@ async fn open_database(
 
         // The database doesn't exist or is empty.
         full_sqlite::DatabaseOpen::Empty(empty) => {
+            let genesis_storage = chain_spec.genesis_storage().into_genesis_items().unwrap(); // TODO: return error instead
+
+            // In order to determine the state_version of the genesis block, we need to compile
+            // the runtime.
+            // TODO: return errors instead of panicking
+            let state_version = executor::host::HostVmPrototype::new(executor::host::Config {
+                module: genesis_storage.value(b":code").unwrap(),
+                heap_pages: executor::storage_heap_pages_to_value(
+                    genesis_storage.value(b":heappages"),
+                )
+                .unwrap(),
+                exec_hint: executor::vm::ExecHint::Oneshot,
+                allow_unresolved_imports: true,
+            })
+            .unwrap()
+            .runtime_version()
+            .decode()
+            .state_version
+            .map(|v| u8::from(v))
+            .unwrap_or(0);
+
             // The finalized block is the genesis block. As such, it has an empty body and
             // no justification.
             let database = empty
                 .initialize(
                     genesis_chain_information,
                     iter::empty(),
                     None,
-                    chain_spec
-                        .genesis_storage()
-                        .into_genesis_items()
-                        .unwrap() // TODO: return error instead
-                        .iter(),
+                    genesis_storage.iter(),
+                    state_version,
                 )
                 .unwrap();
             (database, false)

full-node/src/run/consensus_service.rs

@@ -38,7 +38,7 @@ use smoldot::{
     informant::HashDisplay,
     libp2p,
     network::{self, protocol::BlockData},
-    sync::all,
+    sync::all::{self, TrieEntryVersion},
 };
 use std::{
     collections::BTreeMap,
@@ -128,7 +128,14 @@ impl ConsensusService {
             best_block_number,
             finalized_block_storage,
             finalized_chain_information,
-        ): (_, _, _, _, BTreeMap<Vec<u8>, Vec<u8>>, _) = config
+        ): (
+            _,
+            _,
+            _,
+            _,
+            BTreeMap<Vec<u8>, (Vec<u8>, TrieEntryVersion)>,
+            _,
+        ) = config
             .database
             .with_database({
                 let block_number_bytes = config.block_number_bytes;
@@ -153,9 +160,17 @@ impl ConsensusService {
                     )
                     .unwrap()
                     .number;
-                    let finalized_block_storage = database
+                    let finalized_block_storage: Vec<(Vec<u8>, Vec<u8>, u8)> = database
                         .finalized_block_storage_top_trie(&finalized_block_hash)
                         .unwrap();
+                    // TODO: we copy all entries; it could be more optimal to have a custom implementation of FromIterator that directly does the conversion?
+                    let finalized_block_storage = finalized_block_storage
+                        .into_iter()
+                        .map(|(k, val, vers)| {
+                            let vers = TrieEntryVersion::try_from(vers).unwrap(); // TODO: don't unwrap
+                            (k, (val, vers))
+                        })
+                        .collect();
                     let finalized_chain_information = database
                         .to_chain_information(&finalized_block_hash)
                         .unwrap();
@@ -223,11 +238,11 @@ impl ConsensusService {
                         // Builds the runtime of the finalized block.
                         // Assumed to always be valid, otherwise the block wouldn't have been saved in the
                         // database, hence the large number of unwraps here.
-                        let module = finalized_block_storage.get(&b":code"[..]).unwrap();
+                        let (module, _) = finalized_block_storage.get(&b":code"[..]).unwrap();
                         let heap_pages = executor::storage_heap_pages_to_value(
                             finalized_block_storage
                                 .get(&b":heappages"[..])
-                                .map(|v| &v[..]),
+                                .map(|(v, _)| &v[..]),
                         )
                         .unwrap();
                         executor::host::HostVmPrototype::new(executor::host::Config {
@@ -329,7 +344,7 @@ struct SyncBackground {
     // While reading the storage from the database is an option, doing so considerably slows down
     /// the verification, and also makes it impossible to insert blocks in the database in
     /// parallel of this verification.
-    finalized_block_storage: BTreeMap<Vec<u8>, Vec<u8>>,
+    finalized_block_storage: BTreeMap<Vec<u8>, (Vec<u8>, TrieEntryVersion)>,
 
     sync_state: Arc<Mutex<SyncState>>,
 
@@ -754,10 +769,11 @@ impl SyncBackground {
                             best_block_storage_access.get(key.as_ref(), || {
                                 self.finalized_block_storage
                                     .get(key.as_ref())
-                                    .map(|v| &v[..])
+                                    .map(|(val, vers)| (&val[..], *vers))
                             })
                         };
-                        block_authoring = get.inject_value(value.map(iter::once));
+                        block_authoring =
+                            get.inject_value(value.map(|(val, vers)| (iter::once(val), vers)));
                         continue;
                     }
                     author::build::BuilderAuthoring::NextKey(_) => {
@@ -1137,7 +1153,7 @@ impl SyncBackground {
                                 let value = self
                                     .finalized_block_storage
                                     .get(req.key().as_ref())
-                                    .map(|v| &v[..]);
+                                    .map(|(val, vers)| (&val[..], *vers));
                                 verify = req.inject_value(value);
                             }
                             all::BlockVerification::FinalizedStorageNextKey(req) => {
@@ -1212,16 +1228,21 @@ impl SyncBackground {
 
                             // TODO: maybe write in a separate task? but then we can't access the finalized storage immediately after?
                             for block in &finalized_blocks {
-                                for (key, value) in block
+                                for (key, value, ()) in block
                                     .full
                                     .as_ref()
                                     .unwrap()
                                     .storage_top_trie_changes
                                     .diff_iter_unordered()
                                 {
                                     if let Some(value) = value {
-                                        self.finalized_block_storage
-                                            .insert(key.to_owned(), value.to_owned());
+                                        self.finalized_block_storage.insert(
+                                            key.to_owned(),
+                                            (
+                                                value.to_owned(),
+                                                block.full.as_ref().unwrap().state_trie_version,
+                                            ),
+                                        );
                                     } else {
                                         let _was_there = self.finalized_block_storage.remove(key);
                                         // TODO: if a block inserts a new value, then removes it in the next block, the key will remain in `finalized_block_storage`; either solve this or document this
@@ -1327,7 +1348,9 @@ async fn database_blocks(
                         .as_ref()
                         .unwrap()
                         .storage_top_trie_changes
-                        .diff_iter_unordered(),
+                        .diff_iter_unordered()
+                        .map(|(k, v, ())| (k, v)),
+                    u8::from(block.full.as_ref().unwrap().state_trie_version),
                 );
 
                 match result {

lib/src/author/build.rs

@@ -28,6 +28,8 @@ use crate::{
 use alloc::vec::Vec;
 use core::{num::NonZeroU64, time::Duration};
 
+pub use runtime::TrieEntryVersion;
+
 /// Configuration for a block generation.
 pub struct Config<'a, TLocAuth> {
     /// Consensus-specific configuration.
@@ -310,7 +312,7 @@ impl StorageGet {
     /// Injects the corresponding storage value.
     pub fn inject_value(
         self,
-        value: Option<impl Iterator<Item = impl AsRef<[u8]>>>,
+        value: Option<(impl Iterator<Item = impl AsRef<[u8]>>, TrieEntryVersion)>,
     ) -> BuilderAuthoring {
         self.1.with_runtime_inner(self.0.inject_value(value))
     }

lib/src/author/runtime.rs

@@ -58,6 +58,8 @@ use crate::{
 use alloc::{borrow::ToOwned as _, string::String, vec::Vec};
 use core::{iter, mem};
 
+pub use runtime_host::TrieEntryVersion;
+
 /// Configuration for a block generation.
 pub struct Config<'a> {
     /// Number of bytes used to encode block numbers in the header.
@@ -110,6 +112,9 @@ pub struct Success {
     pub parent_runtime: host::HostVmPrototype,
     /// List of changes to the storage top trie that the block performs.
     pub storage_top_trie_changes: storage_diff::StorageDiff,
+    /// State trie version indicated by the runtime. All the storage changes indicated by
+    /// [`Success::storage_top_trie_changes`] should store this version alongside with them.
+    pub state_trie_version: TrieEntryVersion,
     /// List of changes to the off-chain storage that this block performs.
     pub offchain_storage_changes: storage_diff::StorageDiff,
     /// Cache used for calculating the top trie root of the new block.
@@ -428,6 +433,7 @@ impl BlockBuild {
                         body: shared.block_body,
                         parent_runtime: success.virtual_machine.into_prototype(),
                         storage_top_trie_changes: success.storage_top_trie_changes,
+                        state_trie_version: success.state_trie_version,
                         offchain_storage_changes: success.offchain_storage_changes,
                         top_trie_root_calculation_cache: success.top_trie_root_calculation_cache,
                         logs: shared.logs,
@@ -595,7 +601,10 @@ impl StorageGet {
     }
 
     /// Injects the corresponding storage value.
-    pub fn inject_value(self, value: Option<impl Iterator<Item = impl AsRef<[u8]>>>) -> BlockBuild {
+    pub fn inject_value(
+        self,
+        value: Option<(impl Iterator<Item = impl AsRef<[u8]>>, TrieEntryVersion)>,
+    ) -> BlockBuild {
         BlockBuild::from_inner(self.0.inject_value(value), self.1)
     }
 }

lib/src/author/runtime/tests.rs

@@ -62,7 +62,7 @@ fn block_building_works() {
                     .iter()
                     .find(|(k, _)| *k == get.key().as_ref())
                     .map(|(_, v)| iter::once(v));
-                builder = get.inject_value(value);
+                builder = get.inject_value(value.map(|v| (v, super::TrieEntryVersion::V0)));
             }
             super::BlockBuild::NextKey(_) => unimplemented!(), // Not needed for this test.
             super::BlockBuild::PrefixKeys(prefix) => {

lib/src/chain/blocks_tree/verify.rs

@@ -36,6 +36,8 @@ use super::{
 use alloc::boxed::Box;
 use core::cmp::Ordering;
 
+pub use verify::header_body::TrieEntryVersion;
+
 impl<T> NonFinalizedTree<T> {
     /// Verifies the given block.
     ///
@@ -600,6 +602,7 @@ impl<T> VerifyContext<T> {
                     parent_runtime: success.parent_runtime,
                     new_runtime: success.new_runtime,
                     storage_top_trie_changes: success.storage_top_trie_changes,
+                    state_trie_version: success.state_trie_version,
                     offchain_storage_changes: success.offchain_storage_changes,
                     top_trie_root_calculation_cache: success.top_trie_root_calculation_cache,
                     insert: BodyInsert {
@@ -845,6 +848,10 @@ pub enum BodyVerifyStep2<T> {
         new_runtime: Option<host::HostVmPrototype>,
         /// List of changes to the storage top trie that the block performs.
         storage_top_trie_changes: storage_diff::StorageDiff,
+        /// State trie version indicated by the runtime. All the storage changes indicated by
+        /// [`BodyVerifyStep2::Finished::storage_top_trie_changes`] should store this version
+        /// alongside with them.
+        state_trie_version: TrieEntryVersion,
         /// List of changes to the off-chain storage that this block performs.
         offchain_storage_changes: storage_diff::StorageDiff,
         /// Cache of calculation for the storage trie of the best block.
@@ -940,7 +947,7 @@ impl<T> StorageGet<T> {
     /// Injects the corresponding storage value.
     pub fn inject_value(
         self,
-        value: Option<impl Iterator<Item = impl AsRef<[u8]>>>,
+        value: Option<(impl Iterator<Item = impl AsRef<[u8]>>, TrieEntryVersion)>,
     ) -> BodyVerifyStep2<T> {
         let inner = self.inner.inject_value(value);
         self.context.with_body_verify(inner)

lib/src/chain_spec.rs

@@ -116,12 +116,11 @@ impl ChainSpec {
         let mut chain_information_build = build::ChainInformationBuild::new(build::Config {
             finalized_block_header: build::ConfigFinalizedBlockHeader::Genesis {
                 state_trie_root_hash: {
-                    let state_version = match vm_prototype.runtime_version().decode().state_version
-                    {
-                        Some(0) | None => trie::TrieEntryVersion::V0,
-                        Some(1) => trie::TrieEntryVersion::V1,
-                        Some(_) => return Err(FromGenesisStorageError::UnknownStateVersion),
-                    };
+                    let state_version = vm_prototype
+                        .runtime_version()
+                        .decode()
+                        .state_version
+                        .unwrap_or(trie::TrieEntryVersion::V0);
 
                     match self.genesis_storage() {
                         GenesisStorage::TrieRootHash(hash) => *hash,
@@ -144,7 +143,7 @@ impl ChainSpec {
                                 ) => {
                                     let key: alloc::vec::Vec<u8> = val.key().collect();
                                     let value = genesis_storage.value(&key[..]);
-                                    calculation = val.inject(state_version, value);
+                                    calculation = val.inject(value.map(move |v| (v, state_version)));
                                 }
                             }
                             }
@@ -496,8 +495,6 @@ pub enum FromGenesisStorageError {
     /// Error when initializing the virtual machine.
     #[display(fmt = "Error when initializing the virtual machine: {}", _0)]
     VmInitialization(executor::host::NewErr),
-    /// State version in runtime specification is not supported.
-    UnknownStateVersion,
     /// Chain specification doesn't contain the list of storage items.
     UnknownStorageItems,
 }