Hooks simplification (#361)

examples/demo-nft-module/README.md

@@ -463,15 +463,11 @@ Here's an example of how to do it with `AppTemplate` from `sov-default-stf`:
     fn new(runtime_config: Self::RuntimeConfig) -> Self {
     let runtime = Runtime::new();
     let storage = ZkStorage::with_config(runtime_config).expect("Failed to open zk storage");
-    let tx_verifier = DemoAppTxVerifier::new();
-    let tx_hooks = DemoAppTxHooks::new();
     let app: AppTemplate<
         ZkDefaultContext,
-        DemoAppTxVerifier<ZkDefaultContext>,
         Runtime<ZkDefaultContext>,
-        DemoAppTxHooks<ZkDefaultContext>,
         Vm,
-    > = AppTemplate::new(storage, runtime, tx_verifier, tx_hooks);
+    > = AppTemplate::new(storage, runtime);
     Self(app)
 }
 ```

examples/demo-stf/README.md

@@ -28,97 +28,28 @@ transactions signed by particular private keys. To fill the gap, there's a syste
 bridges between the two layers of abstraction.
 
 The reason the `AppTemplate` is called a "template" is that it's generic. It allows you, the developer, to pass in
-several parameters that specify its exact behavior. In order, these four generics are:
+several parameters that specify its exact behavior. In order, these generics are:
 
 1. `Context`: a per-transaction struct containing the message's sender. This also provides specs for storage access, so we use different `Context`
    implementations for Native and ZK execution. In ZK, we read values non-deterministically from hints and check them against a merkle tree, while in
    native mode we just read values straight from disk.
-2. `TxVerifier`: a struct that verifies the signatures on transactions and deserializes them into messages
-3. `Runtime`: a collection of modules which make up the rollup's public interface
-4. `TxHooks`: a set of functions which are invoked at various points in the transaction lifecycle
+2. `Runtime`: a collection of modules which make up the rollup's public interface
 
-To implement your state transition function, you simply need to specify values for each of these four fields.
+
+To implement your state transition function, you simply need to specify values for each of these fields.
 
 So, a typical app definition looks like this:
 
 ```rust
-pub type MyNativeStf = AppTemplate<DefaultContext, MyTxVerifier<DefaultContext>, MyRuntime<DefaultContext>, MyTxHooks<DefaultContext>>;
-pub type MyZkStf = AppTemplate<ZkDefaultContext, MyTxVerifier<ZkDefaultContext>, MyRuntime<ZkDefaultContext>, MyTxHooks<ZkDefaultContext>>;
+pub type MyNativeStf = AppTemplate<DefaultContext, MyRuntime<DefaultContext>>;
+pub type MyZkStf = AppTemplate<ZkDefaultContext, MyRuntime<ZkDefaultContext>>;
 ```
 
 Note that `DefaultContext` and `ZkDefaultContext` are exported by the `sov_modules_api` crate.
 
 In the remainder of this section, we'll walk you through implementing each of the remaining generics.
 
-### Implementing a TxVerifier
-
-The `TxVerifier` interface is defined in `sov-default-stf`, and has one associated type and one required method:
-
-```rust
-/// TxVerifier encapsulates Transaction verification.
-pub trait TxVerifier {
-    type Transaction;
-    /// Runs stateless checks against a single RawTx.
-    fn verify_tx_stateless(&self, raw_tx: RawTx) -> anyhow::Result<Self::Transaction>;
-```
-
-The semantics of the `TxVerifier` are pretty straightforward - it takes a `RawTx` (a slice of bytes) as an argument, and does
-some work to transform it into some output `Transaction` type _without looking at the current rollup state_. This output transaction
-type will eventually be fed to the `TxHooks` for _stateful_ verification.
-
-A typical workflow for a `TxVerifier` is to deserialize the message, and check its signature. As you can see by looking
-at the implementation in `tx_verifier_impl.rs`, this is exactly what we do:
-
-```rust
-impl<C: Context> TxVerifier for DemoAppTxVerifier<C> {
-	// ...
-    fn verify_tx_stateless(&self, raw_tx: RawTx) -> anyhow::Result<Self::Transaction> {
-        let mut data = Cursor::new(&raw_tx.data);
-        let tx = Transaction::<C>::deserialize_reader(&mut data)?;
-
-        // We check signature against runtime_msg and nonce.
-        let mut hasher = C::Hasher::new();
-        hasher.update(&tx.runtime_msg);
-        hasher.update(&tx.nonce.to_le_bytes());
-
-        let msg_hash = hasher.finalize();
-
-        tx.signature.verify(&tx.pub_key, msg_hash)?;
-        Ok(tx)
-    }
-}
-```
-
-#### Implementing TxHooks
-
-Once a transaction has passed stateless verification, it will get fed into the execution pipeline. In this pipeline there are four places
-where you can inject custom "hooks" using your `TxHooks` implementation.
-
-1. At the beginning of the `apply_blob` function, before the blob is deserialized into a group of transactions. This is a good time to
-   apply per-batch validation logic like ensuring that the sequencer is properly bonded.
-2. Immediately before each transaction is dispatched to the runtime. This is a good time to apply stateful transaction verification, like checking
-   the nonce.
-3. Immediately after each transaction is executed. This is a good place to perform any post-execution operations, like incrementing the nonce.
-4. At the end of the `apply_blob` function. This is a good place to reward sequencers.
-
-To use the `AppTemplate`, you need to provide a `TxHooks` implementation which specifies what needs to happen at each of these four
-stages.
-
-Its common for modules that need access to these hooks to export a `Hooks` struct. If you're relying on an unfamiliar module, be sure to check
-its documentation to make sure that you know about any hooks that it may rely on. Your `TxHooks` implementation will usually
-just be a wrapper which invokes each of these modules hooks. In this demo, we only rely
-on two modules which need access to the hooks - `sov-accounts` and `sequencer-registry`, so our `TxHooks` implementation only has two fields.
-
-```rust
-pub struct DemoAppTxHooks<C: Context> {
-    accounts_hooks: accounts::hooks::Hooks<C>,
-    sequencer_hooks: sov_sequencer_registry::hooks::Hooks<C>,
-}
-```
-
-You can view the full implementation in `tx_hooks_impl.rs`
-
-### Implementing Runtime: Pick Your Modules
+## Implementing Runtime: Pick Your Modules
 
 The final piece of the puzzle is your app's runtime. A runtime is just a list of modules - really, that's it! To add a new
 module to your app, just add an additional field to the runtime.
@@ -143,6 +74,72 @@ initialization code for each module which will get run at your rollup's genesis.
 allow your runtime to dispatch transactions and queries, and tell it which serialization scheme to use.
 We recommend borsh, since it's both fast and safe for hashing.
 
+### Implementing Hooks for the Runtime: 
+The next step is to implement `Hooks` for `MyRuntime`. Hooks are abstractions that allows for the injection of custom logic into the transaction processing pipeline.
+
+There are two kind of hooks:
+
+`TxHooks`, which has the following methods:
+1. `pre_dispatch_tx_hook`: Invoked immediately before each transaction is processed. This is a good time to apply stateful transaction verification, like checking the nonce.
+2. `post_dispatch_tx_hook`: Invoked immediately after each transaction is executed. This is a good place to perform any post-execution operations, like incrementing the nonce.
+
+`ApplyBlobHooks`, which has the following methods: 
+1. `begin_blob_hook `Invoked at the beginning of the `apply_blob` function, before the blob is deserialized into a group of transactions. This is a good time to ensure that the sequencer is properly bonded.
+2. `end_blob_hook` invoked at the end of the `apply_blob` function. This is a good place to reward sequencers.
+
+To use the `AppTemplate`, the runtime needs to provide implementation of these hooks which specifies what needs to happen at each of these four stages.
+
+In this demo, we only rely on two modules which need access to the hooks - `sov-accounts` and `sequencer-registry`. 
+
+The `sov-accounts` module implements `TxHooks` because it needs to check and increment the sender nonce for every transaction.
+The `sequencer-registry` implements `ApplyBlobHooks` since it is responsible for managing the sequencer bond.
+
+The implementation for `MyRuntime` is straightforward because we can leverage the existing hooks provided by `sov-accounts` and `sequencer-registry` and reuse them in our implementation.
+
+```Rust
+impl<C: Context> TxHooks for Runtime<C> {
+    type Context = C;
+
+    fn pre_dispatch_tx_hook(
+        &self,
+        tx: Transaction<Self::Context>,
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
+        self.accounts.pre_dispatch_tx_hook(tx, working_set)
+    }
+
+    fn post_dispatch_tx_hook(
+        &self,
+        tx: &Transaction<Self::Context>,
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<()> {
+        self.accounts.post_dispatch_tx_hook(tx, working_set)
+    }
+}
+```
+
+```Rust
+impl<C: Context> ApplyBlobHooks for Runtime<C> {
+    type Context = C;
+
+    fn lock_sequencer_bond(
+        &self,
+        sequencer: &[u8],
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<()> {
+        self.sequencer.lock_sequencer_bond(sequencer, working_set)
+    }
+
+    fn reward_sequencer(
+        &self,
+        amount: u64,
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<()> {
+        self.sequencer.reward_sequencer(amount, working_set)
+    }
+}
+```
+
 That's it - with those three structs implemented, you can plug them into your `AppTemplate` and get a
 complete State Transition Function!
 
@@ -188,7 +185,7 @@ impl<Vm: Zkvm> StateTransitionRunner<ProverConfig, Vm> for DemoAppRunner<Default
     fn new(runtime_config: Self::RuntimeConfig) -> Self {
         let storage = ProverStorage::with_config(runtime_config.storage)
             .expect("Failed to open prover storage");
-        let app = AppTemplate::new(storage, Runtime::new(), DemoAppTxVerifier::new(), DemoAppTxHooks::new());
+        let app = AppTemplate::new(storage, Runtime::new());
         Self(app)
     }
 	// ...
@@ -200,7 +197,7 @@ impl<Vm: Zkvm> StateTransitionRunner<ZkConfig, Vm> for DemoAppRunner<ZkDefaultCo
 
     fn new(runtime_config: Self::RuntimeConfig) -> Self {
         let storage = ZkStorage::with_config(runtime_config).expect("Failed to open zk storage");
-        let app = AppTemplate::new(storage, Runtime::new(), DemoAppTxVerifier::new(), DemoAppTxHooks::new());
+        let app = AppTemplate::new(storage, Runtime::new());
         Self(app)
     }
 	// ...

examples/demo-stf/src/app.rs

@@ -1,8 +1,6 @@
 #[cfg(feature = "native")]
 use crate::runner_config::Config;
 use crate::runtime::Runtime;
-use crate::tx_hooks_impl::DemoAppTxHooks;
-use crate::tx_verifier_impl::DemoAppTxVerifier;
 use sov_default_stf::AppTemplate;
 pub use sov_default_stf::Batch;
 use sov_default_stf::SequencerOutcome;
@@ -43,7 +41,7 @@ use sov_modules_macros::expose_rpc;
 #[cfg(feature = "native")]
 pub type NativeAppRunner<Vm> = DemoAppRunner<DefaultContext, Vm>;
 
-pub type DemoApp<C, Vm> = AppTemplate<C, DemoAppTxVerifier<C>, Runtime<C>, DemoAppTxHooks<C>, Vm>;
+pub type DemoApp<C, Vm> = AppTemplate<C, Runtime<C>, Vm>;
 
 /// Batch receipt type used by the demo app. We export this type so that it's easily accessible to the full node.
 pub type DemoBatchReceipt = SequencerOutcome;
@@ -60,9 +58,7 @@ impl<Vm: Zkvm> StateTransitionRunner<ProverConfig, Vm> for DemoAppRunner<Default
         let runtime = Runtime::new();
         let storage = ProverStorage::with_config(runtime_config.storage)
             .expect("Failed to open prover storage");
-        let tx_verifier = DemoAppTxVerifier::new();
-        let tx_hooks = DemoAppTxHooks::new();
-        let app = AppTemplate::new(storage, runtime, tx_verifier, tx_hooks);
+        let app = AppTemplate::new(storage, runtime);
         Self(app)
     }
 
@@ -82,15 +78,8 @@ impl<Vm: Zkvm> StateTransitionRunner<ZkConfig, Vm> for DemoAppRunner<ZkDefaultCo
     fn new(runtime_config: Self::RuntimeConfig) -> Self {
         let runtime = Runtime::new();
         let storage = ZkStorage::with_config(runtime_config).expect("Failed to open zk storage");
-        let tx_verifier = DemoAppTxVerifier::new();
-        let tx_hooks = DemoAppTxHooks::new();
-        let app: AppTemplate<
-            ZkDefaultContext,
-            DemoAppTxVerifier<ZkDefaultContext>,
-            Runtime<ZkDefaultContext>,
-            DemoAppTxHooks<ZkDefaultContext>,
-            Vm,
-        > = AppTemplate::new(storage, runtime, tx_verifier, tx_hooks);
+        let app: AppTemplate<ZkDefaultContext, Runtime<ZkDefaultContext>, Vm> =
+            AppTemplate::new(storage, runtime);
         Self(app)
     }
 

examples/demo-stf/src/bank_cmd/main.rs

@@ -1,8 +1,9 @@
 use anyhow::Context;
 use borsh::BorshSerialize;
 use clap::Parser;
-use demo_stf::{runtime::Runtime, sign_tx, Transaction};
+use demo_stf::runtime::Runtime;
 use sov_default_stf::RawTx;
+use sov_modules_api::transaction::Transaction;
 use sov_modules_api::{
     default_context::DefaultContext, default_signature::private_key::DefaultPrivateKey, PublicKey,
     Spec,
@@ -103,7 +104,7 @@ impl SerializedTx {
         let sender_address = sender_priv_key.pub_key().to_address();
         let message = Self::serialize_call_message(call_data_path, &sender_address)?;
 
-        let sig = sign_tx(&sender_priv_key, &message, nonce);
+        let sig = Transaction::<C>::sign(&sender_priv_key, &message, nonce);
         let tx = Transaction::<C>::new(message, sender_priv_key.pub_key(), sig, nonce);
 
         Ok(SerializedTx {
@@ -289,7 +290,7 @@ mod test {
         )
         .inner;
         assert!(
-            matches!(apply_blob_outcome, SequencerOutcome::Rewarded,),
+            matches!(apply_blob_outcome, SequencerOutcome::Rewarded(0),),
             "Sequencer execution should have succeeded but failed "
         );
         StateTransitionFunction::<MockZkvm>::end_slot(demo);

examples/demo-stf/src/hooks_impl.rs

@@ -0,0 +1,56 @@
+use crate::runtime::Runtime;
+use sov_default_stf::SequencerOutcome;
+use sov_modules_api::{
+    hooks::{ApplyBlobHooks, TxHooks},
+    transaction::Transaction,
+    Context, Spec,
+};
+use sov_state::WorkingSet;
+
+impl<C: Context> TxHooks for Runtime<C> {
+    type Context = C;
+
+    fn pre_dispatch_tx_hook(
+        &self,
+        tx: Transaction<Self::Context>,
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
+        self.accounts.pre_dispatch_tx_hook(tx, working_set)
+    }
+
+    fn post_dispatch_tx_hook(
+        &self,
+        tx: &Transaction<Self::Context>,
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<()> {
+        self.accounts.post_dispatch_tx_hook(tx, working_set)
+    }
+}
+
+impl<C: Context> ApplyBlobHooks for Runtime<C> {
+    type Context = C;
+    type BlobResult = SequencerOutcome;
+
+    fn begin_blob_hook(
+        &self,
+        sequencer: &[u8],
+        raw_blob: &[u8],
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<()> {
+        self.sequencer
+            .begin_blob_hook(sequencer, raw_blob, working_set)
+    }
+
+    fn end_blob_hook(
+        &self,
+        result: Self::BlobResult,
+        working_set: &mut WorkingSet<<Self::Context as Spec>::Storage>,
+    ) -> anyhow::Result<()> {
+        let reward = match result {
+            SequencerOutcome::Rewarded(r) => r,
+            SequencerOutcome::Slashed(_) => 0,
+            SequencerOutcome::Ignored => 0,
+        };
+        self.sequencer.end_blob_hook(reward, working_set)
+    }
+}

examples/demo-stf/src/lib.rs

@@ -1,29 +1,10 @@
 pub mod app;
 #[cfg(feature = "native")]
 pub mod genesis_config;
+pub mod hooks_impl;
 #[cfg(feature = "native")]
 pub mod runner_config;
 pub mod runtime;
 #[cfg(test)]
 pub mod tests;
-pub mod tx_hooks_impl;
-pub mod tx_verifier_impl;
-
-#[cfg(feature = "native")]
-use sov_modules_api::{
-    default_context::DefaultContext,
-    default_signature::{private_key::DefaultPrivateKey, DefaultSignature},
-    Hasher, Spec,
-};
-
 pub use sov_state::ArrayWitness;
-pub use tx_verifier_impl::Transaction;
-
-#[cfg(feature = "native")]
-pub fn sign_tx(priv_key: &DefaultPrivateKey, message: &[u8], nonce: u64) -> DefaultSignature {
-    let mut hasher = <DefaultContext as Spec>::Hasher::new();
-    hasher.update(message);
-    hasher.update(&nonce.to_le_bytes());
-    let msg_hash = hasher.finalize();
-    priv_key.sign(msg_hash)
-}