Implement trie version 1 (#2277)

* Implement trie version 1

* Finish update to proof verification code

* CHANGELOG

* Fix doclink

* Rustfmt

* Fix doctest

* Move code around properly

* Link to spec

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

bin/light-base/src/lib.rs

@@ -348,10 +348,11 @@ impl<TChain, TPlat: Platform> Client<TChain, TPlat> {
         // Load the information about the chain from the chain spec. If a light sync state (also
         // known as a checkpoint) is present in the chain spec, it is possible to start syncing at
         // the finalized block it describes.
-        let chain_information = {
+        // TODO: clean up that block
+        let (chain_information, genesis_block_header) = {
             match (
                 chain_spec
-                    .as_chain_information() // TODO: very expensive, don't always call?
+                    .as_chain_information()
                     .map(|(ci, _)| chain::chain_information::ValidChainInformation::try_from(ci)), // TODO: don't just throw away the runtime
                 chain_spec.light_sync_state().map(|s| {
                     chain::chain_information::ValidChainInformation::try_from(
@@ -360,7 +361,22 @@ impl<TChain, TPlat: Platform> Client<TChain, TPlat> {
                 }),
                 finalized_serialize::decode_chain(config.database_content),
             ) {
-                (_, _, Ok((ci, _))) => ci,
+                (Ok(Ok(genesis_ci)), _, Ok((ci, _))) => {
+                    let genesis_header = genesis_ci.as_ref().finalized_block_header.clone();
+                    (ci, genesis_header.into())
+                }
+
+                (Err(chain_spec::FromGenesisStorageError::UnknownStorageItems), _, Ok((ci, _))) => {
+                    let genesis_header = header::Header {
+                        parent_hash: [0; 32],
+                        number: 0,
+                        state_root: *chain_spec.genesis_storage().into_trie_root_hash().unwrap(),
+                        extrinsics_root: smoldot::trie::empty_trie_merkle_value(),
+                        digest: header::DigestRef::empty().into(),
+                    };
+
+                    (ci, genesis_header)
+                }
 
                 (Err(chain_spec::FromGenesisStorageError::UnknownStorageItems), None, _) => {
                     // TODO: we can in theory support chain specs that have neither a checkpoint nor the genesis storage, but it's complicated
@@ -376,7 +392,17 @@ impl<TChain, TPlat: Platform> Client<TChain, TPlat> {
                     Err(chain_spec::FromGenesisStorageError::UnknownStorageItems),
                     Some(Ok(ci)),
                     _,
-                ) => ci,
+                ) => {
+                    let genesis_header = header::Header {
+                        parent_hash: [0; 32],
+                        number: 0,
+                        state_root: *chain_spec.genesis_storage().into_trie_root_hash().unwrap(),
+                        extrinsics_root: smoldot::trie::empty_trie_merkle_value(),
+                        digest: header::DigestRef::empty().into(),
+                    };
+
+                    (ci, genesis_header)
+                }
 
                 (Err(err), _, _) => {
                     return ChainId(self.public_api_chains.insert(PublicApiChain::Erroneous {
@@ -399,16 +425,19 @@ impl<TChain, TPlat: Platform> Client<TChain, TPlat> {
                     }));
                 }
 
-                (_, Some(Ok(ci)), _) => ci,
+                (Ok(Ok(genesis_ci)), Some(Ok(ci)), _) => {
+                    let genesis_header = genesis_ci.as_ref().finalized_block_header.clone();
+                    (ci, genesis_header.into())
+                }
 
-                (Ok(Ok(ci)), None, _) => ci,
+                (Ok(Ok(ci)), None, _) => {
+                    let genesis_header =
+                        header::Header::from(ci.as_ref().finalized_block_header.clone());
+                    (ci, genesis_header)
+                }
             }
         };
 
-        // Even with a checkpoint, knowing the genesis block header is necessary for various
-        // reasons.
-        let genesis_block_header = smoldot::calculate_genesis_block_header(&chain_spec);
-
         // If the chain specification specifies a parachain, find the corresponding relay chain
         // in the list of potential relay chains passed by the user.
         // If no relay chain can be found, the chain creation fails. Exactly one matching relay

bin/light-base/src/runtime_service.rs

@@ -822,7 +822,11 @@ impl<'a> RuntimeCallLock<'a> {
             })
             .map_err(RuntimeCallError::StorageRetrieval)?;
 
-            if node_info.storage_value.is_some() {
+            if matches!(
+                node_info.storage_value,
+                proof_verify::StorageValue::Known(_)
+                    | proof_verify::StorageValue::HashKnownValueMissing(_)
+            ) {
                 assert_eq!(key.len() % 2, 0);
                 output.push(trie::nibbles_to_bytes_extend(key.iter().copied()).collect::<Vec<_>>());
             }

bin/wasm-node/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Added
+
+- Added support for version 1 of the trie. Previously, it wasn't possible to connect to chains that were using version 1. ([#2277](https://github.com/paritytech/smoldot/pull/2277))
+
 ### Changed
 
 - The runtime of the genesis block is now only compiled once when a chain is added, decreasing the time this operation takes. ([#2270](https://github.com/paritytech/smoldot/pull/2270))

src/author/runtime/tests.rs

@@ -28,26 +28,12 @@ fn block_building_works() {
     .unwrap();
     let genesis_storage = chain_specs.genesis_storage().into_genesis_items().unwrap();
 
-    let parent_runtime = {
-        let code = genesis_storage
-            .iter()
-            .find(|(k, _)| k == b":code")
-            .unwrap()
-            .1;
-        crate::executor::host::HostVmPrototype::new(crate::executor::host::Config {
-            module: code,
-            heap_pages: crate::executor::DEFAULT_HEAP_PAGES,
-            exec_hint: crate::executor::vm::ExecHint::Oneshot,
-            allow_unresolved_imports: false,
-        })
-        .unwrap()
-    };
-
-    let parent_hash = crate::calculate_genesis_block_header(&chain_specs).hash();
+    let (chain_info, genesis_runtime) = chain_specs.as_chain_information().unwrap();
+    let genesis_hash = chain_info.finalized_block_header.hash();
 
     let mut builder = super::build_block(super::Config {
-        parent_runtime,
-        parent_hash: &parent_hash,
+        parent_runtime: genesis_runtime,
+        parent_hash: &genesis_hash,
         parent_number: 0,
         block_body_capacity: 0,
         consensus_digest_log_item: super::ConfigPreRuntime::Aura(crate::header::AuraPreDigest {
@@ -61,7 +47,7 @@ fn block_building_works() {
             super::BlockBuild::Finished(Ok(success)) => {
                 let decoded = crate::header::decode(&success.scale_encoded_header).unwrap();
                 assert_eq!(decoded.number, 1);
-                assert_eq!(*decoded.parent_hash, parent_hash);
+                assert_eq!(*decoded.parent_hash, genesis_hash);
                 break;
             }
             super::BlockBuild::Finished(Err(err)) => panic!("{}", err),

src/chain_spec.rs

@@ -39,7 +39,7 @@ use crate::{
         aura_config, babe_genesis_config, grandpa_genesis_config, BabeEpochInformation,
         ChainInformation, ChainInformationConsensus, ChainInformationFinality,
     },
-    executor, libp2p,
+    executor, header, libp2p, trie,
 };
 
 use alloc::{
@@ -179,8 +179,59 @@ impl ChainSpec {
             (finality, vm_prototype)
         };
 
+        let (state_version, vm_prototype) = {
+            match executor::core_version(vm_prototype) {
+                (Ok(runtime_spec), vm_prototype) => {
+                    let state_version = match runtime_spec.decode().state_version {
+                        Some(0) | None => trie::TrieEntryVersion::V0,
+                        Some(1) => trie::TrieEntryVersion::V1,
+                        Some(_) => return Err(FromGenesisStorageError::UnknownStateVersion),
+                    };
+
+                    (state_version, vm_prototype)
+                }
+                (Err(err), _) => return Err(FromGenesisStorageError::CoreVersionLoad(err)),
+            }
+        };
+
         let chain_info = ChainInformation {
-            finalized_block_header: crate::calculate_genesis_block_header(self),
+            finalized_block_header: {
+                let state_root = match self.genesis_storage() {
+                    GenesisStorage::TrieRootHash(hash) => *hash,
+                    GenesisStorage::Items(genesis_storage) => {
+                        let mut calculation = trie::calculate_root::root_merkle_value(None);
+
+                        loop {
+                            match calculation {
+                                trie::calculate_root::RootMerkleValueCalculation::Finished {
+                                    hash,
+                                    ..
+                                } => break hash,
+                                trie::calculate_root::RootMerkleValueCalculation::AllKeys(keys) => {
+                                    calculation = keys.inject(
+                                        genesis_storage.iter().map(|(k, _)| k.iter().copied()),
+                                    );
+                                }
+                                trie::calculate_root::RootMerkleValueCalculation::StorageValue(
+                                    val,
+                                ) => {
+                                    let key: alloc::vec::Vec<u8> = val.key().collect();
+                                    let value = genesis_storage.value(&key[..]);
+                                    calculation = val.inject(state_version, value);
+                                }
+                            }
+                        }
+                    }
+                };
+
+                header::Header {
+                    parent_hash: [0; 32],
+                    number: 0,
+                    state_root,
+                    extrinsics_root: trie::empty_trie_merkle_value(),
+                    digest: header::DigestRef::empty().into(),
+                }
+            },
             consensus,
             finality,
         };
@@ -353,6 +404,14 @@ impl<'a> GenesisStorage<'a> {
             GenesisStorage::TrieRootHash(_) => None,
         }
     }
+
+    /// Returns `Some` for [`GenesisStorage::TrieRootHash`], and `None` otherwise.
+    pub fn into_trie_root_hash(self) -> Option<&'a [u8; 32]> {
+        match self {
+            GenesisStorage::Items(_) => None,
+            GenesisStorage::TrieRootHash(hash) => Some(hash),
+        }
+    }
 }
 
 /// See [`GenesisStorage`].
@@ -475,6 +534,10 @@ pub enum FromGenesisStorageError {
     AuraConfigLoad(aura_config::FromVmPrototypeError),
     /// Error when retrieving the Babe algorithm configuration.
     BabeConfigLoad(babe_genesis_config::FromVmPrototypeError),
+    /// Failed to retrieve the core version of the runtime.
+    CoreVersionLoad(executor::CoreVersionError),
+    /// State version in runtime specification is not supported.
+    UnknownStateVersion,
     /// Multiple consensus algorithms have been detected.
     MultipleConsensusAlgorithms,
     /// Chain specification doesn't contain the list of storage items.

src/executor/host.rs

@@ -791,8 +791,8 @@ impl ReadyToRun {
         macro_rules! expect_state_version {
             ($num:expr) => {{
                 match &params[$num] {
-                    vm::WasmValue::I32(0) => 0,
-                    vm::WasmValue::I32(1) => 1,
+                    vm::WasmValue::I32(0) => trie::TrieEntryVersion::V0,
+                    vm::WasmValue::I32(1) => trie::TrieEntryVersion::V1,
                     v => {
                         return HostVm::Error {
                             error: Error::WrongParamTy {
@@ -932,9 +932,10 @@ impl ReadyToRun {
             HostFunction::ext_storage_root_version_2 => {
                 let state_version = expect_state_version!(0);
                 match state_version {
-                    0 => HostVm::ExternalStorageRoot(ExternalStorageRoot { inner: self.inner }),
-                    1 => host_fn_not_implemented!(), // TODO: https://github.com/paritytech/smoldot/issues/1967
-                    _ => unreachable!(),
+                    trie::TrieEntryVersion::V0 => {
+                        HostVm::ExternalStorageRoot(ExternalStorageRoot { inner: self.inner })
+                    }
+                    trie::TrieEntryVersion::V1 => host_fn_not_implemented!(), // TODO: https://github.com/paritytech/smoldot/issues/1967
                 }
             }
             HostFunction::ext_storage_changes_root_version_1 => {
@@ -1481,13 +1482,12 @@ impl ReadyToRun {
             HostFunction::ext_sandbox_get_global_val_version_1 => host_fn_not_implemented!(),
             HostFunction::ext_trie_blake2_256_root_version_1
             | HostFunction::ext_trie_blake2_256_root_version_2 => {
-                if matches!(host_fn, HostFunction::ext_trie_blake2_256_root_version_2) {
-                    match expect_state_version!(1) {
-                        0 => {}
-                        1 => host_fn_not_implemented!(), // TODO: https://github.com/paritytech/smoldot/issues/1967
-                        _ => unreachable!(),
-                    }
-                }
+                let state_version =
+                    if matches!(host_fn, HostFunction::ext_trie_blake2_256_root_version_2) {
+                        expect_state_version!(1)
+                    } else {
+                        trie::TrieEntryVersion::V0
+                    };
 
                 let result = {
                     let input = expect_pointer_size!(0);
@@ -1514,7 +1514,7 @@ impl ReadyToRun {
                         .map(|(_, parse_result)| parse_result);
 
                     match parsing_result {
-                        Ok(elements) => Ok(trie::trie_root(&elements[..])),
+                        Ok(elements) => Ok(trie::trie_root(state_version, &elements[..])),
                         Err(_) => Err(()),
                     }
                 };
@@ -1531,16 +1531,14 @@ impl ReadyToRun {
             }
             HostFunction::ext_trie_blake2_256_ordered_root_version_1
             | HostFunction::ext_trie_blake2_256_ordered_root_version_2 => {
-                if matches!(
+                let state_version = if matches!(
                     host_fn,
                     HostFunction::ext_trie_blake2_256_ordered_root_version_2
                 ) {
-                    match expect_state_version!(1) {
-                        0 => {}
-                        1 => host_fn_not_implemented!(), // TODO: https://github.com/paritytech/smoldot/issues/1967
-                        _ => unreachable!(),
-                    }
-                }
+                    expect_state_version!(1)
+                } else {
+                    trie::TrieEntryVersion::V0
+                };
 
                 let result = {
                     let input = expect_pointer_size!(0);
@@ -1561,7 +1559,7 @@ impl ReadyToRun {
                         .map(|(_, parse_result)| parse_result);
 
                     match parsing_result {
-                        Ok(elements) => Ok(trie::ordered_root(&elements[..])),
+                        Ok(elements) => Ok(trie::ordered_root(state_version, &elements[..])),
                         Err(_) => Err(()),
                     }
                 };

src/executor/runtime_host.rs

@@ -39,7 +39,7 @@
 
 use crate::{
     executor::{self, host, storage_diff, vm},
-    trie::calculate_root,
+    trie::{self, calculate_root},
     util,
 };
 
@@ -259,7 +259,9 @@ impl StorageGet {
                 if let calculate_root::RootMerkleValueCalculation::StorageValue(value_request) =
                     self.inner.root_calculation.take().unwrap()
                 {
-                    self.inner.root_calculation = Some(value_request.inject(value));
+                    // TODO: we only support V0 for now, see https://github.com/paritytech/smoldot/issues/1967
+                    self.inner.root_calculation =
+                        Some(value_request.inject(trie::TrieEntryVersion::V0, value));
                 } else {
                     // We only create a `StorageGet` if the state is `StorageValue`.
                     panic!()
@@ -621,7 +623,9 @@ impl Inner {
                                 .top_trie_changes
                                 .diff_get(&value_request.key().collect::<Vec<_>>())
                             {
-                                self.root_calculation = Some(value_request.inject(overlay));
+                                // TODO: we only support V0 for now, see https://github.com/paritytech/smoldot/issues/1967
+                                self.root_calculation =
+                                    Some(value_request.inject(trie::TrieEntryVersion::V0, overlay));
                             } else {
                                 self.root_calculation =
                                     Some(calculate_root::RootMerkleValueCalculation::StorageValue(

src/header.rs

@@ -116,7 +116,8 @@ pub fn hash_from_scale_encoded_header_vectored(
 /// Returns the value appropriate for [`Header::extrinsics_root`]. Must be passed the list of
 /// transactions in that block.
 pub fn extrinsics_root(transactions: &[impl AsRef<[u8]>]) -> [u8; 32] {
-    trie::ordered_root(transactions)
+    // The extrinsics root is always calculated with V0 of the trie.
+    trie::ordered_root(trie::TrieEntryVersion::V0, transactions)
 }
 
 /// Attempt to decode the given SCALE-encoded header.