Repository: oasisprotocol/oasis-sdk
Pull request #148 (2021-05-26 17:31:23)

runtime-sdk: Static call/query/message result handler dispatch
Refactors the way method dispatch is performed by removing the method
registries and replacing them with static dispatch in the form of a
MethodHandler trait which needs to be provided by modules.

Additionally the contexts are made more composable by making sure that
the we take Context as a generic type argument everywhere and improving
the way simulation mode is handled.

Each context is now also bound to a specific Runtime type to allow for
nested method dispatch.

pro-wh: my merge in #150
kostko: Cherry-picked some improvements from @nhynes where we have separate traits for `TxContext` and `BatchContext` to make things more clear and avoid unnecessary options/unwraps.
kostko: > but I thought we were trying to unify the context so that we could have operations that we could use in and out of transactions. we might reasonably use gas in these operations, and now Core::use_gas is bifurcated.

That's a fair point, @nhynes any thoughts on this? I agree that having a unified context could be useful. @pro-wh do you have any specific use cases in mind?
nhynes: > unified context could be useful

I agree that it'd be clean, but is it really true that once you downcast into a `TxContext` you can't do batch contextual operations? If you `use_tx_gas` within a transaction, does it always `use_batch_gas`? Alternatively, you can have `tx_ctx.batch_context().use_gas(...)`. All I know is that:
* static types > runtime types
* this is a low-level SDK (by virtue of there being no others) and it might be useful to be able to off-road batch vs tx operations (big maybe)
* unification is backwards-compatible

> operations that we could use in and out of transactions

That describes `trait Context` in this PR. The gas accounting thing is regrettable, but so is code that does a dynamic check for being in a transaction context.
kostko: > but is it really true that once you downcast into a TxContext you can't do batch contextual operations

The only batch operation is entering a transaction context via `with_tx`. I think that it doesn't make much sense to allow `with_tx` inside a transaction context unless you use `with_simulation` which gives you a simulated batch context.

Patch:
diff --git a/runtime-sdk/src/context.rs b/runtime-sdk/src/context.rs
index fce88fa925..7a3242b779 100644
--- a/runtime-sdk/src/context.rs
+++ b/runtime-sdk/src/context.rs
@@ -1,6 +1,11 @@
 //! Execution context.
 use core::fmt;
-use std::{any::Any, collections::BTreeMap, sync::Arc};
+use std::{
+    any::Any,
+    collections::btree_map::{BTreeMap, Entry},
+    marker::PhantomData,
+    sync::Arc,
+};
 
 use io_context::Context as IoContext;
 use slog::{self, o};
@@ -15,9 +20,8 @@ use oasis_core_runtime::{
 
 use crate::{
     event::Event,
-    module::MethodRegistry,
     modules::core::Error,
-    storage,
+    runtime, storage,
     types::{address::Address, message::MessageEventHookInvocation, transaction},
 };
 
@@ -51,8 +55,10 @@ impl From<&Mode> for &'static str {
 
 /// Runtime SDK context.
 pub trait Context {
+    /// Runtime that the context is being invoked in.
+    type Runtime: runtime::Runtime;
     /// Runtime state output type.
-    type S: storage::Store;
+    type Store: storage::Store;
 
     /// Returns a logger.
     fn get_logger(&self, module: &'static str) -> slog::Logger;
@@ -77,7 +83,7 @@ pub trait Context {
     fn runtime_round_results(&self) -> &roothash::RoundResults;
 
     /// Runtime state store.
-    fn runtime_state(&mut self) -> &mut Self::S;
+    fn runtime_state(&mut self) -> &mut Self::Store;
 
     /// Consensus state.
     fn consensus_state(&self) -> &consensus::state::ConsensusState;
@@ -85,22 +91,6 @@ pub trait Context {
     /// Current epoch.
     fn epoch(&self) -> consensus::beacon::EpochTime;
 
-    /// Transaction authentication information.
-    ///
-    /// Only present if this is a transaction processing context.
-    fn tx_auth_info(&self) -> Option<&transaction::AuthInfo>;
-
-    /// Authenticated address of the caller.
-    ///
-    /// In case there are multiple signers of a transaction, this will return the address
-    /// corresponding to the first signer.
-    ///
-    /// Only present if this is a transaction processing context.
-    fn tx_caller_address(&self) -> Option<Address> {
-        self.tx_auth_info()
-            .map(|info| info.signer_info[0].address_spec.address())
-    }
-
     /// Emits an event.
     fn emit_event<E: Event>(&mut self, event: E);
 
@@ -130,75 +120,116 @@ pub trait Context {
     /// consumes the transaction context.
     fn commit(self) -> (Tags, Vec<(roothash::Message, MessageEventHookInvocation)>);
 
-    /// Fetches or sets a value associated with the context.
-    fn value<V>(&mut self, key: &'static str) -> &mut V
-    where
-        V: Any + Default;
+    /// Fetches a value entry associated with the context.
+    fn value<V: Any>(&mut self, key: &'static str) -> ContextValue<'_, V>;
 
-    /// Takes a value associated with the context.
+    /// Executes a function in a simulation context.
     ///
-    /// The previous value is removed so subsequent fetches will return the default value.
-    fn take_value<V>(&mut self, key: &'static str) -> Box<V>
+    /// The simulation context collects its own messages and starts with an empty set of context
+    /// values.
+    fn with_simulation<F, Rs>(&mut self, f: F) -> Rs
     where
-        V: Any + Default;
+        F: FnOnce(
+            RuntimeBatchContext<'_, Self::Runtime, storage::OverlayStore<&mut Self::Store>>,
+        ) -> Rs;
+}
 
-    /// Fetches or sets a value associated with the transaction.
-    ///
-    /// Only present if this is a transaction processing context.
-    fn tx_value<V>(&mut self, key: &'static str) -> Option<&mut V>
+/// Runtime SDK batch-wide context.
+pub trait BatchContext: Context {
+    /// Executes a function in a per-transaction context.
+    fn with_tx<F, Rs>(&mut self, tx: transaction::Transaction, f: F) -> Rs
     where
-        V: Any + Default;
+        F: FnOnce(
+            RuntimeTxContext<'_, '_, <Self as Context>::Runtime, <Self as Context>::Store>,
+            transaction::Call,
+        ) -> Rs;
+}
 
-    /// Takes a value associated with the transaction.
-    ///
-    /// The previous value is removed so subsequent fetches will return the default value.
+/// Runtime SDK transaction context.
+pub trait TxContext: Context {
+    /// Transaction authentication information.
+    fn tx_auth_info(&self) -> &transaction::AuthInfo;
+
+    /// Authenticated address of the caller.
     ///
-    /// Only present if this is a transaction processing context.
-    fn take_tx_value<V>(&mut self, key: &'static str) -> Option<Box<V>>
-    where
-        V: Any + Default;
+    /// In case there are multiple signers of a transaction, this will return the address
+    /// corresponding to the first signer.
+    fn tx_caller_address(&self) -> Address {
+        self.tx_auth_info().signer_info[0].address_spec.address()
+    }
+
+    /// Fetches an entry pointing to a value associated with the transaction.
+    fn tx_value<V: Any>(&mut self, key: &'static str) -> ContextValue<'_, V>;
 }
 
 /// Dispatch context for the whole batch.
-pub struct DispatchContext<'a> {
-    pub(crate) mode: Mode,
+pub struct RuntimeBatchContext<'a, R: runtime::Runtime, S: storage::Store> {
+    mode: Mode,
 
-    pub(crate) runtime_header: &'a roothash::Header,
-    pub(crate) runtime_round_results: &'a roothash::RoundResults,
-    pub(crate) runtime_storage: storage::MKVSStore<&'a mut dyn mkvs::MKVS>,
+    runtime_header: &'a roothash::Header,
+    runtime_round_results: &'a roothash::RoundResults,
+    runtime_storage: S,
     // TODO: linked consensus layer block
-    pub(crate) consensus_state: &'a consensus::state::ConsensusState,
-    pub(crate) epoch: consensus::beacon::EpochTime,
-    pub(crate) io_ctx: Arc<IoContext>,
-    pub(crate) logger: slog::Logger,
-
-    /// The runtime's methods, in case you need to look them up for some reason.
-    pub(crate) methods: &'a MethodRegistry,
+    consensus_state: &'a consensus::state::ConsensusState,
+    epoch: consensus::beacon::EpochTime,
+    io_ctx: Arc<IoContext>,
+    logger: slog::Logger,
 
-    pub(crate) block_tags: Tags,
+    block_tags: Tags,
 
     /// Maximum number of messages that can be emitted.
-    pub(crate) max_messages: u32,
+    max_messages: u32,
     /// Emitted messages.
-    pub(crate) messages: Vec<(roothash::Message, MessageEventHookInvocation)>,
+    messages: Vec<(roothash::Message, MessageEventHookInvocation)>,
 
     /// Per-context values.
-    pub(crate) values: BTreeMap<&'static str, Box<dyn Any>>,
+    values: BTreeMap<&'static str, Box<dyn Any>>,
+
+    _runtime: PhantomData<R>,
 }
 
-impl<'a> DispatchContext<'a> {
+impl<'a, R: runtime::Runtime, S: storage::Store> RuntimeBatchContext<'a, R, S> {
+    /// Create a new dispatch context.
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        mode: Mode,
+        runtime_header: &'a roothash::Header,
+        runtime_round_results: &'a roothash::RoundResults,
+        runtime_storage: S,
+        consensus_state: &'a consensus::state::ConsensusState,
+        epoch: consensus::beacon::EpochTime,
+        io_ctx: Arc<IoContext>,
+        max_messages: u32,
+    ) -> Self {
+        Self {
+            mode,
+            runtime_header,
+            runtime_round_results,
+            runtime_storage,
+            consensus_state,
+            epoch,
+            io_ctx,
+            logger: get_logger("runtime-sdk")
+                .new(o!("ctx" => "dispatch", "mode" => Into::<&'static str>::into(&mode))),
+            block_tags: Tags::new(),
+            max_messages,
+            messages: Vec::new(),
+            values: BTreeMap::new(),
+            _runtime: PhantomData,
+        }
+    }
+
     /// Create a new dispatch context from the low-level runtime context.
     pub(crate) fn from_runtime(
         ctx: &'a RuntimeContext<'_>,
         mkvs: &'a mut dyn mkvs::MKVS,
-        methods: &'a MethodRegistry,
-    ) -> Self {
+    ) -> RuntimeBatchContext<'a, R, storage::MKVSStore<&'a mut dyn mkvs::MKVS>> {
         let mode = if ctx.check_only {
             Mode::CheckTx
         } else {
             Mode::ExecuteTx
         };
-        Self {
+        RuntimeBatchContext {
             mode,
             runtime_header: ctx.header,
             runtime_round_results: ctx.round_results,
@@ -208,81 +239,18 @@ impl<'a> DispatchContext<'a> {
             io_ctx: ctx.io_ctx.clone(),
             logger: get_logger("runtime-sdk")
                 .new(o!("ctx" => "dispatch", "mode" => Into::<&'static str>::into(&mode))),
-            methods,
             block_tags: Tags::new(),
             max_messages: ctx.max_messages,
             messages: Vec::new(),
             values: BTreeMap::new(),
+            _runtime: PhantomData,
         }
     }
-
-    /// Executes a function with the transaction-specific context set.
-    pub fn with_tx<F, R>(&mut self, tx: transaction::Transaction, f: F) -> R
-    where
-        F: FnOnce(TxContext<'_, '_>, transaction::Call) -> R,
-    {
-        // Create a store wrapped by an overlay store so we can either rollback or commit.
-        let store = storage::OverlayStore::new(&mut self.runtime_storage);
-
-        let tx_ctx = TxContext {
-            mode: self.mode,
-            runtime_header: self.runtime_header,
-            runtime_round_results: self.runtime_round_results,
-            consensus_state: self.consensus_state,
-            epoch: self.epoch,
-            store,
-            io_ctx: self.io_ctx.clone(),
-            logger: self
-                .logger
-                .new(o!("ctx" => "transaction", "mode" => Into::<&'static str>::into(&self.mode))),
-            tx_auth_info: tx.auth_info,
-            tags: Tags::new(),
-            // NOTE: Since a limit is enforced (which is a u32) this cast is always safe.
-            max_messages: self.max_messages.saturating_sub(self.messages.len() as u32),
-            messages: Vec::new(),
-            values: &mut self.values,
-            tx_values: BTreeMap::new(),
-        };
-        f(tx_ctx, tx.call)
-    }
-
-    /// Executes a function with the transaction-specific context set in simulation mode.
-    /// The simulation context collects its own messages and starts with an empty set of context
-    /// values.
-    /// Runtime storage is shared with this context, so don't go committing it.
-    pub fn with_tx_simulation<F, R>(&mut self, tx: transaction::Transaction, f: F) -> R
-    where
-        F: FnOnce(TxContext<'_, '_>, transaction::Call) -> R,
-    {
-        // Create a store wrapped by an overlay store so we can either rollback or commit.
-        let store = storage::OverlayStore::new(&mut self.runtime_storage);
-        let mut values = BTreeMap::new();
-
-        let tx_ctx = TxContext {
-            mode: Mode::SimulateTx,
-            runtime_header: self.runtime_header,
-            runtime_round_results: self.runtime_round_results,
-            consensus_state: self.consensus_state,
-            epoch: self.epoch,
-            store,
-            io_ctx: self.io_ctx.clone(),
-            logger: self
-                .logger
-                .new(o!("ctx" => "transaction", "mode" => MODE_SIMULATE_TX)),
-            tx_auth_info: tx.auth_info,
-            // Other than for storage, simulation has a blank slate.
-            tags: Tags::new(),
-            max_messages: self.max_messages,
-            messages: Vec::new(),
-            values: &mut values,
-            tx_values: BTreeMap::new(),
-        };
-        f(tx_ctx, tx.call)
-    }
 }
 
-impl<'a> Context for DispatchContext<'a> {
-    type S = storage::MKVSStore<&'a mut dyn mkvs::MKVS>;
+impl<'a, R: runtime::Runtime, S: storage::Store> Context for RuntimeBatchContext<'a, R, S> {
+    type Runtime = R;
+    type Store = S;
 
     fn get_logger(&self, module: &'static str) -> slog::Logger {
         self.logger.new(o!("sdk_module" => module))
@@ -300,7 +268,7 @@ impl<'a> Context for DispatchContext<'a> {
         &self.runtime_round_results
     }
 
-    fn runtime_state(&mut self) -> &mut Self::S {
+    fn runtime_state(&mut self) -> &mut Self::Store {
         &mut self.runtime_storage
     }
 
@@ -312,10 +280,6 @@ impl<'a> Context for DispatchContext<'a> {
         self.epoch
     }
 
-    fn tx_auth_info(&self) -> Option<&transaction::AuthInfo> {
-        None
-    }
-
     fn emit_event<E: Event>(&mut self, event: E) {
         self.block_tags.push(event.to_tag());
     }
@@ -343,53 +307,86 @@ impl<'a> Context for DispatchContext<'a> {
         (self.block_tags, self.messages)
     }
 
-    fn value<V>(&mut self, key: &'static str) -> &mut V
-    where
-        V: Any + Default,
-    {
-        self.values
-            .entry(key)
-            .or_insert_with(|| Box::new(V::default()))
-            .downcast_mut()
-            .expect("type should stay the same")
+    fn value<V: Any>(&mut self, key: &'static str) -> ContextValue<'_, V> {
+        ContextValue::new(self.values.entry(key))
     }
 
-    fn take_value<V>(&mut self, key: &'static str) -> Box<V>
+    fn with_simulation<F, Rs>(&mut self, f: F) -> Rs
     where
-        V: Any + Default,
+        F: FnOnce(
+            RuntimeBatchContext<'_, Self::Runtime, storage::OverlayStore<&mut Self::Store>>,
+        ) -> Rs,
     {
-        self.values
-            .remove(key)
-            .map(|x| x.downcast().expect("type should stay the same"))
-            .unwrap_or_default()
-    }
+        // Create a store wrapped by an overlay store so any state changes don't leak.
+        let store = storage::OverlayStore::new(&mut self.runtime_storage);
 
-    fn tx_value<V>(&mut self, _key: &'static str) -> Option<&mut V>
-    where
-        V: Any + Default,
-    {
-        None
+        let sim_ctx = RuntimeBatchContext {
+            mode: Mode::SimulateTx,
+            runtime_header: self.runtime_header,
+            runtime_round_results: self.runtime_round_results,
+            runtime_storage: store,
+            consensus_state: self.consensus_state,
+            epoch: self.epoch,
+            io_ctx: self.io_ctx.clone(),
+            logger: self
+                .logger
+                .new(o!("ctx" => "dispatch", "mode" => MODE_SIMULATE_TX)),
+            // Other than for storage, simulation has a blank slate.
+            block_tags: Tags::new(),
+            max_messages: self.max_messages,
+            messages: Vec::new(),
+            values: BTreeMap::new(),
+            _runtime: PhantomData,
+        };
+        f(sim_ctx)
     }
+}
 
-    fn take_tx_value<V>(&mut self, _key: &'static str) -> Option<Box<V>>
+impl<'a, R: runtime::Runtime, S: storage::Store> BatchContext for RuntimeBatchContext<'a, R, S> {
+    fn with_tx<F, Rs>(&mut self, tx: transaction::Transaction, f: F) -> Rs
     where
-        V: Any + Default,
+        F: FnOnce(
+            RuntimeTxContext<'_, '_, <Self as Context>::Runtime, <Self as Context>::Store>,
+            transaction::Call,
+        ) -> Rs,
     {
-        None
+        // Create a store wrapped by an overlay store so we can either rollback or commit.
+        let store = storage::OverlayStore::new(&mut self.runtime_storage);
+
+        let tx_ctx = RuntimeTxContext {
+            mode: self.mode,
+            runtime_header: self.runtime_header,
+            runtime_round_results: self.runtime_round_results,
+            consensus_state: self.consensus_state,
+            epoch: self.epoch,
+            store,
+            io_ctx: self.io_ctx.clone(),
+            logger: self
+                .logger
+                .new(o!("ctx" => "transaction", "mode" => Into::<&'static str>::into(&self.mode))),
+            tx_auth_info: tx.auth_info,
+            tags: Tags::new(),
+            // NOTE: Since a limit is enforced (which is a u32) this cast is always safe.
+            max_messages: self.max_messages.saturating_sub(self.messages.len() as u32),
+            messages: Vec::new(),
+            values: &mut self.values,
+            tx_values: BTreeMap::new(),
+            _runtime: PhantomData,
+        };
+        f(tx_ctx, tx.call)
     }
 }
 
 /// Per-transaction/method dispatch sub-context.
-pub struct TxContext<'a, 'b> {
+pub struct RuntimeTxContext<'round, 'store, R: runtime::Runtime, S: storage::Store> {
     mode: Mode,
 
-    runtime_header: &'a roothash::Header,
-    runtime_round_results: &'a roothash::RoundResults,
-    consensus_state: &'a consensus::state::ConsensusState,
+    runtime_header: &'round roothash::Header,
+    runtime_round_results: &'round roothash::RoundResults,
+    consensus_state: &'round consensus::state::ConsensusState,
     epoch: consensus::beacon::EpochTime,
     // TODO: linked consensus layer block
-    store: storage::OverlayStore<&'b mut storage::MKVSStore<&'a mut dyn mkvs::MKVS>>,
-
+    store: storage::OverlayStore<&'store mut S>,
     io_ctx: Arc<IoContext>,
     logger: slog::Logger,
 
@@ -405,14 +402,19 @@ pub struct TxContext<'a, 'b> {
     messages: Vec<(roothash::Message, MessageEventHookInvocation)>,
 
     /// Per-context values.
-    values: &'b mut BTreeMap<&'static str, Box<dyn Any>>,
+    values: &'store mut BTreeMap<&'static str, Box<dyn Any>>,
 
     /// Per-transaction values.
     tx_values: BTreeMap<&'static str, Box<dyn Any>>,
+
+    _runtime: PhantomData<R>,
 }
 
-impl<'a, 'b> Context for TxContext<'a, 'b> {
-    type S = storage::OverlayStore<&'b mut storage::MKVSStore<&'a mut dyn mkvs::MKVS>>;
+impl<'round, 'store, R: runtime::Runtime, S: storage::Store> Context
+    for RuntimeTxContext<'round, 'store, R, S>
+{
+    type Runtime = R;
+    type Store = storage::OverlayStore<&'store mut S>;
 
     fn get_logger(&self, module: &'static str) -> slog::Logger {
         self.logger.new(o!("sdk_module" => module))
@@ -430,7 +432,7 @@ impl<'a, 'b> Context for TxContext<'a, 'b> {
         self.runtime_round_results
     }
 
-    fn runtime_state(&mut self) -> &mut Self::S {
+    fn runtime_state(&mut self) -> &mut Self::Store {
         &mut self.store
     }
 
@@ -442,10 +444,6 @@ impl<'a, 'b> Context for TxContext<'a, 'b> {
         self.epoch
     }
 
-    fn tx_auth_info(&self) -> Option<&transaction::AuthInfo> {
-        Some(&self.tx_auth_info)
-    }
-
     fn emit_event<E: Event>(&mut self, event: E) {
         self.tags.push(event.to_tag());
     }
@@ -474,58 +472,148 @@ impl<'a, 'b> Context for TxContext<'a, 'b> {
         (self.tags, self.messages)
     }
 
-    fn value<V>(&mut self, key: &'static str) -> &mut V
-    where
-        V: Any + Default,
-    {
-        self.values
-            .entry(key)
-            .or_insert_with(|| Box::new(V::default()))
-            .downcast_mut()
-            .expect("type should stay the same")
+    fn value<V: Any>(&mut self, key: &'static str) -> ContextValue<'_, V> {
+        ContextValue::new(self.values.entry(key))
     }
 
-    fn take_value<V>(&mut self, key: &'static str) -> Box<V>
+    fn with_simulation<F, Rs>(&mut self, f: F) -> Rs
     where
-        V: Any + Default,
+        F: FnOnce(
+            RuntimeBatchContext<'_, Self::Runtime, storage::OverlayStore<&mut Self::Store>>,
+        ) -> Rs,
     {
-        self.values
-            .remove(key)
-            .map(|x| x.downcast().expect("type should stay the same"))
-            .unwrap_or_default()
+        // Create a store wrapped by an overlay store so any state changes don't leak.
+        let store = storage::OverlayStore::new(&mut self.store);
+
+        let sim_ctx = RuntimeBatchContext {
+            mode: Mode::SimulateTx,
+            runtime_header: self.runtime_header,
+            runtime_round_results: self.runtime_round_results,
+            runtime_storage: store,
+            consensus_state: self.consensus_state,
+            epoch: self.epoch,
+            io_ctx: self.io_ctx.clone(),
+            logger: self
+                .logger
+                .new(o!("ctx" => "dispatch", "mode" => MODE_SIMULATE_TX)),
+            // Other than for storage, simulation has a blank slate.
+            block_tags: Tags::new(),
+            max_messages: self.max_messages,
+            messages: Vec::new(),
+            values: BTreeMap::new(),
+            _runtime: PhantomData,
+        };
+        f(sim_ctx)
+    }
+}
+
+impl<R: runtime::Runtime, S: storage::Store> TxContext for RuntimeTxContext<'_, '_, R, S> {
+    fn tx_auth_info(&self) -> &transaction::AuthInfo {
+        &self.tx_auth_info
     }
 
-    /// Fetches or sets a value associated with the transaction.
-    fn tx_value<V>(&mut self, key: &'static str) -> Option<&mut V>
-    where
-        V: Any + Default,
-    {
-        Some(
-            self.tx_values
-                .entry(key)
-                .or_insert_with(|| Box::new(V::default()))
-                .downcast_mut()
-                .expect("type should stay the same"),
-        )
+    fn tx_value<V: Any>(&mut self, key: &'static str) -> ContextValue<'_, V> {
+        ContextValue::new(self.tx_values.entry(key))
     }
+}
 
-    /// Takes a value associated with the transaction.
+/// A per-context arbitrary value.
+pub struct ContextValue<'a, V> {
+    inner: Entry<'a, &'static str, Box<dyn Any>>,
+    _value: PhantomData<V>,
+}
+
+impl<'a, V: Any> ContextValue<'a, V> {
+    fn new(inner: Entry<'a, &'static str, Box<dyn Any>>) -> Self {
+        Self {
+            inner,
+            _value: PhantomData,
+        }
+    }
+
+    /// Gets a reference to the specified per-context value.
     ///
-    /// The previous value is removed so subsequent fetches will return the default value.
-    fn take_tx_value<V>(&mut self, key: &'static str) -> Option<Box<V>>
-    where
-        V: Any + Default,
-    {
-        Some(
-            self.tx_values
-                .remove(key)
-                .map(|x| x.downcast().expect("type should stay the same"))
-                .unwrap_or_default(),
-        )
+    /// # Panics
+    ///
+    /// Panics if the retrieved type is not the type that was stored.
+    pub fn get(self) -> Option<&'a V> {
+        match self.inner {
+            Entry::Occupied(oe) => Some(
+                oe.into_mut()
+                    .downcast_ref()
+                    .expect("type should stay the same"),
+            ),
+            _ => None,
+        }
+    }
+
+    /// Gets a mutable reference to the specifeid per-context value.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the retrieved type is not the type that was stored.
+    pub fn get_mut(&mut self) -> Option<&mut V> {
+        match &mut self.inner {
+            Entry::Occupied(oe) => Some(
+                oe.get_mut()
+                    .downcast_mut()
+                    .expect("type should stay the same"),
+            ),
+            _ => None,
+        }
+    }
+
+    /// Sets the context value, returning a mutable reference to the set value.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the retrieved type is not the type that was stored.
+    pub fn set(self, value: V) -> &'a mut V {
+        let value = Box::new(value);
+        match self.inner {
+            Entry::Occupied(mut oe) => {
+                oe.insert(value);
+                oe.into_mut()
+            }
+            Entry::Vacant(ve) => ve.insert(value),
+        }
+        .downcast_mut()
+        .expect("type should stay the same")
+    }
+
+    /// Takes the context value, if it exists.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the retrieved type is not the type that was stored.
+    pub fn take(self) -> Option<V> {
+        match self.inner {
+            Entry::Occupied(oe) => {
+                Some(*oe.remove().downcast().expect("type should stay the same"))
+            }
+            Entry::Vacant(_) => None,
+        }
+    }
+}
+
+impl<'a, V: Any + Default> ContextValue<'a, V> {
+    /// Retrieves the existing value or inserts and returns the default.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the retrieved type is not the type that was stored.
+    pub fn or_default(self) -> &'a mut V {
+        match self.inner {
+            Entry::Occupied(oe) => oe.into_mut(),
+            Entry::Vacant(ve) => ve.insert(Box::new(V::default())),
+        }
+        .downcast_mut()
+        .expect("type should stay the same")
     }
 }
 
 #[cfg(test)]
+#[allow(clippy::many_single_char_names)]
 mod test {
     use oasis_core_runtime::{common::cbor, consensus::staking};
 
@@ -537,18 +625,19 @@ mod test {
         let mut mock = Mock::default();
         let mut ctx = mock.create_ctx();
 
-        let x: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(x, &None, "default value should be created");
-        *x = Some(42);
+        let x = ctx.value::<u64>("module.TestKey").get();
+        assert_eq!(x, None);
+
+        ctx.value::<u64>("module.TestKey").set(42);
 
-        let y: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(y, &Some(42));
+        let y = ctx.value::<u64>("module.TestKey").get();
+        assert_eq!(y, Some(&42u64));
 
-        let z: Box<Option<u64>> = ctx.take_value("module.TestKey");
-        assert_eq!(z, Box::new(Some(42)));
+        let z = ctx.value::<u64>("module.TestKey").take();
+        assert_eq!(z, Some(42u64));
 
-        let y: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(y, &None);
+        let y = ctx.value::<u64>("module.TestKey").get();
+        assert_eq!(y, None);
     }
 
     #[test]
@@ -557,12 +646,8 @@ mod test {
         let mut mock = Mock::default();
         let mut ctx = mock.create_ctx();
 
-        let x: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(x, &None, "default value should be created");
-        *x = Some(42);
-
-        // Changing the type of a key should result in a panic.
-        ctx.value::<Option<u32>>("module.TestKey");
+        ctx.value::<u64>("module.TestKey").or_default();
+        ctx.value::<u32>("module.TestKey").get();
     }
 
     #[test]
@@ -570,9 +655,7 @@ mod test {
         let mut mock = Mock::default();
         let mut ctx = mock.create_ctx();
 
-        let x: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(x, &None, "default value should be created");
-        *x = Some(42);
+        ctx.value("module.TestKey").set(42u64);
 
         let tx = transaction::Transaction {
             version: 1,
@@ -589,40 +672,39 @@ mod test {
             },
         };
         ctx.with_tx(tx.clone(), |mut tx_ctx, _call| {
-            let y: &mut Option<u64> = tx_ctx.value("module.TestKey");
-            assert_eq!(y, &Some(42));
-
-            *y = Some(48);
-
-            let a: &mut Option<u64> = tx_ctx.tx_value("module.TestTxKey").unwrap();
-            assert_eq!(a, &None);
-
-            *a = Some(65);
-
-            let b: &mut Option<u64> = tx_ctx.tx_value("module.TestTxKey").unwrap();
-            assert_eq!(b, &Some(65));
-
-            let c: &mut Option<u64> = tx_ctx.tx_value("module.TestTakeTxKey").unwrap();
-            *c = Some(67);
-            let d: Box<Option<u64>> = tx_ctx.take_tx_value("module.TestTakeTxKey").unwrap();
-            assert_eq!(d, Box::new(Some(67)));
-            let e: &mut Option<u64> = tx_ctx.tx_value("module.TestTakeTxKey").unwrap();
-            assert_eq!(e, &None);
+            let mut y = tx_ctx.value::<u64>("module.TestKey");
+            let y = y.get_mut().unwrap();
+            assert_eq!(*y, 42);
+            *y = 48;
+
+            let a = tx_ctx.tx_value::<u64>("module.TestTxKey").get();
+            assert_eq!(a, None);
+            tx_ctx.tx_value::<u64>("module.TestTxKey").set(65);
+
+            let b = tx_ctx.tx_value::<u64>("module.TestTxKey").get();
+            assert_eq!(b, Some(&65));
+
+            let c = tx_ctx.tx_value::<u64>("module.TestTakeTxKey").or_default();
+            *c = 67;
+            let d = tx_ctx.tx_value::<u64>("module.TestTakeTxKey").take();
+            assert_eq!(d, Some(67));
+            let e = tx_ctx.tx_value::<u64>("module.TestTakeTxKey").get();
+            assert_eq!(e, None);
         });
 
-        let x: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(x, &Some(48));
+        let x = ctx.value::<u64>("module.TestKey").get();
+        assert_eq!(x, Some(&48));
 
         ctx.with_tx(tx, |mut tx_ctx, _call| {
-            let z: Box<Option<u64>> = tx_ctx.take_value("module.TestKey");
-            assert_eq!(z, Box::new(Some(48)));
+            let z = tx_ctx.value::<u64>("module.TestKey").take();
+            assert_eq!(z, Some(48));
 
-            let a: &mut Option<u64> = tx_ctx.tx_value("module.TestTxKey").unwrap();
-            assert_eq!(a, &None);
+            let a = tx_ctx.tx_value::<u64>("module.TestTxKey").get();
+            assert_eq!(a, None);
         });
 
-        let y: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(y, &None);
+        let y = ctx.value::<u64>("module.TestKey").get();
+        assert_eq!(y, None);
     }
 
     #[test]
@@ -631,9 +713,8 @@ mod test {
         let mut mock = Mock::default();
         let mut ctx = mock.create_ctx();
 
-        let x: &mut Option<u64> = ctx.value("module.TestKey");
-        assert_eq!(x, &None, "default value should be created");
-        *x = Some(42);
+        let x = ctx.value::<u64>("module.TestKey").set(0);
+        *x = 42;
 
         let tx = transaction::Transaction {
             version: 1,
@@ -651,7 +732,7 @@ mod test {
         };
         ctx.with_tx(tx, |mut tx_ctx, _call| {
             // Changing the type of a key should result in a panic.
-            tx_ctx.value::<Option<u32>>("module.TestKey");
+            tx_ctx.value::<Option<u32>>("module.TestKey").get();
         });
     }
 
diff --git a/runtime-sdk/src/dispatcher.rs b/runtime-sdk/src/dispatcher.rs
index 1ccfed1d8d..ceaeaa9596 100644
--- a/runtime-sdk/src/dispatcher.rs
+++ b/runtime-sdk/src/dispatcher.rs
@@ -11,8 +11,7 @@ use thiserror::Error;
 use oasis_core_runtime::{
     self,
     common::cbor,
-    consensus::roothash::MessageEvent,
-    storage::context::StorageContext,
+    storage::{context::StorageContext, mkvs},
     transaction::{
         self,
         dispatcher::{ExecuteBatchResult, ExecuteTxResult},
@@ -23,9 +22,9 @@ use oasis_core_runtime::{
 };
 
 use crate::{
-    context::{Context, DispatchContext},
+    context::{BatchContext, Context, RuntimeBatchContext},
     error::{Error as _, RuntimeError},
-    module::{AuthHandler, BlockHandler, MessageHandlerRegistry, MethodRegistry},
+    module::{self, AuthHandler, BlockHandler, MethodHandler},
     modules,
     runtime::Runtime,
     storage, types,
@@ -47,9 +46,10 @@ pub enum Error {
     MalformedTransactionInBatch(#[source] modules::core::Error),
 }
 
-struct DispatchResult {
-    result: types::transaction::CallResult,
-    tags: Tags,
+/// Result of dispatching a transaction.
+pub struct DispatchResult {
+    pub result: types::transaction::CallResult,
+    pub tags: Tags,
 }
 
 impl From<types::transaction::CallResult> for DispatchResult {
@@ -61,30 +61,25 @@ impl From<types::transaction::CallResult> for DispatchResult {
     }
 }
 
+/// The runtime dispatcher.
 pub struct Dispatcher<R: Runtime> {
-    /// Method registry.
-    methods: MethodRegistry,
-    /// Handlers registered for consensus messages.
-    consensus_message_handlers: MessageHandlerRegistry,
-
     _runtime: PhantomData<R>,
 }
 
 impl<R: Runtime> Dispatcher<R> {
-    pub(super) fn new(
-        methods: MethodRegistry,
-        consensus_message_handlers: MessageHandlerRegistry,
-    ) -> Self {
+    /// Create a new instance of the dispatcher for the given runtime.
+    ///
+    /// Note that the dispatcher is fully static and the constructor is only needed so that the
+    /// instance can be used directly with the dispatcher system provided by Oasis Core.
+    pub(super) fn new() -> Self {
         Self {
-            methods,
-            consensus_message_handlers,
             _runtime: PhantomData,
         }
     }
 
-    fn decode_tx(
-        &self,
-        ctx: &mut DispatchContext<'_>,
+    /// Decode a runtime transaction.
+    pub fn decode_tx<C: Context>(
+        ctx: &mut C,
         tx: &[u8],
     ) -> Result<types::transaction::Transaction, modules::core::Error> {
         // TODO: Check against transaction size limit.
@@ -102,9 +97,9 @@ impl<R: Runtime> Dispatcher<R> {
             .map_err(|_| modules::core::Error::MalformedTransaction)
     }
 
-    fn dispatch_tx(
-        &self,
-        ctx: &mut DispatchContext<'_>,
+    /// Dispatch a runtime transaction in the given context.
+    pub fn dispatch_tx<C: BatchContext>(
+        ctx: &mut C,
         tx: types::transaction::Transaction,
     ) -> Result<DispatchResult, Error> {
         // Run pre-processing hooks.
@@ -112,20 +107,16 @@ impl<R: Runtime> Dispatcher<R> {
             return Ok(err.to_call_result().into());
         }
 
-        // Perform transaction method lookup.
-        let method_info = match self.methods.lookup_callable(&tx.call.method) {
-            Some(method_info) => method_info,
-            None => {
-                // Method not found.
-                return Ok(modules::core::Error::InvalidMethod.to_call_result().into());
-            }
-        };
-
         let (result, messages) = ctx.with_tx(tx, |mut ctx, call| {
-            let result = (method_info.handler)(&method_info, &mut ctx, call.body);
-            if !result.is_success() {
-                return (result.into(), Vec::new());
-            }
+            let result = match R::Modules::dispatch_call(&mut ctx, &call.method, call.body) {
+                module::DispatchResult::Handled(result) => result,
+                module::DispatchResult::Unhandled(_) => {
+                    return (
+                        modules::core::Error::InvalidMethod.to_call_result().into(),
+                        Vec::new(),
+                    )
+                }
+            };
 
             // Commit store and return emitted tags and messages.
             let (tags, messages) = ctx.commit();
@@ -140,8 +131,9 @@ impl<R: Runtime> Dispatcher<R> {
         Ok(result)
     }
 
-    fn check_tx(&self, ctx: &mut DispatchContext<'_>, tx: &[u8]) -> Result<CheckTxResult, Error> {
-        let tx = match self.decode_tx(ctx, &tx) {
+    /// Check whether the given transaction is valid.
+    pub fn check_tx<C: BatchContext>(ctx: &mut C, tx: &[u8]) -> Result<CheckTxResult, Error> {
+        let tx = match Self::decode_tx(ctx, &tx) {
             Ok(tx) => tx,
             Err(err) => {
                 return Ok(CheckTxResult {
@@ -155,7 +147,7 @@ impl<R: Runtime> Dispatcher<R> {
             }
         };
 
-        match self.dispatch_tx(ctx, tx)?.result {
+        match Self::dispatch_tx(ctx, tx)?.result {
             types::transaction::CallResult::Ok(value) => Ok(CheckTxResult {
                 error: Default::default(),
                 meta: Some(value),
@@ -176,21 +168,16 @@ impl<R: Runtime> Dispatcher<R> {
         }
     }
 
-    fn execute_tx(
-        &self,
-        ctx: &mut DispatchContext<'_>,
-        tx: &[u8],
-    ) -> Result<ExecuteTxResult, Error> {
+    /// Execute the given transaction.
+    pub fn execute_tx<C: BatchContext>(ctx: &mut C, tx: &[u8]) -> Result<ExecuteTxResult, Error> {
         // It is an error to include a malformed transaction in a batch. So instead of only
         // reporting a failed execution result, we fail the whole batch. This will make the compute
         // node vote for failure and the round will fail.
         //
         // Correct proposers should only include transactions which have passed check_tx.
-        let tx = self
-            .decode_tx(ctx, &tx)
-            .map_err(Error::MalformedTransactionInBatch)?;
+        let tx = Self::decode_tx(ctx, &tx).map_err(Error::MalformedTransactionInBatch)?;
 
-        let dispatch_result = self.dispatch_tx(ctx, tx)?;
+        let dispatch_result = Self::dispatch_tx(ctx, tx)?;
 
         Ok(ExecuteTxResult {
             output: cbor::to_vec(&dispatch_result.result),
@@ -198,28 +185,7 @@ impl<R: Runtime> Dispatcher<R> {
         })
     }
 
-    fn dispatch_message(
-        &self,
-        mut ctx: &mut DispatchContext<'_>,
-        handler_name: String,
-        message_event: MessageEvent,
-        handler_ctx: cbor::Value,
-    ) -> Result<(), modules::core::Error> {
-        // Perform message handler lookup.
-        let method_info = self
-            .consensus_message_handlers
-            .lookup_handler(&handler_name)
-            .ok_or(modules::core::Error::InvalidMethod)?;
-
-        (method_info.handler)(&method_info, &mut ctx, message_event, handler_ctx);
-
-        Ok(())
-    }
-
-    fn handle_last_round_messages(
-        &self,
-        ctx: &mut DispatchContext<'_>,
-    ) -> Result<(), modules::core::Error> {
+    fn handle_last_round_messages<C: Context>(ctx: &mut C) -> Result<(), modules::core::Error> {
         let message_events = ctx.runtime_round_results().messages.clone();
 
         let store = storage::TypedStore::new(storage::PrefixStore::new(
@@ -234,7 +200,16 @@ impl<R: Runtime> Dispatcher<R> {
             let handler = handlers
                 .remove(&event.index)
                 .ok_or(modules::core::Error::MessageHandlerMissing(event.index))?;
-            self.dispatch_message(ctx, handler.hook_name, event, handler.payload)?;
+
+            R::Modules::dispatch_message_result(
+                ctx,
+                &handler.hook_name,
+                types::message::MessageResult {
+                    event,
+                    context: handler.payload,
+                },
+            )
+            .ok_or(modules::core::Error::InvalidMethod)?;
         }
 
         if !handlers.is_empty() {
@@ -246,7 +221,6 @@ impl<R: Runtime> Dispatcher<R> {
     }
 
     fn save_emitted_message_handlers<S: storage::Store>(
-        &self,
         store: S,
         handlers: Vec<types::message::MessageEventHookInvocation>,
     ) {
@@ -262,10 +236,6 @@ impl<R: Runtime> Dispatcher<R> {
         ));
         store.insert(&modules::core::state::MESSAGE_HANDLERS, &message_handlers);
     }
-
-    fn maybe_init_state(&self, ctx: &mut DispatchContext<'_>) {
-        R::migrate(ctx)
-    }
 }
 
 impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
@@ -277,12 +247,15 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
         // TODO: Get rid of StorageContext (pass mkvs in ctx).
         StorageContext::with_current(|mkvs, _| {
             // Prepare dispatch context.
-            let mut ctx = DispatchContext::from_runtime(&rt_ctx, mkvs, &self.methods);
+            let mut ctx =
+                RuntimeBatchContext::<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>>::from_runtime(
+                    &rt_ctx, mkvs,
+                );
             // Perform state migrations if required.
-            self.maybe_init_state(&mut ctx);
+            R::migrate(&mut ctx);
 
             // Handle last round message results.
-            self.handle_last_round_messages(&mut ctx)?;
+            Self::handle_last_round_messages(&mut ctx)?;
 
             // Run begin block hooks.
             R::Modules::begin_block(&mut ctx);
@@ -290,7 +263,7 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
             // Execute the batch.
             let mut results = Vec::with_capacity(batch.len());
             for tx in batch.iter() {
-                results.push(self.execute_tx(&mut ctx, &tx)?);
+                results.push(Self::execute_tx(&mut ctx, &tx)?);
             }
 
             // Run end block hooks.
@@ -301,7 +274,7 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
             let (messages, handlers) = messages.into_iter().unzip();
 
             let state = storage::MKVSStore::new(rt_ctx.io_ctx.clone(), mkvs);
-            self.save_emitted_message_handlers(state, handlers);
+            Self::save_emitted_message_handlers(state, handlers);
 
             Ok(ExecuteBatchResult {
                 results,
@@ -319,14 +292,17 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
         // TODO: Get rid of StorageContext (pass mkvs in ctx).
         StorageContext::with_current(|mkvs, _| {
             // Prepare dispatch context.
-            let mut ctx = DispatchContext::from_runtime(&ctx, mkvs, &self.methods);
+            let mut ctx =
+                RuntimeBatchContext::<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>>::from_runtime(
+                    &ctx, mkvs,
+                );
             // Perform state migrations if required.
-            self.maybe_init_state(&mut ctx);
+            R::migrate(&mut ctx);
 
             // Check the batch.
             let mut results = Vec::with_capacity(batch.len());
             for tx in batch.iter() {
-                results.push(self.check_tx(&mut ctx, &tx)?);
+                results.push(Self::check_tx(&mut ctx, &tx)?);
             }
 
             Ok(results)
@@ -346,16 +322,16 @@ impl<R: Runtime> transaction::dispatcher::Dispatcher for Dispatcher<R> {
         // TODO: Get rid of StorageContext (pass mkvs in ctx).
         StorageContext::with_current(|mkvs, _| {
             // Prepare dispatch context.
-            let mut ctx = DispatchContext::from_runtime(&ctx, mkvs, &self.methods);
+            let mut ctx =
+                RuntimeBatchContext::<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>>::from_runtime(
+                    &ctx, mkvs,
+                );
             // Perform state migrations if required.
-            self.maybe_init_state(&mut ctx);
+            R::migrate(&mut ctx);
 
             // Execute the query.
-            let method_info = self
-                .methods
-                .lookup_query(method)
-                .ok_or(modules::core::Error::InvalidMethod)?;
-            (method_info.handler)(&method_info, &mut ctx, args)
+            R::Modules::dispatch_query(&mut ctx, method, args)
+                .ok_or(modules::core::Error::InvalidMethod)?
         })
     }
 }
diff --git a/runtime-sdk/src/lib.rs b/runtime-sdk/src/lib.rs
index 8c348484ed..b6d3e296cc 100644
--- a/runtime-sdk/src/lib.rs
+++ b/runtime-sdk/src/lib.rs
@@ -15,7 +15,7 @@ pub mod testing;
 pub mod types;
 
 pub use crate::{
-    context::{Context, DispatchContext},
+    context::{BatchContext, Context, TxContext},
     core::common::version::Version,
     module::Module,
     runtime::Runtime,
diff --git a/runtime-sdk/src/module.rs b/runtime-sdk/src/module.rs
index ef1b14e036..5cff8172b0 100644
--- a/runtime-sdk/src/module.rs
+++ b/runtime-sdk/src/module.rs
@@ -1,5 +1,5 @@
 //! Runtime modules.
-use std::{collections::BTreeMap, fmt::Debug};
+use std::fmt::Debug;
 
 use impl_trait_for_tuples::impl_for_tuples;
 use serde::{de::DeserializeOwned, Serialize};
@@ -7,138 +7,113 @@ use serde::{de::DeserializeOwned, Serialize};
 use oasis_core_runtime::common::cbor;
 
 use crate::{
-    context::{DispatchContext, TxContext},
+    context::{Context, TxContext},
     error, event, modules, storage,
     storage::Store,
     types::{
-        message::MessageEvent,
+        message::MessageResult,
         transaction::{CallResult, Transaction, UnverifiedTransaction},
     },
 };
 
-/// Metadata of a callable method.
-pub struct CallableMethodInfo {
-    /// Method name.
-    pub name: &'static str,
-
-    /// Method handler function.
-    pub handler: fn(&CallableMethodInfo, &mut TxContext<'_, '_>, cbor::Value) -> CallResult,
-}
-
-/// Metadata of a query method.
-pub struct QueryMethodInfo {
-    /// Method name.
-    pub name: &'static str,
-
-    /// Method handler function.
-    pub handler: fn(
-        &QueryMethodInfo,
-        &mut DispatchContext<'_>,
-        cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError>,
-}
-
-/// Registry of methods exposed by the modules.
-pub struct MethodRegistry {
-    callable_methods: BTreeMap<&'static str, CallableMethodInfo>,
-    query_methods: BTreeMap<&'static str, QueryMethodInfo>,
+/// Result of invoking the method handler.
+pub enum DispatchResult<B, R> {
+    Handled(R),
+    Unhandled(B),
 }
 
-impl MethodRegistry {
-    /// Create a new method registry.
-    #[allow(clippy::new_without_default)]
-    pub fn new() -> Self {
-        Self {
-            callable_methods: BTreeMap::new(),
-            query_methods: BTreeMap::new(),
+impl<B, R> DispatchResult<B, R> {
+    /// Transforms `DispatchResult<B, R>` into `Result<R, E>`, mapping `Handled(r)` to `Ok(r)` and
+    /// `Unhandled(_)` to `Err(err)`.
+    pub fn ok_or<E>(self, err: E) -> Result<R, E> {
+        match self {
+            DispatchResult::Handled(result) => Ok(result),
+            DispatchResult::Unhandled(_) => Err(err),
         }
     }
+}
 
-    /// Register a new callable method.
-    ///
-    /// # Panics
-    ///
-    /// This method will panic in case a method with the same name is already registered.
-    pub fn register_callable(&mut self, mi: CallableMethodInfo) {
-        if self.callable_methods.contains_key(mi.name) {
-            panic!("callable method already exists: {}", mi.name);
-        }
-        self.callable_methods.insert(mi.name, mi);
-    }
-
-    /// Looks up a previously registered callable method.
-    pub fn lookup_callable(&self, name: &str) -> Option<&CallableMethodInfo> {
-        self.callable_methods.get(name)
+/// Method handler.
+pub trait MethodHandler {
+    /// Dispatch a call.
+    fn dispatch_call<C: TxContext>(
+        _ctx: &mut C,
+        _method: &str,
+        body: cbor::Value,
+    ) -> DispatchResult<cbor::Value, CallResult> {
+        // Default implementation indicates that the call was not handled.
+        DispatchResult::Unhandled(body)
     }
 
-    /// Register a new query method.
-    ///
-    /// # Panics
-    ///
-    /// This method will panic in case a method with the same name is already registered.
-    pub fn register_query(&mut self, mi: QueryMethodInfo) {
-        if self.query_methods.contains_key(mi.name) {
-            panic!("query method already exists: {}", mi.name);
-        }
-        self.query_methods.insert(mi.name, mi);
+    /// Dispatch a query.
+    fn dispatch_query<C: Context>(
+        _ctx: &mut C,
+        _method: &str,
+        args: cbor::Value,
+    ) -> DispatchResult<cbor::Value, Result<cbor::Value, error::RuntimeError>> {
+        // Default implementation indicates that the query was not handled.
+        DispatchResult::Unhandled(args)
     }
 
-    /// Looks up a previously registered callable method.
-    pub fn lookup_query(&self, name: &str) -> Option<&QueryMethodInfo> {
-        self.query_methods.get(name)
+    /// Dispatch a message result.
+    fn dispatch_message_result<C: Context>(
+        _ctx: &mut C,
+        _handler_name: &str,
+        result: MessageResult,
+    ) -> DispatchResult<MessageResult, ()> {
+        // Default implementation indicates that the query was not handled.
+        DispatchResult::Unhandled(result)
     }
 }
 
-/// Method registration handler.
 #[impl_for_tuples(30)]
-pub trait MethodRegistrationHandler {
-    /// Register any methods exported by the module.
-    fn register_methods(_methods: &mut MethodRegistry) {
-        // Default implementation doesn't do anything.
+impl MethodHandler for Tuple {
+    fn dispatch_call<C: TxContext>(
+        ctx: &mut C,
+        method: &str,
+        body: cbor::Value,
+    ) -> DispatchResult<cbor::Value, CallResult> {
+        // Return on first handler that can handle the method.
+        for_tuples!( #(
+            let body = match Tuple::dispatch_call::<C>(ctx, method, body) {
+                DispatchResult::Handled(result) => return DispatchResult::Handled(result),
+                DispatchResult::Unhandled(body) => body,
+            };
+        )* );
+
+        DispatchResult::Unhandled(body)
     }
-}
 
-/// Metadata of the consensus message handling callback.
-pub struct MessageHandlerInfo {
-    /// Message handler name.
-    pub name: &'static str,
-    /// Message handler.
-    pub handler: fn(&MessageHandlerInfo, &mut DispatchContext<'_>, MessageEvent, cbor::Value),
-}
-
-/// Registry of message handlers registered by the module.
-pub struct MessageHandlerRegistry {
-    handlers: BTreeMap<&'static str, MessageHandlerInfo>,
-}
-
-impl MessageHandlerRegistry {
-    /// Create a new method registry.
-    #[allow(clippy::new_without_default)]
-    pub fn new() -> Self {
-        Self {
-            handlers: BTreeMap::new(),
-        }
-    }
-
-    /// Register a message handler exposed by the module.
-    pub fn register_handler(&mut self, method_info: MessageHandlerInfo) {
-        if self.handlers.contains_key(method_info.name) {
-            panic!("message handler already exists: {}", method_info.name);
-        }
-        self.handlers.insert(method_info.name, method_info);
-    }
-
-    /// Looks up a previously registered message handler.
-    pub fn lookup_handler(&self, name: &str) -> Option<&MessageHandlerInfo> {
-        self.handlers.get(name)
+    fn dispatch_query<C: Context>(
+        ctx: &mut C,
+        method: &str,
+        args: cbor::Value,
+    ) -> DispatchResult<cbor::Value, Result<cbor::Value, error::RuntimeError>> {
+        // Return on first handler that can handle the method.
+        for_tuples!( #(
+            let args = match Tuple::dispatch_query::<C>(ctx, method, args) {
+                DispatchResult::Handled(result) => return DispatchResult::Handled(result),
+                DispatchResult::Unhandled(args) => args,
+            };
+        )* );
+
+        DispatchResult::Unhandled(args)
     }
-}
 
-#[impl_for_tuples(30)]
-pub trait MessageHookRegistrationHandler {
-    /// Register any handlers defined by the module.
-    fn register_handlers(_handlers: &mut MessageHandlerRegistry) {
-        // Default implementation doesn't do anything.
+    fn dispatch_message_result<C: Context>(
+        ctx: &mut C,
+        handler_name: &str,
+        result: MessageResult,
+    ) -> DispatchResult<MessageResult, ()> {
+        // Return on first handler that can handle the method.
+        for_tuples!( #(
+            let result = match Tuple::dispatch_message_result::<C>(ctx, handler_name, result) {
+                DispatchResult::Handled(result) => return DispatchResult::Handled(result),
+                DispatchResult::Unhandled(result) => result,
+            };
+        )* );
+
+        DispatchResult::Unhandled(result)
     }
 }
 
@@ -146,8 +121,8 @@ pub trait MessageHookRegistrationHandler {
 pub trait AuthHandler {
     /// Judge if an unverified transaction is good enough to undergo verification.
     /// This takes place before even verifying signatures.
-    fn approve_unverified_tx(
-        _ctx: &mut DispatchContext<'_>,
+    fn approve_unverified_tx<C: Context>(
+        _ctx: &mut C,
         _utx: &UnverifiedTransaction,
     ) -> Result<(), modules::core::Error> {
         // Default implementation doesn't do any checks.
@@ -157,8 +132,8 @@ pub trait AuthHandler {
     /// Authenticate a transaction.
     ///
     /// Note that any signatures have already been verified.
-    fn authenticate_tx(
-        _ctx: &mut DispatchContext<'_>,
+    fn authenticate_tx<C: Context>(
+        _ctx: &mut C,
         _tx: &Transaction,
     ) -> Result<(), modules::core::Error> {
         // Default implementation doesn't do any checks.
@@ -168,16 +143,16 @@ pub trait AuthHandler {
 
 #[impl_for_tuples(30)]
 impl AuthHandler for Tuple {
-    fn approve_unverified_tx(
-        ctx: &mut DispatchContext<'_>,
+    fn approve_unverified_tx<C: Context>(
+        ctx: &mut C,
         utx: &UnverifiedTransaction,
     ) -> Result<(), modules::core::Error> {
         for_tuples!( #( Tuple::approve_unverified_tx(ctx, utx)?; )* );
         Ok(())
     }
 
-    fn authenticate_tx(
-        ctx: &mut DispatchContext<'_>,
+    fn authenticate_tx<C: Context>(
+        ctx: &mut C,
         tx: &Transaction,
     ) -> Result<(), modules::core::Error> {
         for_tuples!( #( Tuple::authenticate_tx(ctx, tx)?; )* );
@@ -196,11 +171,14 @@ pub trait MigrationHandler {
     /// Initialize state from genesis or perform a migration.
     ///
     /// Should return true in case metadata has been changed.
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
-        meta: &mut modules::core::types::Metadata,
-        genesis: &Self::Genesis,
-    ) -> bool;
+    fn init_or_migrate<C: Context>(
+        _ctx: &mut C,
+        _meta: &mut modules::core::types::Metadata,
+        _genesis: &Self::Genesis,
+    ) -> bool {
+        // Default implementation doesn't perform any migrations.
+        false
+    }
 }
 
 #[allow(clippy::type_complexity)]
@@ -208,8 +186,8 @@ pub trait MigrationHandler {
 impl MigrationHandler for Tuple {
     for_tuples!( type Genesis = ( #( Tuple::Genesis ),* ); );
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut modules::core::types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {
@@ -224,13 +202,13 @@ impl MigrationHandler for Tuple {
 pub trait BlockHandler {
     /// Perform any common actions at the start of the block (before any transactions have been
     /// executed).
-    fn begin_block(_ctx: &mut DispatchContext<'_>) {
+    fn begin_block<C: Context>(_ctx: &mut C) {
         // Default implementation doesn't do anything.
     }
 
     /// Perform any common actions at the end of the block (after all transactions have been
     /// executed).
-    fn end_block(_ctx: &mut DispatchContext<'_>) {
+    fn end_block<C: Context>(_ctx: &mut C) {
         // Default implementation doesn't do anything.
     }
 }
diff --git a/runtime-sdk/src/modules/accounts/mod.rs b/runtime-sdk/src/modules/accounts/mod.rs
index be6d2d0003..d558e20399 100644
--- a/runtime-sdk/src/modules/accounts/mod.rs
+++ b/runtime-sdk/src/modules/accounts/mod.rs
@@ -9,11 +9,11 @@ use thiserror::Error;
 use oasis_core_runtime::common::cbor;
 
 use crate::{
-    context::{Context, DispatchContext, TxContext},
+    context::{Context, TxContext},
     crypto::signature::PublicKey,
     error::{self, Error as _},
     module,
-    module::{CallableMethodInfo, Module as _, QueryMethodInfo},
+    module::Module as _,
     modules,
     modules::core::{Module as Core, API as _},
     storage,
@@ -114,36 +114,6 @@ pub struct Genesis {
     pub total_supplies: BTreeMap<token::Denomination, token::Quantity>,
 }
 
-// TODO: Add a custom macro for easier module derivation.
-
-/*
-module!{
-    #[module(name = MODULE_NAME)]
-    impl Module {
-        type Error = Error;
-        type Event = Event;
-
-        #[module::callable_method(name = "Transfer")]
-        fn tx_transfer(ctx: &mut Context, body: u64) -> Result<(), Error> {
-            //
-            Ok(())
-        }
-
-        #[module::api]
-        fn transfer(ctx: &mut Context, body: u64) -> Result<(), Error> {
-            //
-            Ok(())
-        }
-
-        #[module::api]
-        fn mint(ctx: &mut Context, msg: messages::Mint) -> Result<(), Error> {
-            //
-            Ok(())
-        }
-    }
-}
-*/
-
 /// Interface that can be called from other modules.
 pub trait API {
     /// Transfer an amount from one account to the other.
@@ -333,7 +303,7 @@ impl API for Module {
 }
 
 impl Module {
-    fn tx_transfer(ctx: &mut TxContext<'_, '_>, body: types::Transfer) -> Result<(), Error> {
+    fn tx_transfer<C: TxContext>(ctx: &mut C, body: types::Transfer) -> Result<(), Error> {
         let params = Self::params(ctx.runtime_state());
 
         // Reject transfers when they are disabled.
@@ -341,65 +311,25 @@ impl Module {
             return Err(Error::Forbidden);
         }
 
-        Core::use_gas(ctx, params.gas_costs.tx_transfer)?;
+        Core::use_tx_gas(ctx, params.gas_costs.tx_transfer)?;
 
-        Self::transfer(
-            ctx,
-            ctx.tx_caller_address().expect("transaction context"),
-            body.to,
-            &body.amount,
-        )?;
+        Self::transfer(ctx, ctx.tx_caller_address(), body.to, &body.amount)?;
 
         Ok(())
     }
 
-    fn query_nonce(ctx: &mut DispatchContext<'_>, args: types::NonceQuery) -> Result<u64, Error> {
+    fn query_nonce<C: Context>(ctx: &mut C, args: types::NonceQuery) -> Result<u64, Error> {
         Self::get_nonce(ctx.runtime_state(), args.address)
     }
 
-    fn query_balances(
-        ctx: &mut DispatchContext<'_>,
+    fn query_balances<C: Context>(
+        ctx: &mut C,
         args: types::BalancesQuery,
     ) -> Result<types::AccountBalances, Error> {
         Self::get_balances(ctx.runtime_state(), args.address)
     }
 }
 
-impl Module {
-    fn _callable_transfer_handler(
-        _mi: &CallableMethodInfo,
-        ctx: &mut TxContext<'_, '_>,
-        body: cbor::Value,
-    ) -> CallResult {
-        let result = || -> Result<cbor::Value, Error> {
-            let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
-            Ok(cbor::to_value(&Self::tx_transfer(ctx, args)?))
-        }();
-        match result {
-            Ok(value) => CallResult::Ok(value),
-            Err(err) => err.to_call_result(),
-        }
-    }
-
-    fn _query_nonce_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext<'_>,
-        args: cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError> {
-        let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::query_nonce(ctx, args)?))
-    }
-
-    fn _query_balances_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext<'_>,
-        args: cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError> {
-        let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::query_balances(ctx, args)?))
-    }
-}
-
 impl module::Module for Module {
     const NAME: &'static str = MODULE_NAME;
     type Error = Error;
@@ -407,31 +337,49 @@ impl module::Module for Module {
     type Parameters = Parameters;
 }
 
-impl module::MessageHookRegistrationHandler for Module {}
-
-impl module::MethodRegistrationHandler for Module {
-    fn register_methods(methods: &mut module::MethodRegistry) {
-        // Callable methods.
-        methods.register_callable(module::CallableMethodInfo {
-            name: "accounts.Transfer",
-            handler: Self::_callable_transfer_handler,
-        });
+impl module::MethodHandler for Module {
+    fn dispatch_call<C: TxContext>(
+        ctx: &mut C,
+        method: &str,
+        body: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, CallResult> {
+        match method {
+            "accounts.Transfer" => {
+                let result = || -> Result<cbor::Value, Error> {
+                    let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
+                    Ok(cbor::to_value(&Self::tx_transfer(ctx, args)?))
+                }();
+                match result {
+                    Ok(value) => module::DispatchResult::Handled(CallResult::Ok(value)),
+                    Err(err) => module::DispatchResult::Handled(err.to_call_result()),
+                }
+            }
+            _ => module::DispatchResult::Unhandled(body),
+        }
+    }
 
-        // Queries.
-        methods.register_query(module::QueryMethodInfo {
-            name: "accounts.Nonce",
-            handler: Self::_query_nonce_handler,
-        });
-        methods.register_query(module::QueryMethodInfo {
-            name: "accounts.Balances",
-            handler: Self::_query_balances_handler,
-        });
+    fn dispatch_query<C: Context>(
+        ctx: &mut C,
+        method: &str,
+        args: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, Result<cbor::Value, error::RuntimeError>> {
+        match method {
+            "accounts.Nonce" => module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_nonce(ctx, args)?))
+            })()),
+            "accounts.Balances" => module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_balances(ctx, args)?))
+            })()),
+            _ => module::DispatchResult::Unhandled(args),
+        }
     }
 }
 
 impl Module {
     /// Initialize state from genesis.
-    fn init(ctx: &mut DispatchContext<'_>, genesis: &Genesis) {
+    fn init<C: Context>(ctx: &mut C, genesis: &Genesis) {
         // Create accounts.
         let mut store = storage::PrefixStore::new(ctx.runtime_state(), &MODULE_NAME);
         let mut accounts =
@@ -486,7 +434,7 @@ impl Module {
     }
 
     /// Migrate state from a previous version.
-    fn migrate(_ctx: &mut DispatchContext<'_>, _from: u32) -> bool {
+    fn migrate<C: Context>(_ctx: &mut C, _from: u32) -> bool {
         // No migrations currently supported.
         false
     }
@@ -495,8 +443,8 @@ impl Module {
 impl module::MigrationHandler for Module {
     type Genesis = Genesis;
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut modules::core::types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {
@@ -534,8 +482,8 @@ impl FeeAccumulator {
 const CONTEXT_KEY_FEE_ACCUMULATOR: &str = "accounts.FeeAccumulator";
 
 impl module::AuthHandler for Module {
-    fn authenticate_tx(
-        ctx: &mut DispatchContext<'_>,
+    fn authenticate_tx<C: Context>(
+        ctx: &mut C,
         tx: &Transaction,
     ) -> Result<(), modules::core::Error> {
         // Fetch information about each signer.
@@ -569,6 +517,7 @@ impl module::AuthHandler for Module {
                 .map_err(|_| modules::core::Error::InsufficientFeeBalance)?;
 
             ctx.value::<FeeAccumulator>(CONTEXT_KEY_FEE_ACCUMULATOR)
+                .or_default()
                 .add(&tx.auth_info.fee.amount);
 
             // TODO: Emit event that fee has been paid.
@@ -578,7 +527,7 @@ impl module::AuthHandler for Module {
 }
 
 impl module::BlockHandler for Module {
-    fn end_block(ctx: &mut DispatchContext<'_>) {
+    fn end_block<C: Context>(ctx: &mut C) {
         // Determine the fees that are available for disbursement from the last block.
         let mut previous_fees = Self::get_balances(ctx.runtime_state(), *ADDRESS_FEE_ACCUMULATOR)
             .expect("get_balances must succeed")
@@ -635,7 +584,10 @@ impl module::BlockHandler for Module {
         }
 
         // Fees for the active block should be transferred to the fee accumulator address.
-        let acc = ctx.take_value::<FeeAccumulator>(CONTEXT_KEY_FEE_ACCUMULATOR);
+        let acc = ctx
+            .value::<FeeAccumulator>(CONTEXT_KEY_FEE_ACCUMULATOR)
+            .take()
+            .unwrap_or_default();
         for (denom, amount) in acc.total_fees.into_iter() {
             Self::add_amount(
                 ctx.runtime_state(),
diff --git a/runtime-sdk/src/modules/consensus/mod.rs b/runtime-sdk/src/modules/consensus/mod.rs
index 6b47c926b6..abc4fb7de1 100644
--- a/runtime-sdk/src/modules/consensus/mod.rs
+++ b/runtime-sdk/src/modules/consensus/mod.rs
@@ -14,7 +14,7 @@ use oasis_core_runtime::consensus::{
 };
 
 use crate::{
-    context::{Context, DispatchContext, TxContext},
+    context::{Context, TxContext},
     crypto::signature::PublicKey,
     module,
     module::Module as _,
@@ -122,7 +122,7 @@ pub trait API {
     fn consensus_denomination<C: Context>(ctx: &mut C) -> Result<token::Denomination, Error>;
 
     /// Ensures transaction signer is consensus compatible.
-    fn ensure_compatible_tx_signer(ctx: &TxContext<'_, '_>) -> Result<(), Error>;
+    fn ensure_compatible_tx_signer<C: TxContext>(ctx: &C) -> Result<(), Error>;
 
     /// Query consensus account info.
     fn account<C: Context>(ctx: &C, addr: Address) -> Result<ConsensusAccount, Error>;
@@ -237,13 +237,8 @@ impl API for Module {
         Ok(params.consensus_denomination)
     }
 
-    fn ensure_compatible_tx_signer(ctx: &TxContext<'_, '_>) -> Result<(), Error> {
-        match ctx
-            .tx_auth_info()
-            .expect("should be called with a transaction ctx")
-            .signer_info[0]
-            .address_spec
-        {
+    fn ensure_compatible_tx_signer<C: TxContext>(ctx: &C) -> Result<(), Error> {
+        match ctx.tx_auth_info().signer_info[0].address_spec {
             AddressSpec::Signature(PublicKey::Ed25519(_)) => Ok(()),
             _ => Err(Error::ConsensusIncompatibleSigner),
         }
@@ -265,14 +260,13 @@ impl module::Module for Module {
     type Parameters = Parameters;
 }
 
-impl module::MethodRegistrationHandler for Module {}
-impl module::MessageHookRegistrationHandler for Module {}
+impl module::MethodHandler for Module {}
 
 impl module::MigrationHandler for Module {
     type Genesis = Genesis;
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut modules::core::types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {
diff --git a/runtime-sdk/src/modules/consensus_accounts/mod.rs b/runtime-sdk/src/modules/consensus_accounts/mod.rs
index d52e5d2fd1..60d6224f56 100644
--- a/runtime-sdk/src/modules/consensus_accounts/mod.rs
+++ b/runtime-sdk/src/modules/consensus_accounts/mod.rs
@@ -8,15 +8,15 @@ use thiserror::Error;
 use oasis_core_runtime::{common::cbor, consensus::staking::Account as ConsensusAccount};
 
 use crate::{
-    context::{Context, DispatchContext, TxContext},
+    context::{Context, TxContext},
     error::{self, Error as _},
     module,
-    module::{CallableMethodInfo, Module as _, QueryMethodInfo},
+    module::Module as _,
     modules,
     modules::core::{Module as Core, API as _},
     types::{
         address::Address,
-        message::{MessageEvent, MessageEventHookInvocation},
+        message::{MessageEvent, MessageEventHookInvocation, MessageResult},
         token,
         transaction::CallResult,
     },
@@ -192,14 +192,11 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
     Module<Accounts, Consensus>
 {
     /// Deposit in the runtime.
-    fn tx_deposit(ctx: &mut TxContext<'_, '_>, body: types::Deposit) -> Result<(), Error> {
+    fn tx_deposit<C: TxContext>(ctx: &mut C, body: types::Deposit) -> Result<(), Error> {
         let params = Self::params(ctx.runtime_state());
-        Core::use_gas(ctx, params.gas_costs.tx_deposit)?;
+        Core::use_tx_gas(ctx, params.gas_costs.tx_deposit)?;
 
-        let signer = &ctx
-            .tx_auth_info()
-            .expect("should be called with a transaction ctx")
-            .signer_info[0];
+        let signer = &ctx.tx_auth_info().signer_info[0];
         Consensus::ensure_compatible_tx_signer(ctx)?;
 
         let address = signer.address_spec.address();
@@ -207,23 +204,20 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
     }
 
     /// Withdraw from the runtime.
-    fn tx_withdraw(ctx: &mut TxContext<'_, '_>, body: types::Withdraw) -> Result<(), Error> {
+    fn tx_withdraw<C: TxContext>(ctx: &mut C, body: types::Withdraw) -> Result<(), Error> {
         let params = Self::params(ctx.runtime_state());
-        Core::use_gas(ctx, params.gas_costs.tx_withdraw)?;
+        Core::use_tx_gas(ctx, params.gas_costs.tx_withdraw)?;
 
         // Signer.
-        let signer = &ctx
-            .tx_auth_info()
-            .expect("should be called with a transaction context")
-            .signer_info[0];
+        let signer = &ctx.tx_auth_info().signer_info[0];
         Consensus::ensure_compatible_tx_signer(ctx)?;
 
         let address = signer.address_spec.address();
         Self::withdraw(ctx, address, body.amount)
     }
 
-    fn query_balance(
-        ctx: &mut DispatchContext<'_>,
+    fn query_balance<C: Context>(
+        ctx: &mut C,
         args: types::BalanceQuery,
     ) -> Result<types::AccountBalance, Error> {
         let denomination = Consensus::consensus_denomination(ctx)?;
@@ -238,63 +232,39 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
         })
     }
 
-    fn query_consensus_account(
-        ctx: &mut DispatchContext<'_>,
+    fn query_consensus_account<C: Context>(
+        ctx: &mut C,
         args: types::ConsensusAccountQuery,
     ) -> Result<ConsensusAccount, Error> {
         Consensus::account(ctx, args.address).map_err(|_| Error::InvalidArgument)
     }
-}
 
-impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
-    Module<Accounts, Consensus>
-{
-    fn _callable_deposit_handler(
-        _mi: &CallableMethodInfo,
-        ctx: &mut TxContext<'_, '_>,
-        body: cbor::Value,
-    ) -> CallResult {
-        let result = || -> Result<cbor::Value, Error> {
-            let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
-            Ok(cbor::to_value(&Self::tx_deposit(ctx, args)?))
-        }();
-        match result {
-            Ok(value) => CallResult::Ok(value),
-            Err(err) => err.to_call_result(),
+    fn message_result_transfer<C: Context>(
+        ctx: &mut C,
+        me: MessageEvent,
+        context: types::ConsensusTransferContext,
+    ) {
+        if !me.is_success() {
+            // Transfer out failed.
+            return;
         }
-    }
 
-    fn _callable_withdraw_handler(
-        _mi: &CallableMethodInfo,
-        ctx: &mut TxContext<'_, '_>,
-        body: cbor::Value,
-    ) -> CallResult {
-        let result = || -> Result<cbor::Value, Error> {
-            let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
-            Ok(cbor::to_value(&Self::tx_withdraw(ctx, args)?))
-        }();
-        match result {
-            Ok(value) => CallResult::Ok(value),
-            Err(err) => err.to_call_result(),
-        }
+        // Update runtime state.
+        Accounts::burn(ctx, context.address, &context.amount).expect("should have enough balance");
     }
 
-    fn _query_balance_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext<'_>,
-        args: cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError> {
-        let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::query_balance(ctx, args)?))
-    }
+    fn message_result_withdraw<C: Context>(
+        ctx: &mut C,
+        me: MessageEvent,
+        context: types::ConsensusWithdrawContext,
+    ) {
+        if !me.is_success() {
+            // Transfer out failed.
+            return;
+        }
 
-    fn _query_consensus_account_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext<'_>,
-        args: cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError> {
-        let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::query_consensus_account(ctx, args)?))
+        // Update runtime state.
+        Accounts::mint(ctx, context.address, &context.amount).unwrap();
     }
 }
 
@@ -309,84 +279,81 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> modul
 }
 
 /// Module methods.
-impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
-    module::MethodRegistrationHandler for Module<Accounts, Consensus>
-{
-    fn register_methods(methods: &mut module::MethodRegistry) {
-        // Callable methods.
-        methods.register_callable(module::CallableMethodInfo {
-            name: "consensus.Deposit",
-            handler: Self::_callable_deposit_handler,
-        });
-        methods.register_callable(module::CallableMethodInfo {
-            name: "consensus.Withdraw",
-            handler: Self::_callable_withdraw_handler,
-        });
-
-        // Queries.
-        methods.register_query(module::QueryMethodInfo {
-            name: "consensus.Balance",
-            handler: Self::_query_balance_handler,
-        });
-
-        methods.register_query(module::QueryMethodInfo {
-            name: "consensus.Account",
-            handler: Self::_query_consensus_account_handler,
-        })
-    }
-}
-
-impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
-    Module<Accounts, Consensus>
+impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> module::MethodHandler
+    for Module<Accounts, Consensus>
 {
-    fn _consensus_transfer_handler(
-        _mi: &module::MessageHandlerInfo,
-        ctx: &mut DispatchContext<'_>,
-        me: MessageEvent,
-        h_ctx: cbor::Value,
-    ) {
-        let h_ctx: types::ConsensusTransferContext =
-            cbor::from_value(h_ctx).expect("invalid message handler context");
-
-        // Transfer out succeed.
-        if me.is_success() {
-            // Update runtime state.
-            Accounts::burn(ctx, h_ctx.address, &h_ctx.amount).expect("should have enough balance");
+    fn dispatch_call<C: TxContext>(
+        ctx: &mut C,
+        method: &str,
+        body: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, CallResult> {
+        match method {
+            "consensus.Deposit" => {
+                let result = || -> Result<cbor::Value, Error> {
+                    let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
+                    Ok(cbor::to_value(&Self::tx_deposit(ctx, args)?))
+                }();
+                match result {
+                    Ok(value) => module::DispatchResult::Handled(CallResult::Ok(value)),
+                    Err(err) => module::DispatchResult::Handled(err.to_call_result()),
+                }
+            }
+            "consensus.Withdraw" => {
+                let result = || -> Result<cbor::Value, Error> {
+                    let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
+                    Ok(cbor::to_value(&Self::tx_withdraw(ctx, args)?))
+                }();
+                match result {
+                    Ok(value) => module::DispatchResult::Handled(CallResult::Ok(value)),
+                    Err(err) => module::DispatchResult::Handled(err.to_call_result()),
+                }
+            }
+            _ => module::DispatchResult::Unhandled(body),
         }
     }
 
-    fn _consensus_withdraw_handler(
-        _mi: &module::MessageHandlerInfo,
-        ctx: &mut DispatchContext<'_>,
-        me: MessageEvent,
-        h_ctx: cbor::Value,
-    ) {
-        let h_ctx: types::ConsensusWithdrawContext =
-            cbor::from_value(h_ctx).expect("invalid message handler context");
-
-        // Deposit in succeed.
-        if me.is_success() {
-            // Update runtime state.
-            Accounts::mint(ctx, h_ctx.address, &h_ctx.amount).unwrap();
+    fn dispatch_query<C: Context>(
+        ctx: &mut C,
+        method: &str,
+        args: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, Result<cbor::Value, error::RuntimeError>> {
+        match method {
+            "consensus.Balance" => module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_balance(ctx, args)?))
+            })()),
+            "consensus.Account" => module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_consensus_account(ctx, args)?))
+            })()),
+            _ => module::DispatchResult::Unhandled(args),
         }
     }
-}
 
-/// Module message handlers.
-impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
-    module::MessageHookRegistrationHandler for Module<Accounts, Consensus>
-{
-    // Register message handlers.
-    fn register_handlers(handlers: &mut module::MessageHandlerRegistry) {
-        handlers.register_handler(module::MessageHandlerInfo {
-            name: CONSENSUS_TRANSFER_HANDLER,
-            handler: Self::_consensus_transfer_handler,
-        });
-
-        handlers.register_handler(module::MessageHandlerInfo {
-            name: CONSENSUS_WITHDRAW_HANDLER,
-            handler: Self::_consensus_withdraw_handler,
-        });
+    fn dispatch_message_result<C: Context>(
+        ctx: &mut C,
+        handler_name: &str,
+        result: MessageResult,
+    ) -> module::DispatchResult<MessageResult, ()> {
+        match handler_name {
+            CONSENSUS_TRANSFER_HANDLER => {
+                Self::message_result_transfer(
+                    ctx,
+                    result.event,
+                    cbor::from_value(result.context).expect("invalid message handler context"),
+                );
+                module::DispatchResult::Handled(())
+            }
+            CONSENSUS_WITHDRAW_HANDLER => {
+                Self::message_result_withdraw(
+                    ctx,
+                    result.event,
+                    cbor::from_value(result.context).expect("invalid message handler context"),
+                );
+                module::DispatchResult::Handled(())
+            }
+            _ => module::DispatchResult::Unhandled(result),
+        }
     }
 }
 
@@ -395,8 +362,8 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> modul
 {
     type Genesis = Genesis;
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut modules::core::types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {
diff --git a/runtime-sdk/src/modules/core/mod.rs b/runtime-sdk/src/modules/core/mod.rs
index 464da08668..358e1a69a4 100644
--- a/runtime-sdk/src/modules/core/mod.rs
+++ b/runtime-sdk/src/modules/core/mod.rs
@@ -5,9 +5,9 @@ use thiserror::Error;
 use oasis_core_runtime::common::cbor;
 
 use crate::{
-    context::{Context, DispatchContext},
-    error,
-    module::{self, Module as _, QueryMethodInfo},
+    context::{Context, TxContext},
+    dispatcher, error,
+    module::{self, Module as _},
     types::transaction::{self, AuthProof, UnverifiedTransaction},
 };
 
@@ -99,12 +99,15 @@ impl module::Parameters for Parameters {
 }
 
 pub trait API {
-    /// Attempt to use gas. Gas limits are per-batch (max_batch_gas in Parameters) and
-    /// per-transaction (.ai.fee.gas). If the gas specified would cause either total used to exceed
+    /// Attempt to use gas. If the gas specified would cause either total used to exceed
     /// its limit, fails with Error::OutOfGas or Error::BatchOutOfGas, and neither gas usage is
-    /// increased. Per-transaction gas is not assessed when C is not a transaction processing
-    /// context.
-    fn use_gas<C: Context>(ctx: &mut C, gas: u64) -> Result<(), Error>;
+    /// increased.
+    fn use_batch_gas<C: Context>(ctx: &mut C, gas: u64) -> Result<(), Error>;
+
+    /// Attempt to use gas. If the gas specified would cause either total used to exceed
+    /// its limit, fails with Error::OutOfGas or Error::BatchOutOfGas, and neither gas usage is
+    /// increased.
+    fn use_tx_gas<C: TxContext>(ctx: &mut C, gas: u64) -> Result<(), Error>;
 }
 
 /// Genesis state for the accounts module.
@@ -129,37 +132,22 @@ const CONTEXT_KEY_GAS_USED: &str = "core.GasUsed";
 
 impl Module {
     /// Initialize state from genesis.
-    fn init(ctx: &mut DispatchContext<'_>, genesis: &Genesis) {
+    fn init<C: Context>(ctx: &mut C, genesis: &Genesis) {
         // Set genesis parameters.
         Self::set_params(ctx.runtime_state(), &genesis.parameters);
     }
 
     /// Migrate state from a previous version.
-    fn migrate(_ctx: &mut DispatchContext<'_>, _from: u32) -> bool {
+    fn migrate<C: Context>(_ctx: &mut C, _from: u32) -> bool {
         // No migrations currently supported.
         false
     }
 }
 
 impl API for Module {
-    fn use_gas<C: Context>(ctx: &mut C, gas: u64) -> Result<(), Error> {
-        let new_gas_used = match (
-            ctx.tx_auth_info().map(|ai| ai.fee.gas),
-            ctx.tx_value::<u64>(CONTEXT_KEY_GAS_USED).copied(),
-        ) {
-            (Some(gas_limit), Some(gas_used)) => {
-                let sum = gas_used.checked_add(gas).ok_or(Error::GasOverflow)?;
-                if sum > gas_limit {
-                    return Err(Error::OutOfGas);
-                }
-                Some(sum)
-            }
-            (None, None) => None,
-            _ => panic!("inconsistent tx availability"),
-        };
-
+    fn use_batch_gas<C: Context>(ctx: &mut C, gas: u64) -> Result<(), Error> {
         let batch_gas_limit = Self::params(ctx.runtime_state()).max_batch_gas;
-        let batch_gas_used = ctx.value::<u64>(CONTEXT_KEY_GAS_USED);
+        let batch_gas_used = ctx.value::<u64>(CONTEXT_KEY_GAS_USED).or_default();
         let batch_new_gas_used = batch_gas_used
             .checked_add(gas)
             .ok_or(Error::BatchGasOverflow)?;
@@ -167,13 +155,26 @@ impl API for Module {
             return Err(Error::BatchOutOfGas);
         }
 
-        match (new_gas_used, ctx.tx_value::<u64>(CONTEXT_KEY_GAS_USED)) {
-            (Some(sum), Some(gas_used)) => *gas_used = sum,
-            (None, None) => {}
-            _ => panic!("inconsistent tx availability"),
-        }
-        let batch_gas_used = ctx.value::<u64>(CONTEXT_KEY_GAS_USED);
-        *batch_gas_used = batch_new_gas_used;
+        ctx.value::<u64>(CONTEXT_KEY_GAS_USED)
+            .set(batch_new_gas_used);
+
+        Ok(())
+    }
+
+    fn use_tx_gas<C: TxContext>(ctx: &mut C, gas: u64) -> Result<(), Error> {
+        let gas_limit = ctx.tx_auth_info().fee.gas;
+        let gas_used = ctx.tx_value::<u64>(CONTEXT_KEY_GAS_USED).or_default();
+        let new_gas_used = {
+            let sum = gas_used.checked_add(gas).ok_or(Error::GasOverflow)?;
+            if sum > gas_limit {
+                return Err(Error::OutOfGas);
+            }
+            sum
+        };
+
+        Self::use_batch_gas(ctx, gas)?;
+
+        *ctx.tx_value::<u64>(CONTEXT_KEY_GAS_USED).or_default() = new_gas_used;
 
         Ok(())
     }
@@ -185,34 +186,19 @@ impl Module {
     /// estimate that and return successfully, so do not use this query to see if a transaction will
     /// succeed. Failure due to OutOfGas are included, so it's best to set the query argument
     /// transaction's gas to something high.
-    fn query_estimate_gas(
-        ctx: &mut DispatchContext<'_>,
+    fn query_estimate_gas<C: Context>(
+        ctx: &mut C,
         args: transaction::Transaction,
     ) -> Result<u64, Error> {
-        let mi = ctx
-            .methods
-            .lookup_callable(&args.call.method)
-            .ok_or(Error::InvalidMethod)?;
-        ctx.with_tx_simulation(args, |mut tx_ctx, call| {
-            (mi.handler)(&mi, &mut tx_ctx, call.body);
+        ctx.with_simulation(|mut ctx| {
+            dispatcher::Dispatcher::<C::Runtime>::dispatch_tx(&mut ctx, args).ok();
             // Warning: we don't report success or failure. If the call fails, we still report
             // how much gas it uses while it fails.
-            Ok(*tx_ctx.tx_value::<u64>(CONTEXT_KEY_GAS_USED).unwrap())
+            Ok(*ctx.value::<u64>(CONTEXT_KEY_GAS_USED).or_default())
         })
     }
 }
 
-impl Module {
-    fn _query_estimate_gas_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext<'_>,
-        args: cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError> {
-        let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::query_estimate_gas(ctx, args)?))
-    }
-}
-
 impl module::Module for Module {
     const NAME: &'static str = MODULE_NAME;
     type Error = Error;
@@ -221,8 +207,8 @@ impl module::Module for Module {
 }
 
 impl module::AuthHandler for Module {
-    fn approve_unverified_tx(
-        ctx: &mut DispatchContext<'_>,
+    fn approve_unverified_tx<C: Context>(
+        ctx: &mut C,
         utx: &UnverifiedTransaction,
     ) -> Result<(), Error> {
         let params = Self::params(ctx.runtime_state());
@@ -243,8 +229,8 @@ impl module::AuthHandler for Module {
 impl module::MigrationHandler for Module {
     type Genesis = Genesis;
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {
@@ -261,16 +247,20 @@ impl module::MigrationHandler for Module {
     }
 }
 
-impl module::MethodRegistrationHandler for Module {
-    fn register_methods(methods: &mut module::MethodRegistry) {
-        // Queries.
-        methods.register_query(module::QueryMethodInfo {
-            name: "core.EstimateGas",
-            handler: Self::_query_estimate_gas_handler,
-        });
+impl module::MethodHandler for Module {
+    fn dispatch_query<C: Context>(
+        ctx: &mut C,
+        method: &str,
+        args: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, Result<cbor::Value, error::RuntimeError>> {
+        match method {
+            "core.EstimateGas" => module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_estimate_gas(ctx, args)?))
+            })()),
+            _ => module::DispatchResult::Unhandled(args),
+        }
     }
 }
 
 impl module::BlockHandler for Module {}
-
-impl module::MessageHookRegistrationHandler for Module {}
diff --git a/runtime-sdk/src/modules/rewards/mod.rs b/runtime-sdk/src/modules/rewards/mod.rs
index a625f6c7b2..d37e99f1ff 100644
--- a/runtime-sdk/src/modules/rewards/mod.rs
+++ b/runtime-sdk/src/modules/rewards/mod.rs
@@ -7,11 +7,11 @@ use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::{
-    context::{Context, DispatchContext},
+    context::Context,
     core::{common::cbor, consensus::beacon},
     crypto::signature::PublicKey,
     error,
-    module::{self, Module as _, Parameters as _, QueryMethodInfo},
+    module::{self, Module as _, Parameters as _},
     modules, storage,
     types::address::Address,
 };
@@ -96,22 +96,11 @@ pub static ADDRESS_REWARD_POOL: Lazy<Address> =
     Lazy::new(|| Address::from_module(MODULE_NAME, "reward-pool"));
 
 impl<Accounts: modules::accounts::API> Module<Accounts> {
-    fn query_parameters(ctx: &mut DispatchContext<'_>, _args: ()) -> Result<Parameters, Error> {
+    fn query_parameters<C: Context>(ctx: &mut C, _args: ()) -> Result<Parameters, Error> {
         Ok(Self::params(ctx.runtime_state()))
     }
 }
 
-impl<Accounts: modules::accounts::API> Module<Accounts> {
-    fn _query_parameters_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext<'_>,
-        args: cbor::Value,
-    ) -> Result<cbor::Value, error::RuntimeError> {
-        let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::query_parameters(ctx, args)?))
-    }
-}
-
 impl<Accounts: modules::accounts::API> module::Module for Module<Accounts> {
     const NAME: &'static str = MODULE_NAME;
     type Error = Infallible;
@@ -119,21 +108,25 @@ impl<Accounts: modules::accounts::API> module::Module for Module<Accounts> {
     type Parameters = Parameters;
 }
 
-impl<Accounts: modules::accounts::API> module::MessageHookRegistrationHandler for Module<Accounts> {}
-
-impl<Accounts: modules::accounts::API> module::MethodRegistrationHandler for Module<Accounts> {
-    fn register_methods(methods: &mut module::MethodRegistry) {
-        // Queries.
-        methods.register_query(module::QueryMethodInfo {
-            name: "rewards.Parameters",
-            handler: Self::_query_parameters_handler,
-        });
+impl<Accounts: modules::accounts::API> module::MethodHandler for Module<Accounts> {
+    fn dispatch_query<C: Context>(
+        ctx: &mut C,
+        method: &str,
+        args: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, Result<cbor::Value, error::RuntimeError>> {
+        match method {
+            "rewards.Parameters" => module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_parameters(ctx, args)?))
+            })()),
+            _ => module::DispatchResult::Unhandled(args),
+        }
     }
 }
 
 impl<Accounts: modules::accounts::API> Module<Accounts> {
     /// Initialize state from genesis.
-    fn init(ctx: &mut DispatchContext<'_>, genesis: &Genesis) {
+    fn init<C: Context>(ctx: &mut C, genesis: &Genesis) {
         genesis
             .parameters
             .validate_basic()
@@ -144,7 +137,7 @@ impl<Accounts: modules::accounts::API> Module<Accounts> {
     }
 
     /// Migrate state from a previous version.
-    fn migrate(_ctx: &mut DispatchContext<'_>, _from: u32) -> bool {
+    fn migrate<C: Context>(_ctx: &mut C, _from: u32) -> bool {
         // No migrations currently supported.
         false
     }
@@ -153,8 +146,8 @@ impl<Accounts: modules::accounts::API> Module<Accounts> {
 impl<Accounts: modules::accounts::API> module::MigrationHandler for Module<Accounts> {
     type Genesis = Genesis;
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext<'_>,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut modules::core::types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {
@@ -174,7 +167,7 @@ impl<Accounts: modules::accounts::API> module::MigrationHandler for Module<Accou
 impl<Accounts: modules::accounts::API> module::AuthHandler for Module<Accounts> {}
 
 impl<Accounts: modules::accounts::API> module::BlockHandler for Module<Accounts> {
-    fn end_block(ctx: &mut DispatchContext<'_>) {
+    fn end_block<C: Context>(ctx: &mut C) {
         let epoch = ctx.epoch();
 
         // Load previous epoch.
diff --git a/runtime-sdk/src/runtime.rs b/runtime-sdk/src/runtime.rs
index fd865a974d..40ae7e408b 100644
--- a/runtime-sdk/src/runtime.rs
+++ b/runtime-sdk/src/runtime.rs
@@ -6,12 +6,9 @@ use oasis_core_runtime::{
 };
 
 use crate::{
-    context::{Context, DispatchContext},
+    context::Context,
     crypto, dispatcher,
-    module::{
-        AuthHandler, BlockHandler, MessageHandlerRegistry, MessageHookRegistrationHandler,
-        MethodRegistrationHandler, MethodRegistry, MigrationHandler,
-    },
+    module::{AuthHandler, BlockHandler, MethodHandler, MigrationHandler},
     modules, storage,
 };
 
@@ -20,17 +17,13 @@ pub trait Runtime {
     /// Runtime version.
     const VERSION: version::Version;
 
-    type Modules: AuthHandler
-        + MigrationHandler
-        + MethodRegistrationHandler
-        + BlockHandler
-        + MessageHookRegistrationHandler;
+    type Modules: AuthHandler + MigrationHandler + MethodHandler + BlockHandler;
 
     /// Genesis state for the runtime.
     fn genesis_state() -> <Self::Modules as MigrationHandler>::Genesis;
 
     /// Perform state migrations if required.
-    fn migrate(ctx: &mut DispatchContext<'_>) {
+    fn migrate<C: Context>(ctx: &mut C) {
         let store = storage::TypedStore::new(storage::PrefixStore::new(
             ctx.runtime_state(),
             &modules::core::MODULE_NAME,
@@ -72,13 +65,7 @@ pub trait Runtime {
             );
 
             // Register runtime's methods.
-            let mut methods = MethodRegistry::new();
-            Self::Modules::register_methods(&mut methods);
-
-            let mut handlers = MessageHandlerRegistry::new();
-            Self::Modules::register_handlers(&mut handlers);
-
-            let dispatcher = dispatcher::Dispatcher::<Self>::new(methods, handlers);
+            let dispatcher = dispatcher::Dispatcher::<Self>::new();
             Some(Box::new(dispatcher))
         };
 
diff --git a/runtime-sdk/src/types/message.rs b/runtime-sdk/src/types/message.rs
index 3f95470b84..7b929f619b 100644
--- a/runtime-sdk/src/types/message.rs
+++ b/runtime-sdk/src/types/message.rs
@@ -9,8 +9,12 @@ pub type MessageEvent = consensus::roothash::MessageEvent;
 
 /// Handler name and context to be called after message is executed.
 #[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
 pub struct MessageEventHookInvocation {
+    #[serde(rename = "hook_name")]
     pub hook_name: String,
+
+    #[serde(rename = "payload")]
     pub payload: cbor::Value,
 }
 
@@ -23,3 +27,11 @@ impl MessageEventHookInvocation {
         }
     }
 }
+
+/// Result of a message being processed by the consensus layer combined with the context for the
+/// result handler.
+#[derive(Clone, Debug)]
+pub struct MessageResult {
+    pub event: MessageEvent,
+    pub context: cbor::Value,
+}


Test patch:
diff --git a/runtime-sdk/src/modules/accounts/test.rs b/runtime-sdk/src/modules/accounts/test.rs
index 020bb1aaba..9918b684ed 100644
--- a/runtime-sdk/src/modules/accounts/test.rs
+++ b/runtime-sdk/src/modules/accounts/test.rs
@@ -4,7 +4,7 @@ use std::collections::BTreeMap;
 use oasis_core_runtime::common::cbor;
 
 use crate::{
-    context::{Context, DispatchContext},
+    context::{BatchContext, Context},
     module::{AuthHandler, BlockHandler},
     modules::core,
     testing::{keys, mock},
@@ -199,7 +199,7 @@ fn test_api_tx_transfer_disabled() {
     });
 }
 
-pub(crate) fn init_accounts(ctx: &mut DispatchContext<'_>) {
+pub(crate) fn init_accounts<C: Context>(ctx: &mut C) {
     Accounts::init(
         ctx,
         &Genesis {
diff --git a/runtime-sdk/src/modules/consensus/test.rs b/runtime-sdk/src/modules/consensus/test.rs
index 8d89f40c43..7d4476aa06 100644
--- a/runtime-sdk/src/modules/consensus/test.rs
+++ b/runtime-sdk/src/modules/consensus/test.rs
@@ -11,7 +11,7 @@ use oasis_core_runtime::{
 };
 
 use crate::{
-    context::Context,
+    context::{BatchContext, Context},
     modules::consensus::Module as Consensus,
     testing::{keys, mock},
     types::{
diff --git a/runtime-sdk/src/modules/consensus_accounts/test.rs b/runtime-sdk/src/modules/consensus_accounts/test.rs
index ed4b12a8f0..ef51ebe6ca 100644
--- a/runtime-sdk/src/modules/consensus_accounts/test.rs
+++ b/runtime-sdk/src/modules/consensus_accounts/test.rs
@@ -6,6 +6,7 @@ use oasis_core_runtime::consensus::{
 };
 
 use crate::{
+    context::BatchContext,
     module::MigrationHandler,
     modules::{
         accounts::{Genesis as AccountsGenesis, Module as Accounts, API},
@@ -324,23 +325,13 @@ fn test_consensus_transfer_handler() {
     );
     Module::<Accounts, Consensus>::init_or_migrate(&mut ctx, &mut meta, &Default::default());
 
-    // Invoke message handler.
-    let mi = module::MessageHandlerInfo {
-        name: CONSENSUS_TRANSFER_HANDLER,
-        handler: Module::<Accounts, Consensus>::_consensus_transfer_handler,
-    };
-    // Successful event.
+    // Simulate successful event.
     let me = Default::default();
     let h_ctx = types::ConsensusTransferContext {
         address: keys::alice::address(),
         amount: BaseUnits::new(999_999.into(), denom.clone()),
     };
-    Module::<Accounts, Consensus>::_consensus_transfer_handler(
-        &mi,
-        &mut ctx,
-        me,
-        cbor::to_value(h_ctx),
-    );
+    Module::<Accounts, Consensus>::message_result_transfer(&mut ctx, me, h_ctx);
 
     // Ensure runtime balance is updated.
     let bals = Accounts::get_balances(ctx.runtime_state(), keys::alice::address()).unwrap();
@@ -384,23 +375,13 @@ fn test_consensus_withdraw_handler() {
     );
     Module::<Accounts, Consensus>::init_or_migrate(&mut ctx, &mut meta, &Default::default());
 
-    // Invoke message handler.
-    let mi = module::MessageHandlerInfo {
-        name: CONSENSUS_WITHDRAW_HANDLER,
-        handler: Module::<Accounts, Consensus>::_consensus_withdraw_handler,
-    };
-    // Successful event.
+    // Simulate successful event.
     let me = Default::default();
     let h_ctx = types::ConsensusWithdrawContext {
         address: keys::alice::address(),
         amount: BaseUnits::new(1.into(), denom.clone()),
     };
-    Module::<Accounts, Consensus>::_consensus_withdraw_handler(
-        &mi,
-        &mut ctx,
-        me,
-        cbor::to_value(h_ctx),
-    );
+    Module::<Accounts, Consensus>::message_result_withdraw(&mut ctx, me, h_ctx);
 
     // Ensure runtime balance is updated.
     let bals = Accounts::get_balances(ctx.runtime_state(), keys::alice::address()).unwrap();
diff --git a/runtime-sdk/src/modules/core/test.rs b/runtime-sdk/src/modules/core/test.rs
index 45d469f309..f081114855 100644
--- a/runtime-sdk/src/modules/core/test.rs
+++ b/runtime-sdk/src/modules/core/test.rs
@@ -1,9 +1,9 @@
-use oasis_core_runtime::common::cbor;
-
 use crate::{
-    context::{Context, Mode},
+    context::{BatchContext, Context, Mode, TxContext},
+    core::common::{cbor, version::Version},
     module,
     module::{AuthHandler as _, Module as _},
+    runtime::Runtime,
     testing::{keys, mock},
     types::{token, transaction},
 };
@@ -25,70 +25,113 @@ fn test_use_gas() {
         },
     );
 
-    Core::use_gas(&mut ctx, 1).expect("using batch gas under limit should succeed");
+    Core::use_batch_gas(&mut ctx, 1).expect("using batch gas under limit should succeed");
 
     let mut tx = mock::transaction();
     tx.auth_info.fee.gas = MAX_GAS;
 
     ctx.with_tx(tx.clone(), |mut tx_ctx, _call| {
-        Core::use_gas(&mut tx_ctx, MAX_GAS).expect("using gas under limit should succeed");
+        Core::use_tx_gas(&mut tx_ctx, MAX_GAS).expect("using gas under limit should succeed");
     });
 
     ctx.with_tx(tx.clone(), |mut tx_ctx, _call| {
-        Core::use_gas(&mut tx_ctx, MAX_GAS)
+        Core::use_tx_gas(&mut tx_ctx, MAX_GAS)
             .expect("gas across separate transactions shouldn't accumulate");
     });
 
     ctx.with_tx(tx.clone(), |mut tx_ctx, _call| {
-        Core::use_gas(&mut tx_ctx, MAX_GAS).unwrap();
-        Core::use_gas(&mut tx_ctx, 1).expect_err("gas in same transaction should accumulate");
+        Core::use_tx_gas(&mut tx_ctx, MAX_GAS).unwrap();
+        Core::use_tx_gas(&mut tx_ctx, 1).expect_err("gas in same transaction should accumulate");
     });
 
     ctx.with_tx(tx.clone(), |mut tx_ctx, _call| {
-        Core::use_gas(&mut tx_ctx, 1).unwrap();
-        Core::use_gas(&mut tx_ctx, u64::MAX).expect_err("overflow should cause error");
+        Core::use_tx_gas(&mut tx_ctx, 1).unwrap();
+        Core::use_tx_gas(&mut tx_ctx, u64::MAX).expect_err("overflow should cause error");
     });
 
     let mut big_tx = tx.clone();
     big_tx.auth_info.fee.gas = u64::MAX;
     ctx.with_tx(big_tx, |mut tx_ctx, _call| {
-        Core::use_gas(&mut tx_ctx, u64::MAX).expect_err("batch overflow should cause error");
+        Core::use_tx_gas(&mut tx_ctx, u64::MAX).expect_err("batch overflow should cause error");
     });
 
     ctx.with_tx(tx.clone(), |mut tx_ctx, _call| {
-        Core::use_gas(&mut tx_ctx, 1).expect_err("batch gas should accumulate");
+        Core::use_tx_gas(&mut tx_ctx, 1).expect_err("batch gas should accumulate");
     });
 
-    Core::use_gas(&mut ctx, 1).expect_err("batch gas should accumulate outside tx");
+    Core::use_batch_gas(&mut ctx, 1).expect_err("batch gas should accumulate outside tx");
+}
+
+// Module that implements the gas waster method.
+struct GasWasterModule;
+
+impl GasWasterModule {
+    const MAX_GAS: u64 = 100;
+    const METHOD_WASTE_GAS: &'static str = "test.WasteGas";
+}
+
+impl module::Module for GasWasterModule {
+    const NAME: &'static str = "gaswaster";
+    type Error = std::convert::Infallible;
+    type Event = ();
+    type Parameters = ();
+}
+
+impl module::MethodHandler for GasWasterModule {
+    fn dispatch_call<C: TxContext>(
+        ctx: &mut C,
+        method: &str,
+        body: cbor::Value,
+    ) -> module::DispatchResult<cbor::Value, transaction::CallResult> {
+        match method {
+            Self::METHOD_WASTE_GAS => {
+                Core::use_tx_gas(ctx, Self::MAX_GAS).expect("use_gas should succeed");
+                module::DispatchResult::Handled(transaction::CallResult::Ok(cbor::Value::Null))
+            }
+            _ => module::DispatchResult::Unhandled(body),
+        }
+    }
+}
+
+impl module::BlockHandler for GasWasterModule {}
+impl module::AuthHandler for GasWasterModule {}
+impl module::MigrationHandler for GasWasterModule {
+    type Genesis = ();
+}
+
+// Runtime that knows how to waste gas.
+struct GasWasterRuntime;
+
+impl Runtime for GasWasterRuntime {
+    const VERSION: Version = Version::new(0, 0, 0);
+
+    type Modules = (Core, GasWasterModule);
+
+    fn genesis_state() -> (super::Genesis, ()) {
+        (
+            super::Genesis {
+                parameters: super::Parameters {
+                    max_batch_gas: u64::MAX,
+                    max_tx_signers: 8,
+                    max_multisig_signers: 8,
+                },
+            },
+            (),
+        )
+    }
 }
 
 #[test]
 fn test_query_estimate_gas() {
-    const MAX_GAS: u64 = 100;
-    const METHOD_WASTE_GAS: &str = "test.WasteGas";
     let mut mock = mock::Mock::default();
-    mock.methods.register_callable(module::CallableMethodInfo {
-        name: METHOD_WASTE_GAS,
-        handler: |_mi, ctx, _args| {
-            Core::use_gas(ctx, MAX_GAS).expect("use_gas should succeed");
-            transaction::CallResult::Ok(cbor::Value::Null)
-        },
-    });
-    let mut ctx = mock.create_ctx();
-    ctx.mode = Mode::CheckTx;
-    Core::set_params(
-        ctx.runtime_state(),
-        &super::Parameters {
-            max_batch_gas: u64::MAX,
-            max_tx_signers: 8,
-            max_multisig_signers: 8,
-        },
-    );
+    let mut ctx = mock.create_ctx_for_runtime::<GasWasterRuntime>(Mode::CheckTx);
+
+    GasWasterRuntime::migrate(&mut ctx);
 
     let tx = transaction::Transaction {
         version: 1,
         call: transaction::Call {
-            method: METHOD_WASTE_GAS.to_owned(),
+            method: GasWasterModule::METHOD_WASTE_GAS.to_owned(),
             body: cbor::Value::Null,
         },
         auth_info: transaction::AuthInfo {
@@ -101,7 +144,11 @@ fn test_query_estimate_gas() {
     };
 
     let est = Core::query_estimate_gas(&mut ctx, tx).expect("query_estimate_gas should succeed");
-    assert_eq!(est, MAX_GAS, "estimated gas should be correct");
+    assert_eq!(
+        est,
+        GasWasterModule::MAX_GAS,
+        "estimated gas should be correct"
+    );
 }
 
 #[test]
diff --git a/runtime-sdk/src/modules/rewards/test.rs b/runtime-sdk/src/modules/rewards/test.rs
index f6feb217bc..4ad8dc8f52 100644
--- a/runtime-sdk/src/modules/rewards/test.rs
+++ b/runtime-sdk/src/modules/rewards/test.rs
@@ -2,7 +2,7 @@
 use std::collections::BTreeMap;
 
 use crate::{
-    context::{Context, DispatchContext},
+    context::Context,
     module::{BlockHandler, MigrationHandler},
     modules::{
         accounts::{self, Module as Accounts, API as _},
@@ -16,7 +16,7 @@ use super::{types, Genesis, Parameters, ADDRESS_REWARD_POOL};
 
 type Rewards = super::Module<Accounts>;
 
-fn init_accounts(ctx: &mut DispatchContext<'_>) {
+fn init_accounts<C: Context>(ctx: &mut C) {
     Accounts::init_or_migrate(
         ctx,
         &mut core::types::Metadata::default(),
diff --git a/runtime-sdk/src/testing/mock.rs b/runtime-sdk/src/testing/mock.rs
index 7ae28e1d68..47fb6a5ae4 100644
--- a/runtime-sdk/src/testing/mock.rs
+++ b/runtime-sdk/src/testing/mock.rs
@@ -1,22 +1,34 @@
 //! Mock dispatch context for use in tests.
-use std::collections::BTreeMap;
-
 use io_context::Context as IoContext;
 
 use oasis_core_runtime::{
-    common::{cbor, logger::get_logger},
+    common::{cbor, version::Version},
     consensus::{beacon, roothash, state::ConsensusState},
     storage::mkvs,
-    transaction::tags::Tags,
 };
 
 use crate::{
-    context::{DispatchContext, Mode},
-    module::MethodRegistry,
+    context::{Mode, RuntimeBatchContext},
+    module::MigrationHandler,
+    modules,
+    runtime::Runtime,
     storage,
     types::transaction,
 };
 
+/// A mock runtime that only has the core module.
+pub struct EmptyRuntime;
+
+impl Runtime for EmptyRuntime {
+    const VERSION: Version = Version::new(0, 0, 0);
+
+    type Modules = modules::core::Module;
+
+    fn genesis_state() -> <Self::Modules as MigrationHandler>::Genesis {
+        Default::default()
+    }
+}
+
 /// Mock dispatch context factory.
 pub struct Mock {
     pub runtime_header: roothash::Header,
@@ -25,32 +37,32 @@ pub struct Mock {
     pub consensus_state: ConsensusState,
     pub epoch: beacon::EpochTime,
 
-    pub methods: MethodRegistry,
-
     pub max_messages: u32,
 }
 
 impl Mock {
     /// Create a new mock dispatch context.
-    pub fn create_ctx(&mut self) -> DispatchContext<'_> {
-        DispatchContext {
-            mode: Mode::ExecuteTx,
-            runtime_header: &self.runtime_header,
-            runtime_round_results: &self.runtime_round_results,
-            runtime_storage: storage::MKVSStore::new(
-                IoContext::background().freeze(),
-                self.mkvs.as_mut(),
-            ),
-            consensus_state: &self.consensus_state,
-            epoch: self.epoch,
-            io_ctx: IoContext::background().freeze(),
-            methods: &self.methods,
-            logger: get_logger("mock"),
-            block_tags: Tags::new(),
-            messages: Vec::new(),
-            max_messages: self.max_messages,
-            values: BTreeMap::new(),
-        }
+    pub fn create_ctx(
+        &mut self,
+    ) -> RuntimeBatchContext<'_, EmptyRuntime, storage::MKVSStore<&mut dyn mkvs::MKVS>> {
+        self.create_ctx_for_runtime(Mode::ExecuteTx)
+    }
+
+    /// Create a new mock dispatch context.
+    pub fn create_ctx_for_runtime<R: Runtime>(
+        &mut self,
+        mode: Mode,
+    ) -> RuntimeBatchContext<'_, R, storage::MKVSStore<&mut dyn mkvs::MKVS>> {
+        RuntimeBatchContext::new(
+            mode,
+            &self.runtime_header,
+            &self.runtime_round_results,
+            storage::MKVSStore::new(IoContext::background().freeze(), self.mkvs.as_mut()),
+            &self.consensus_state,
+            self.epoch,
+            IoContext::background().freeze(),
+            self.max_messages,
+        )
     }
 }
 
@@ -71,7 +83,6 @@ impl Default for Mock {
             mkvs: Box::new(mkvs),
             consensus_state: ConsensusState::new(consensus_tree),
             epoch: 1,
-            methods: MethodRegistry::new(),
             max_messages: 32,
         }
     }
diff --git a/tests/runtimes/simple-keyvalue/src/keyvalue.rs b/tests/runtimes/simple-keyvalue/src/keyvalue.rs
index 803f580a39..8353b9992d 100644
--- a/tests/runtimes/simple-keyvalue/src/keyvalue.rs
+++ b/tests/runtimes/simple-keyvalue/src/keyvalue.rs
@@ -3,10 +3,10 @@ use thiserror::Error;
 
 use oasis_runtime_sdk::{
     self as sdk,
-    context::{Context, DispatchContext, TxContext},
+    context::{Context, TxContext},
     core::common::cbor,
     error::{Error as _, RuntimeError},
-    module::{CallableMethodInfo, Module as _, QueryMethodInfo},
+    module::Module as _,
     modules::{
         core,
         core::{Module as Core, API as _},
@@ -89,72 +89,57 @@ impl sdk::module::Module for Module {
 
 impl sdk::module::AuthHandler for Module {}
 impl sdk::module::BlockHandler for Module {}
-impl sdk::module::MessageHookRegistrationHandler for Module {}
-
-impl sdk::module::MethodRegistrationHandler for Module {
-    /// Register all supported methods.
-    fn register_methods(methods: &mut sdk::module::MethodRegistry) {
-        methods.register_callable(sdk::module::CallableMethodInfo {
-            name: "keyvalue.Insert",
-            handler: Self::_callable_insert_handler,
-        });
-        methods.register_callable(sdk::module::CallableMethodInfo {
-            name: "keyvalue.Remove",
-            handler: Self::_callable_remove_handler,
-        });
-        methods.register_query(sdk::module::QueryMethodInfo {
-            name: "keyvalue.Get",
-            handler: Self::_query_get_handler,
-        });
-    }
-}
 
-// Boilerplate.
-impl Module {
-    fn _callable_insert_handler(
-        _mi: &CallableMethodInfo,
-        ctx: &mut TxContext,
+impl sdk::module::MethodHandler for Module {
+    fn dispatch_call<C: TxContext>(
+        ctx: &mut C,
+        method: &str,
         body: cbor::Value,
-    ) -> CallResult {
-        let result = || -> Result<cbor::Value, Error> {
-            let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
-            Ok(cbor::to_value(&Self::insert(ctx, args)?))
-        }();
-        match result {
-            Ok(value) => CallResult::Ok(value),
-            Err(err) => err.to_call_result(),
+    ) -> sdk::module::DispatchResult<cbor::Value, CallResult> {
+        match method {
+            "keyvalue.Insert" => {
+                let result = || -> Result<cbor::Value, Error> {
+                    let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
+                    Ok(cbor::to_value(&Self::tx_insert(ctx, args)?))
+                }();
+                match result {
+                    Ok(value) => sdk::module::DispatchResult::Handled(CallResult::Ok(value)),
+                    Err(err) => sdk::module::DispatchResult::Handled(err.to_call_result()),
+                }
+            }
+            "keyvalue.Remove" => {
+                let result = || -> Result<cbor::Value, Error> {
+                    let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
+                    Ok(cbor::to_value(&Self::tx_remove(ctx, args)?))
+                }();
+                match result {
+                    Ok(value) => sdk::module::DispatchResult::Handled(CallResult::Ok(value)),
+                    Err(err) => sdk::module::DispatchResult::Handled(err.to_call_result()),
+                }
+            }
+            _ => sdk::module::DispatchResult::Unhandled(body),
         }
     }
 
-    fn _callable_remove_handler(
-        _mi: &CallableMethodInfo,
-        ctx: &mut TxContext,
-        body: cbor::Value,
-    ) -> CallResult {
-        let result = || -> Result<cbor::Value, Error> {
-            let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
-            Ok(cbor::to_value(&Self::remove(ctx, args)?))
-        }();
-        match result {
-            Ok(value) => CallResult::Ok(value),
-            Err(err) => err.to_call_result(),
+    fn dispatch_query<C: Context>(
+        ctx: &mut C,
+        method: &str,
+        args: cbor::Value,
+    ) -> sdk::module::DispatchResult<cbor::Value, Result<cbor::Value, RuntimeError>> {
+        match method {
+            "keyvalue.Get" => sdk::module::DispatchResult::Handled((|| {
+                let args = cbor::from_value(args).map_err(|_| Error::InvalidArgument)?;
+                Ok(cbor::to_value(&Self::query_get(ctx, args)?))
+            })()),
+            _ => sdk::module::DispatchResult::Unhandled(args),
         }
     }
-
-    fn _query_get_handler(
-        _mi: &QueryMethodInfo,
-        ctx: &mut DispatchContext,
-        body: cbor::Value,
-    ) -> Result<cbor::Value, RuntimeError> {
-        let args = cbor::from_value(body).map_err(|_| Error::InvalidArgument)?;
-        Ok(cbor::to_value(&Self::get(ctx, args)?))
-    }
 }
 
 // Actual implementation of this runtime's externally-callable methods.
 impl Module {
-    // Insert given keyvalue into storage.
-    fn insert(ctx: &mut TxContext, body: types::KeyValue) -> Result<(), Error> {
+    /// Insert given keyvalue into storage.
+    fn tx_insert<C: TxContext>(ctx: &mut C, body: types::KeyValue) -> Result<(), Error> {
         if ctx.is_check_only() {
             return Ok(());
         }
@@ -168,7 +153,7 @@ impl Module {
             Some(_) => params.gas_costs.insert_existing,
         };
         // We must drop ts and store so that use_gas can borrow ctx.
-        Core::use_gas(ctx, cost)?;
+        Core::use_tx_gas(ctx, cost)?;
 
         // Recreate store and ts after we get ctx back
         let mut store = sdk::storage::PrefixStore::new(ctx.runtime_state(), &MODULE_NAME);
@@ -179,8 +164,8 @@ impl Module {
         Ok(())
     }
 
-    // Remove keyvalue from storage using given key.
-    fn remove(ctx: &mut TxContext, body: types::Key) -> Result<(), Error> {
+    /// Remove keyvalue from storage using given key.
+    fn tx_remove<C: TxContext>(ctx: &mut C, body: types::Key) -> Result<(), Error> {
         if ctx.is_check_only() {
             return Ok(());
         }
@@ -194,7 +179,7 @@ impl Module {
             Some(_) => params.gas_costs.remove_existing,
         };
         // We must drop ts and store so that use_gas can borrow ctx.
-        Core::use_gas(ctx, cost)?;
+        Core::use_tx_gas(ctx, cost)?;
 
         // Recreate store and ts after we get ctx back
         let mut store = sdk::storage::PrefixStore::new(ctx.runtime_state(), &MODULE_NAME);
@@ -205,8 +190,8 @@ impl Module {
         Ok(())
     }
 
-    // Fetch keyvalue from storage using given key.
-    fn get(ctx: &mut DispatchContext, body: types::Key) -> Result<types::KeyValue, Error> {
+    /// Fetch keyvalue from storage using given key.
+    fn query_get<C: Context>(ctx: &mut C, body: types::Key) -> Result<types::KeyValue, Error> {
         let mut store = sdk::storage::PrefixStore::new(ctx.runtime_state(), &MODULE_NAME);
         let ts = sdk::storage::TypedStore::new(&mut store);
         let v: Vec<u8> = ts.get(body.key.clone()).ok_or(Error::InvalidArgument)?;
@@ -220,8 +205,8 @@ impl Module {
 impl sdk::module::MigrationHandler for Module {
     type Genesis = Genesis;
 
-    fn init_or_migrate(
-        ctx: &mut DispatchContext,
+    fn init_or_migrate<C: Context>(
+        ctx: &mut C,
         meta: &mut sdk::modules::core::types::Metadata,
         genesis: &Self::Genesis,
     ) -> bool {