Fix demo-rollup restart issue (Sovereign-Labs#360)

* Fix demo-rollup restart issue

examples/demo-rollup/src/main.rs

@@ -27,6 +27,7 @@ use tracing::{debug, info};
 // RPC related imports
 use demo_stf::app::get_rpc_methods;
 use sov_modules_api::RpcRunner;
+use sov_state::Storage;
 
 // The rollup stores its data in the namespace b"sov-test" on Celestia
 // You can change this constant to point your rollup at a different namespace
@@ -93,6 +94,7 @@ async fn main() -> Result<(), anyhow::Error> {
 
     // Our state transition also implements the RpcRunner interface, so we use that to initialize the RPC server.
     let storj = demo_runner.get_storage();
+    let is_storage_empty = storj.is_empty();
     let mut methods = get_rpc_methods(storj);
     let ledger_rpc_module =
         ledger_rpc::get_ledger_rpc::<DemoBatchReceipt, DemoTxReceipt>(ledger_db.clone());
@@ -118,10 +120,8 @@ async fn main() -> Result<(), anyhow::Error> {
     });
 
     let demo = demo_runner.inner_mut();
-    // Check if the rollup has previously processed any data. If not, run it's "genesis" initialization code
-    let item_numbers = ledger_db.get_next_items_numbers();
-    let last_slot_processed_before_shutdown = item_numbers.slot_number - 1;
-    if last_slot_processed_before_shutdown == 0 {
+    // Check if the rollup has previously been initialized
+    if is_storage_empty {
         info!("No history detected. Initializing chain...");
         demo.init_chain(get_genesis_config());
         info!("Chain initialization is done.");
@@ -136,6 +136,8 @@ async fn main() -> Result<(), anyhow::Error> {
     let mut prev_state_root = prev_state_root.0;
 
     // Start the main rollup loop
+    let item_numbers = ledger_db.get_next_items_numbers();
+    let last_slot_processed_before_shutdown = item_numbers.slot_number - 1;
     let start_height = rollup_config.start_height + last_slot_processed_before_shutdown;
 
     for height in start_height.. {

full-node/db/sov-schema-db/src/lib.rs

@@ -388,7 +388,7 @@ pub mod temppath {
         }
     }
 
-    impl std::convert::AsRef<Path> for TempPath {
+    impl AsRef<Path> for TempPath {
         fn as_ref(&self) -> &Path {
             self.path()
         }

module-system/sov-state/src/prover_storage.rs

@@ -136,6 +136,11 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
         self.db.inc_next_version();
         Ok(new_root.0)
     }
+
+    // Based on assumption `validate_and_commit` increments version.
+    fn is_empty(&self) -> bool {
+        self.db.get_next_version() <= 1
+    }
 }
 
 pub fn delete_storage(path: impl AsRef<Path>) {
@@ -211,4 +216,34 @@ mod test {
             }
         }
     }
+
+    #[test]
+    fn test_restart_lifecycle() {
+        let path = sov_schema_db::temppath::TempPath::new();
+        {
+            let prover_storage = ProverStorage::<DefaultStorageSpec>::with_path(&path).unwrap();
+            assert!(prover_storage.is_empty());
+        }
+
+        let key = StorageKey::from("some_key");
+        let value = StorageValue::from("some_value");
+        // First restart
+        {
+            let prover_storage = ProverStorage::<DefaultStorageSpec>::with_path(&path).unwrap();
+            assert!(prover_storage.is_empty());
+            let mut storage = WorkingSet::new(prover_storage.clone());
+            storage.set(key.clone(), value.clone());
+            let (cache, witness) = storage.freeze();
+            prover_storage
+                .validate_and_commit(cache, &witness)
+                .expect("storage is valid");
+        }
+
+        // Correctly restart from disk
+        {
+            let prover_storage = ProverStorage::<DefaultStorageSpec>::with_path(&path).unwrap();
+            assert!(!prover_storage.is_empty());
+            assert_eq!(value, prover_storage.get(key, &Default::default()).unwrap());
+        }
+    }
 }

module-system/sov-state/src/storage.rs

@@ -111,6 +111,10 @@ pub trait Storage: Clone {
         state_accesses: OrderedReadsAndWrites,
         witness: &Self::Witness,
     ) -> Result<[u8; 32], anyhow::Error>;
+
+    /// Indicates if storage is empty or not.
+    /// Useful during initialization
+    fn is_empty(&self) -> bool;
 }
 
 // Used only in tests.

module-system/sov-state/src/zk_storage.rs

@@ -91,4 +91,8 @@ impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
 
         Ok(new_root.0)
     }
+
+    fn is_empty(&self) -> bool {
+        unimplemented!("Needs simplification in JellyfishMerkleTree: https://github.com/Sovereign-Labs/sovereign-sdk/issues/362")
+    }
 }