Test for runner initialization

Sovereign-Labs · Dec 18, 2023 · 847f43b · 847f43b
1 parent 50b6a55
commit 847f43b
Show file tree

Hide file tree

Showing 6 changed files with 246 additions and 110 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/full-node/sov-stf-runner/Cargo.toml b/full-node/sov-stf-runner/Cargo.toml
@@ -30,6 +30,7 @@ sov-rollup-interface = { path = "../../rollup-interface", version = "0.3" }
 
 [dev-dependencies]
 tempfile = { workspace = true }
+sha2 = { workspace = true }
 
 sov-sequencer-registry = { path = "../../module-system/module-implementations/sov-sequencer-registry", features = ["native"] }
 sov-bank = { path = "../../module-system/module-implementations/sov-bank", features = ["native"] }

diff --git a/full-node/sov-stf-runner/src/runner.rs b/full-node/sov-stf-runner/src/runner.rs
@@ -57,7 +57,9 @@ pub enum InitVariant<Stf: StateTransitionFunction<Vm, Da>, Vm: Zkvm, Da: DaSpec>
     /// From empty state root
     Genesis {
         /// Genesis block header should be finalized at init moment
-        genesis_block_header: Da::BlockHeader,
+        block_header: Da::BlockHeader,
+        /// Genesis params for Stf::init
+        genesis_params: GenesisParams<Stf, Vm, Da>,
     },
 }
 
@@ -89,7 +91,6 @@ where
         stf: Stf,
         mut storage_manager: Sm,
         init_variant: InitVariant<Stf, Vm, Da::Spec>,
-        genesis_config: GenesisParams<Stf, Vm, Da::Spec>,
         prover_service: Ps,
     ) -> Result<Self, anyhow::Error> {
         let rpc_config = runner_config.rpc_config;
@@ -100,16 +101,17 @@ where
                 state_root
             }
             InitVariant::Genesis {
-                genesis_block_header,
+                block_header,
+                genesis_params: params,
             } => {
                 info!(
                     "No history detected. Initializing chain on block_header={:?}...",
-                    genesis_block_header
+                    block_header
                 );
-                let storage = storage_manager.create_storage_on(&genesis_block_header)?;
-                let (genesis_root, initialized_storage) = stf.init_chain(storage, genesis_config);
-                storage_manager.save_change_set(&genesis_block_header, initialized_storage)?;
-                storage_manager.finalize(&genesis_block_header)?;
+                let storage = storage_manager.create_storage_on(&block_header)?;
+                let (genesis_root, initialized_storage) = stf.init_chain(storage, params);
+                storage_manager.save_change_set(&block_header, initialized_storage)?;
+                storage_manager.finalize(&block_header)?;
                 info!(
                     "Chain initialization is done. Genesis root: 0x{}",
                     hex::encode(genesis_root.as_ref()),
@@ -265,4 +267,9 @@ where
 
         Ok(())
     }
+
+    /// Allows to read current state root
+    pub fn get_state_root(&self) -> &Stf::StateRoot {
+        &self.state_root
+    }
 }
diff --git a/full-node/sov-stf-runner/tests/runner_initialization.rs b/full-node/sov-stf-runner/tests/runner_initialization.rs
@@ -1,65 +1,189 @@
+use sha2::Digest;
 use sov_db::ledger_db::LedgerDB;
 use sov_mock_da::{
     MockAddress, MockBlockHeader, MockDaConfig, MockDaService, MockDaSpec, MockDaVerifier,
     MockValidityCond,
 };
 use sov_mock_zkvm::MockZkvm;
-use sov_rollup_interface::da::DaSpec;
+use sov_prover_storage_manager::{ProverStorageManager, SnapshotManager};
+use sov_rollup_interface::da::{BlobReaderTrait, DaSpec};
+use sov_rollup_interface::stf::{SlotResult, StateTransitionFunction};
 use sov_rollup_interface::storage::HierarchicalStorageManager;
-use sov_stf_runner::mock::MockStf;
+use sov_rollup_interface::zk::{ValidityCondition, Zkvm};
+use sov_state::storage::{StorageKey, StorageValue};
+use sov_state::{
+    ArrayWitness, DefaultStorageSpec, OrderedReadsAndWrites, Prefix, ProverStorage, Storage,
+};
 use sov_stf_runner::{
     InitVariant, ParallelProverService, RollupConfig, RollupProverConfig, RpcConfig, RunnerConfig,
     StateTransitionRunner, StorageConfig,
 };
 
-struct MockStorageManager<Da: DaSpec> {
-    phantom_data: std::marker::PhantomData<Da>,
+type MockInitVariant = InitVariant<HashStf<MockValidityCond>, MockZkvm, MockDaSpec>;
+
+type S = DefaultStorageSpec;
+type Q = SnapshotManager;
+
+type StorageManager = ProverStorageManager<MockDaSpec, S>;
+
+#[derive(Default, Clone)]
+struct HashStf<Cond> {
+    phantom_data: std::marker::PhantomData<Cond>,
 }
 
-impl<Da: DaSpec> Default for MockStorageManager<Da> {
-    fn default() -> Self {
+impl<Cond> HashStf<Cond> {
+    fn new() -> Self {
         Self {
             phantom_data: std::marker::PhantomData,
         }
     }
-}
 
-impl<Da: DaSpec> HierarchicalStorageManager<Da> for MockStorageManager<Da> {
-    type NativeStorage = ();
-    type NativeChangeSet = ();
-
-    fn create_storage_on(
-        &mut self,
-        _block_header: &Da::BlockHeader,
-    ) -> anyhow::Result<Self::NativeStorage> {
-        Ok(())
+    fn hash_key() -> StorageKey {
+        let prefix = Prefix::new(b"root".to_vec());
+        StorageKey::singleton(&prefix)
     }
 
-    fn create_finalized_storage(&mut self) -> anyhow::Result<Self::NativeStorage> {
-        Ok(())
+    fn save_from_hasher(
+        hasher: sha2::Sha256,
+        storage: ProverStorage<S, Q>,
+        witness: &ArrayWitness,
+    ) -> ([u8; 32], ProverStorage<S, Q>) {
+        let result = hasher.finalize();
+
+        let hash_key = HashStf::<Cond>::hash_key();
+        let hash_value = StorageValue::from(result.as_slice().to_vec());
+
+        let ordered_reads_writes = OrderedReadsAndWrites {
+            ordered_reads: Vec::default(),
+            ordered_writes: vec![(hash_key.to_cache_key(), Some(hash_value.into_cache_value()))],
+        };
+
+        let (jmt_root_hash, state_update) = storage
+            .compute_state_update(ordered_reads_writes, witness)
+            .unwrap();
+
+        storage.commit(&state_update, &OrderedReadsAndWrites::default());
+
+        let mut root_hash = [0u8; 32];
+
+        for (i, &byte) in jmt_root_hash.as_ref().iter().enumerate().take(32) {
+            root_hash[i] = byte;
+        }
+
+        (root_hash, storage)
     }
+}
 
-    fn save_change_set(
-        &mut self,
-        _block_header: &Da::BlockHeader,
-        _change_set: Self::NativeChangeSet,
-    ) -> anyhow::Result<()> {
-        Ok(())
+/// Outcome of the apply_slot method.
+
+impl<Vm: Zkvm, Cond: ValidityCondition, Da: DaSpec> StateTransitionFunction<Vm, Da>
+    for HashStf<Cond>
+{
+    type StateRoot = [u8; 32];
+    type GenesisParams = Vec<u8>;
+    type PreState = ProverStorage<S, Q>;
+    type ChangeSet = ProverStorage<S, Q>;
+    type TxReceiptContents = ();
+    type BatchReceiptContents = [u8; 32];
+    type Witness = ArrayWitness;
+    type Condition = Cond;
+
+    fn init_chain(
+        &self,
+        genesis_state: Self::PreState,
+        params: Self::GenesisParams,
+    ) -> (Self::StateRoot, Self::ChangeSet) {
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(params);
+
+        HashStf::<Cond>::save_from_hasher(hasher, genesis_state, &ArrayWitness::default())
     }
 
-    fn finalize(&mut self, _block_header: &Da::BlockHeader) -> anyhow::Result<()> {
-        Ok(())
+    fn apply_slot<'a, I>(
+        &self,
+        _pre_state_root: &Self::StateRoot,
+        storage: Self::PreState,
+        witness: Self::Witness,
+        _slot_header: &Da::BlockHeader,
+        _validity_condition: &Da::ValidityCondition,
+        blobs: I,
+    ) -> SlotResult<
+        Self::StateRoot,
+        Self::ChangeSet,
+        Self::BatchReceiptContents,
+        Self::TxReceiptContents,
+        Self::Witness,
+    >
+    where
+        I: IntoIterator<Item = &'a mut Da::BlobTransaction>,
+    {
+        let mut hasher = sha2::Sha256::new();
+
+        let hash_key = HashStf::<Cond>::hash_key();
+        let existing_cache = storage.get(&hash_key, None, &witness).unwrap();
+        hasher.update(existing_cache.value());
+
+        for blob in blobs {
+            let data = blob.verified_data();
+            hasher.update(data);
+        }
+
+        let (state_root, storage) = HashStf::<Cond>::save_from_hasher(hasher, storage, &witness);
+
+        SlotResult {
+            state_root,
+            change_set: storage,
+            // TODO: Add batch receipts to inspection
+            batch_receipts: vec![],
+            witness,
+        }
     }
 }
 
-type Stf = MockStf<MockValidityCond>;
 #[tokio::test]
 async fn init_and_restart() {
     let tmpdir = tempfile::tempdir().unwrap();
+    let genesis_params = vec![1, 2, 3, 4, 5];
+    let init_variant: MockInitVariant = InitVariant::Genesis {
+        block_header: MockBlockHeader::from_height(0),
+        genesis_params,
+    };
+
+    let state_root_after_genesis = {
+        let runner = initialize_runner(tmpdir.path(), init_variant);
+        *runner.get_state_root()
+    };
+
+    let init_variant_2: MockInitVariant = InitVariant::Initialized(state_root_after_genesis);
+
+    let runner_2 = initialize_runner(tmpdir.path(), init_variant_2);
+
+    let state_root_2 = *runner_2.get_state_root();
+
+    assert_eq!(state_root_after_genesis, state_root_2);
+}
+
+type MockProverService = ParallelProverService<
+    [u8; 32],
+    ArrayWitness,
+    MockDaService,
+    MockZkvm,
+    HashStf<MockValidityCond>,
+>;
+fn initialize_runner(
+    path: &std::path::Path,
+    init_variant: MockInitVariant,
+) -> StateTransitionRunner<
+    HashStf<MockValidityCond>,
+    StorageManager,
+    MockDaService,
+    MockZkvm,
+    MockProverService,
+> {
     let address = MockAddress::new([11u8; 32]);
     let rollup_config = RollupConfig::<MockDaConfig> {
         storage: StorageConfig {
-            path: tmpdir.path().to_path_buf(),
+            path: path.to_path_buf(),
         },
         runner: RunnerConfig {
             start_height: 1,
@@ -75,37 +199,38 @@ async fn init_and_restart() {
 
     let da_service = MockDaService::new(address);
 
-    let ledger_db = LedgerDB::with_path(tmpdir.path()).unwrap();
-
-    let stf = Stf::default();
+    let ledger_db = LedgerDB::with_path(path).unwrap();
 
-    let storage_manager = MockStorageManager::<MockDaSpec>::default();
+    let stf = HashStf::<MockValidityCond>::new();
 
-    let init_variant: InitVariant<MockStf<MockValidityCond>, MockZkvm, MockDaSpec> =
-        InitVariant::Genesis {
-            genesis_block_header: MockBlockHeader::from_height(0),
-        };
-
-    let genesis_config = ();
+    let storage_config = sov_state::config::Config {
+        path: path.to_path_buf(),
+    };
+    let mut storage_manager = ProverStorageManager::new(storage_config).unwrap();
 
     let vm = MockZkvm::default();
     let verifier = MockDaVerifier::default();
 
     let prover_config = RollupProverConfig::Prove;
 
-    let prover_service =
-        ParallelProverService::new(vm, stf.clone(), verifier, prover_config, (), 1);
-
-    // TODO: Extend test, probably with different STF
-    let _runner = StateTransitionRunner::new(
+    let prover_service = ParallelProverService::new(
+        vm,
+        stf.clone(),
+        verifier,
+        prover_config,
+        // Should be ZkStorage, but we don't need it for this test
+        storage_manager.create_finalized_storage().unwrap(),
+        1,
+    );
+
+    StateTransitionRunner::new(
         rollup_config.runner,
         da_service,
         ledger_db,
         stf,
         storage_manager,
         init_variant,
-        genesis_config,
         prover_service,
     )
-    .unwrap();
+    .unwrap()
 }