From 5664adc3443125ddb7e210a532dee1035359df58 Mon Sep 17 00:00:00 2001
From: Farhad Shabani <Farhad.Shabani@gmail.com>
Date: Tue, 3 Oct 2023 13:24:18 -0700
Subject: [PATCH] Hooks impl for `demo-stf` should include `chain_state` (#974)

* fix: include chain-state in hooks-impl for demo-stf

* Fix mock validity cond; improve guest logging (#15)

---------

Co-authored-by: Preston Evans <32944016+preston-evans98@users.noreply.github.com>
---
 examples/demo-stf/src/hooks_impl.rs            | 12 ++++++++++++
 full-node/sov-stf-runner/src/runner.rs         | 18 +++++++++---------
 rollup-interface/src/state_machine/mocks/da.rs |  2 +-
 .../state_machine/mocks/validity_condition.rs  | 18 ++++++++----------
 4 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/examples/demo-stf/src/hooks_impl.rs b/examples/demo-stf/src/hooks_impl.rs
index c3f55d754..38e786a69 100644
--- a/examples/demo-stf/src/hooks_impl.rs
+++ b/examples/demo-stf/src/hooks_impl.rs
@@ -90,6 +90,13 @@ impl<C: Context, Da: DaSpec> SlotHooks<Da> for Runtime<C, Da> {
         #[cfg(feature = "experimental")]
         self.evm
             .begin_slot_hook(slot_header.hash().into(), working_set);
+
+        self.chain_state.begin_slot_hook(
+            slot_header,
+            validity_condition,
+            pre_state_root,
+            working_set,
+        );
     }
 
     fn end_slot_hook(
@@ -98,6 +105,8 @@ impl<C: Context, Da: DaSpec> SlotHooks<Da> for Runtime<C, Da> {
     ) {
         #[cfg(feature = "experimental")]
         self.evm.end_slot_hook(working_set);
+
+        self.chain_state.end_slot_hook(working_set);
     }
 }
 
@@ -111,5 +120,8 @@ impl<C: Context, Da: sov_modules_api::DaSpec> FinalizeHook<Da> for Runtime<C, Da
     ) {
         #[cfg(feature = "experimental")]
         self.evm.finalize_hook(root_hash, accessory_working_set);
+
+        self.chain_state
+            .finalize_hook(root_hash, accessory_working_set);
     }
 }
diff --git a/full-node/sov-stf-runner/src/runner.rs b/full-node/sov-stf-runner/src/runner.rs
index 7edf82019..0e6fcc613 100644
--- a/full-node/sov-stf-runner/src/runner.rs
+++ b/full-node/sov-stf-runner/src/runner.rs
@@ -196,16 +196,16 @@ where
                         state_transition_witness: slot_result.witness,
                     };
                 vm.add_hint(transition_data);
-
-                match config {
-                    ProofGenConfig::Simulate(verifier) => {
-                        verifier.run_block(vm.simulate_with_hints()).map_err(|e| {
+                tracing::info_span!("guest_execution").in_scope(|| match config {
+                    ProofGenConfig::Simulate(verifier) => verifier
+                        .run_block(vm.simulate_with_hints())
+                        .map_err(|e| {
                             anyhow::anyhow!("Guest execution must succeed but failed with {:?}", e)
-                        })?;
-                    }
-                    ProofGenConfig::Execute => vm.run(false)?,
-                    ProofGenConfig::Prover => vm.run(true)?,
-                }
+                        })
+                        .map(|_| ()),
+                    ProofGenConfig::Execute => vm.run(false),
+                    ProofGenConfig::Prover => vm.run(true),
+                })?;
             }
             let next_state_root = slot_result.state_root;
 
diff --git a/rollup-interface/src/state_machine/mocks/da.rs b/rollup-interface/src/state_machine/mocks/da.rs
index 20c38d2eb..d4229fb3c 100644
--- a/rollup-interface/src/state_machine/mocks/da.rs
+++ b/rollup-interface/src/state_machine/mocks/da.rs
@@ -373,7 +373,7 @@ impl DaVerifier for MockDaVerifier {
         _inclusion_proof: <Self::Spec as DaSpec>::InclusionMultiProof,
         _completeness_proof: <Self::Spec as DaSpec>::CompletenessProof,
     ) -> Result<<Self::Spec as DaSpec>::ValidityCondition, Self::Error> {
-        Ok(MockValidityCond { is_valid: true })
+        Ok(Default::default())
     }
 }
 
diff --git a/rollup-interface/src/state_machine/mocks/validity_condition.rs b/rollup-interface/src/state_machine/mocks/validity_condition.rs
index e96b72df4..f7ad4ceb2 100644
--- a/rollup-interface/src/state_machine/mocks/validity_condition.rs
+++ b/rollup-interface/src/state_machine/mocks/validity_condition.rs
@@ -9,22 +9,20 @@ use crate::zk::{ValidityCondition, ValidityConditionChecker};
 
 /// A trivial test validity condition structure that only contains a boolean
 #[derive(
-    Debug,
-    BorshDeserialize,
-    BorshSerialize,
-    Serialize,
-    Deserialize,
-    PartialEq,
-    Clone,
-    Copy,
-    Default,
-    Eq,
+    Debug, BorshDeserialize, BorshSerialize, Serialize, Deserialize, PartialEq, Clone, Copy, Eq,
 )]
 pub struct MockValidityCond {
     /// The associated validity condition field. If it is true, the validity condition is verified
     pub is_valid: bool,
 }
 
+// Validity conditions should usually be valid
+impl Default for MockValidityCond {
+    fn default() -> Self {
+        Self { is_valid: true }
+    }
+}
+
 impl ValidityCondition for MockValidityCond {
     type Error = Error;
     fn combine<H: Digest>(&self, rhs: Self) -> Result<Self, Self::Error> {