From 3fe22343f1b96ef97c4199741e0b44b6567799b4 Mon Sep 17 00:00:00 2001
From: Federico Gimenez <fgimenez@users.noreply.github.com>
Date: Mon, 4 Nov 2024 13:48:32 +0100
Subject: [PATCH] feat(engine): add StateRootTask skeleton (#12305)

---
 Cargo.lock                          |  2 +
 crates/engine/tree/Cargo.toml       |  3 ++
 crates/engine/tree/src/tree/mod.rs  |  2 +
 crates/engine/tree/src/tree/root.rs | 60 +++++++++++++++++++++++++++++
 4 files changed, 67 insertions(+)
 create mode 100644 crates/engine/tree/src/tree/root.rs

diff --git a/Cargo.lock b/Cargo.lock
index 9b5a5350c16c..a885081cd36b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7218,8 +7218,10 @@ dependencies = [
  "reth-tracing",
  "reth-trie",
  "reth-trie-parallel",
+ "revm-primitives",
  "thiserror",
  "tokio",
+ "tokio-stream",
  "tracing",
 ]
 
diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml
index dee0bcaf7ce1..293883c036e7 100644
--- a/crates/engine/tree/Cargo.toml
+++ b/crates/engine/tree/Cargo.toml
@@ -39,9 +39,12 @@ alloy-primitives.workspace = true
 alloy-eips.workspace = true
 alloy-rpc-types-engine.workspace = true
 
+revm-primitives.workspace = true
+
 # common
 futures.workspace = true
 tokio = { workspace = true, features = ["macros", "sync"] }
+tokio-stream.workspace = true
 thiserror.workspace = true
 
 # metrics
diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index bc070d87345e..c3e922d11c9f 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -75,6 +75,8 @@ pub use invalid_block_hook::{InvalidBlockHooks, NoopInvalidBlockHook};
 pub use persistence_state::PersistenceState;
 pub use reth_engine_primitives::InvalidBlockHook;
 
+mod root;
+
 /// Keeps track of the state of the tree.
 ///
 /// ## Invariants
diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
new file mode 100644
index 000000000000..48b2eccdf147
--- /dev/null
+++ b/crates/engine/tree/src/tree/root.rs
@@ -0,0 +1,60 @@
+//! State root task related functionality.
+
+use reth_provider::providers::ConsistentDbView;
+use reth_trie::{updates::TrieUpdates, TrieInput};
+use reth_trie_parallel::parallel_root::ParallelStateRootError;
+use revm_primitives::{EvmState, B256};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+};
+use tokio_stream::wrappers::UnboundedReceiverStream;
+
+/// Standalone task that receives a transaction state stream and updates relevant
+/// data structures to calculate state root.
+///
+/// It is responsile of  initializing a blinded sparse trie and subscribe to
+/// transaction state stream. As it receives transaction execution results, it
+/// fetches the proofs for relevant accounts from the database and reveal them
+/// to the tree.
+/// Then it updates relevant leaves according to the result of the transaction.
+#[allow(dead_code)]
+pub(crate) struct StateRootTask<Factory> {
+    /// View over the state in the database.
+    consistent_view: ConsistentDbView<Factory>,
+    /// Incoming state updates.
+    state_stream: UnboundedReceiverStream<EvmState>,
+    /// Latest trie input.
+    input: Arc<TrieInput>,
+}
+
+#[allow(dead_code)]
+impl<Factory> StateRootTask<Factory> {
+    /// Creates a new `StateRootTask`.
+    pub(crate) const fn new(
+        consistent_view: ConsistentDbView<Factory>,
+        input: Arc<TrieInput>,
+        state_stream: UnboundedReceiverStream<EvmState>,
+    ) -> Self {
+        Self { consistent_view, state_stream, input }
+    }
+
+    /// Handles state updates.
+    pub(crate) fn on_state_update(&self, _update: EvmState) {
+        // TODO: calculate hashed state update and dispatch proof gathering for it.
+    }
+}
+
+impl<Factory> Future for StateRootTask<Factory> {
+    type Output = Result<(B256, TrieUpdates), ParallelStateRootError>;
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // TODO:
+        //    * poll incoming state updates stream
+        //    * keep track of proof calculation
+        //    * keep track of intermediate root computation
+        Poll::Pending
+    }
+}