feat(driver): Multi-block derivation (#888)

Co-authored-by: refcell <[email protected]>

bin/client/src/interop.rs

@@ -128,6 +128,8 @@ where
     let cursor =
         new_pipeline_cursor(&boot.rollup_config, safe_head, &mut l1_provider, &mut l2_provider)
             .await?;
+    l2_provider.set_cursor(cursor.clone());
+
     let cfg = Arc::new(boot.rollup_config.clone());
     let pipeline = OraclePipeline::new(
         cfg.clone(),

bin/client/src/single.rs

@@ -6,9 +6,7 @@ use alloy_primitives::B256;
 use core::fmt::Debug;
 use kona_driver::{Driver, DriverError};
 use kona_executor::{ExecutorError, KonaHandleRegister, TrieDBProvider};
-use kona_preimage::{
-    CommsClient, HintWriterClient, PreimageKey, PreimageKeyType, PreimageOracleClient,
-};
+use kona_preimage::{CommsClient, HintWriterClient, PreimageKeyType, PreimageOracleClient};
 use kona_proof::{
     errors::OracleProviderError,
     executor::KonaExecutor,
@@ -109,6 +107,8 @@ where
     let cursor =
         new_pipeline_cursor(&boot.rollup_config, safe_head, &mut l1_provider, &mut l2_provider)
             .await?;
+    l2_provider.set_cursor(cursor.clone());
+
     let cfg = Arc::new(boot.rollup_config.clone());
     let pipeline = OraclePipeline::new(
         cfg.clone(),
@@ -152,25 +152,22 @@ where
 
 /// Fetches the safe head hash of the L2 chain based on the agreed upon L2 output root in the
 /// [BootInfo].
-async fn fetch_safe_head_hash<O>(
+pub async fn fetch_safe_head_hash<O>(
     caching_oracle: &O,
     boot_info: &BootInfo,
 ) -> Result<B256, OracleProviderError>
 where
     O: CommsClient,
 {
-    caching_oracle
-        .write(&HintType::StartingL2Output.encode_with(&[boot_info.agreed_l2_output_root.as_ref()]))
-        .await
-        .map_err(OracleProviderError::Preimage)?;
     let mut output_preimage = [0u8; 128];
-    caching_oracle
-        .get_exact(
-            PreimageKey::new(*boot_info.agreed_l2_output_root, PreimageKeyType::Keccak256),
+    HintType::StartingL2Output
+        .get_exact_preimage(
+            caching_oracle,
+            boot_info.agreed_l2_output_root,
+            PreimageKeyType::Keccak256,
             &mut output_preimage,
         )
-        .await
-        .map_err(OracleProviderError::Preimage)?;
+        .await?;
 
     output_preimage[96..128].try_into().map_err(OracleProviderError::SliceConversion)
 }

crates/driver/Cargo.toml

@@ -29,6 +29,7 @@ op-alloy-consensus.workspace = true
 op-alloy-rpc-types-engine.workspace = true
 
 # Misc
-tracing.workspace = true
-thiserror .workspace = true
 async-trait.workspace = true
+spin.workspace = true
+thiserror .workspace = true
+tracing.workspace = true

crates/driver/src/core.rs

@@ -1,6 +1,7 @@
 //! The driver of the kona derivation pipeline.
 
-use alloc::vec::Vec;
+use crate::{DriverError, DriverPipeline, DriverResult, Executor, PipelineCursor, TipCursor};
+use alloc::{sync::Arc, vec::Vec};
 use alloy_consensus::{BlockBody, Sealable};
 use alloy_primitives::B256;
 use alloy_rlp::Decodable;
@@ -14,8 +15,7 @@ use maili_genesis::RollupConfig;
 use maili_protocol::L2BlockInfo;
 use op_alloy_consensus::{OpBlock, OpTxEnvelope, OpTxType};
 use op_alloy_rpc_types_engine::OpAttributesWithParent;
-
-use crate::{DriverError, DriverPipeline, DriverResult, Executor, PipelineCursor, TipCursor};
+use spin::RwLock;
 
 /// The Rollup Driver entrypoint.
 #[derive(Debug)]
@@ -32,7 +32,7 @@ where
     /// A pipeline abstraction.
     pub pipeline: DP,
     /// Cursor to keep track of the L2 tip
-    pub cursor: PipelineCursor,
+    pub cursor: Arc<RwLock<PipelineCursor>>,
     /// The Executor.
     pub executor: E,
 }
@@ -44,7 +44,7 @@ where
     P: Pipeline + SignalReceiver + Send + Sync + Debug,
 {
     /// Creates a new [Driver].
-    pub const fn new(cursor: PipelineCursor, executor: E, pipeline: DP) -> Self {
+    pub const fn new(cursor: Arc<RwLock<PipelineCursor>>, executor: E, pipeline: DP) -> Self {
         Self {
             _marker: core::marker::PhantomData,
             _marker2: core::marker::PhantomData,
@@ -76,20 +76,22 @@ where
     ) -> DriverResult<(u64, B256, B256), E::Error> {
         loop {
             // Check if we have reached the target block number.
+            let pipeline_cursor = self.cursor.read();
+            let tip_cursor = pipeline_cursor.tip();
             if let Some(tb) = target {
-                if self.cursor.l2_safe_head().block_info.number >= tb {
+                if tip_cursor.l2_safe_head.block_info.number >= tb {
                     info!(target: "client", "Derivation complete, reached L2 safe head.");
                     return Ok((
-                        self.cursor.l2_safe_head().block_info.number,
-                        self.cursor.l2_safe_head().block_info.hash,
-                        *self.cursor.l2_safe_head_output_root(),
+                        tip_cursor.l2_safe_head.block_info.number,
+                        tip_cursor.l2_safe_head.block_info.hash,
+                        tip_cursor.l2_safe_head_output_root,
                     ));
                 }
             }
 
             let OpAttributesWithParent { mut attributes, .. } = match self
                 .pipeline
-                .produce_payload(*self.cursor.l2_safe_head())
+                .produce_payload(tip_cursor.l2_safe_head)
                 .await
             {
                 Ok(attrs) => attrs,
@@ -99,7 +101,7 @@ where
                     // Adjust the target block number to the current safe head, as no more blocks
                     // can be produced.
                     if target.is_some() {
-                        target = Some(self.cursor.l2_safe_head().block_info.number);
+                        target = Some(tip_cursor.l2_safe_head.block_info.number);
                     };
                     continue;
                 }
@@ -109,7 +111,7 @@ where
                 }
             };
 
-            self.executor.update_safe_head(self.cursor.l2_safe_head_header().clone());
+            self.executor.update_safe_head(tip_cursor.l2_safe_head_header.clone());
             let header = match self.executor.execute_payload(attributes.clone()).await {
                 Ok(header) => header,
                 Err(e) => {
@@ -133,7 +135,7 @@ where
                         });
 
                         // Retry the execution.
-                        self.executor.update_safe_head(self.cursor.l2_safe_head_header().clone());
+                        self.executor.update_safe_head(tip_cursor.l2_safe_head_header.clone());
                         match self.executor.execute_payload(attributes.clone()).await {
                             Ok(header) => header,
                             Err(e) => {
@@ -166,18 +168,21 @@ where
                 },
             };
 
-            // Get the pipeline origin and update the cursor.
+            // Get the pipeline origin and update the tip cursor.
             let origin = self.pipeline.origin().ok_or(PipelineError::MissingOrigin.crit())?;
             let l2_info = L2BlockInfo::from_block_and_genesis(
                 &block,
                 &self.pipeline.rollup_config().genesis,
             )?;
-            let cursor = TipCursor::new(
+            let tip_cursor = TipCursor::new(
                 l2_info,
                 header.clone().seal_slow(),
                 self.executor.compute_output_root().map_err(DriverError::Executor)?,
             );
-            self.cursor.advance(origin, cursor);
+
+            // Advance the derivation pipeline cursor
+            drop(pipeline_cursor);
+            self.cursor.write().advance(origin, tip_cursor);
         }
     }
 }

crates/proof-sdk/proof/src/hint.rs

@@ -1,12 +1,13 @@
 //! This module contains the [HintType] enum.
 
-use crate::errors::HintParsingError;
+use crate::errors::{HintParsingError, OracleProviderError};
 use alloc::{
     string::{String, ToString},
     vec::Vec,
 };
-use alloy_primitives::{hex, Bytes};
+use alloy_primitives::{hex, Bytes, B256};
 use core::fmt::Display;
+use kona_preimage::{CommsClient, PreimageKey, PreimageKeyType};
 
 /// A [Hint] is parsed in the format `<hint_type> <hint_data>`, where `<hint_type>` is a string that
 /// represents the type of hint, and `<hint_data>` is the data associated with the hint (bytes
@@ -80,6 +81,39 @@ impl HintType {
         let concatenated = hex::encode(data.iter().copied().flatten().copied().collect::<Vec<_>>());
         alloc::format!("{} {}", self, concatenated)
     }
+    /// Retrieves a preimage through an oracle
+    pub async fn get_preimage<T: CommsClient>(
+        &self,
+        oracle: &T,
+        image: B256,
+        preimage_key_type: PreimageKeyType,
+    ) -> Result<Vec<u8>, OracleProviderError> {
+        oracle
+            .write(&self.encode_with(&[image.as_ref()]))
+            .await
+            .map_err(OracleProviderError::Preimage)?;
+        oracle
+            .get(PreimageKey::new(*image, preimage_key_type))
+            .await
+            .map_err(OracleProviderError::Preimage)
+    }
+    /// Retrieves a preimage through an oracle
+    pub async fn get_exact_preimage<T: CommsClient>(
+        &self,
+        oracle: &T,
+        image: B256,
+        preimage_key_type: PreimageKeyType,
+        buf: &mut [u8],
+    ) -> Result<(), OracleProviderError> {
+        oracle
+            .write(&self.encode_with(&[image.as_ref()]))
+            .await
+            .map_err(OracleProviderError::Preimage)?;
+        oracle
+            .get_exact(PreimageKey::new(*image, preimage_key_type), buf)
+            .await
+            .map_err(OracleProviderError::Preimage)
+    }
 }
 
 impl TryFrom<&str> for HintType {

crates/proof-sdk/proof/src/l1/chain_provider.rs

@@ -33,18 +33,10 @@ impl<T: CommsClient + Sync + Send> ChainProvider for OracleL1ChainProvider<T> {
     type Error = OracleProviderError;
 
     async fn header_by_hash(&mut self, hash: B256) -> Result<Header, Self::Error> {
-        // Send a hint for the block header.
-        self.oracle
-            .write(&HintType::L1BlockHeader.encode_with(&[hash.as_ref()]))
-            .await
-            .map_err(OracleProviderError::Preimage)?;
-
         // Fetch the header RLP from the oracle.
-        let header_rlp = self
-            .oracle
-            .get(PreimageKey::new(*hash, PreimageKeyType::Keccak256))
-            .await
-            .map_err(OracleProviderError::Preimage)?;
+        let header_rlp = HintType::L1BlockHeader
+            .get_preimage(self.oracle.as_ref(), hash, PreimageKeyType::Keccak256)
+            .await?;
 
         // Decode the header RLP into a Header.
         Header::decode(&mut header_rlp.as_slice()).map_err(OracleProviderError::Rlp)

crates/proof-sdk/proof/src/l1/pipeline.rs

@@ -21,6 +21,7 @@ use kona_preimage::CommsClient;
 use maili_genesis::{RollupConfig, SystemConfig};
 use maili_protocol::{BlockInfo, L2BlockInfo};
 use op_alloy_rpc_types_engine::OpAttributesWithParent;
+use spin::RwLock;
 
 /// An oracle-backed derivation pipeline.
 pub type OracleDerivationPipeline<O, B> = DerivationPipeline<
@@ -73,7 +74,7 @@ where
     /// Constructs a new oracle-backed derivation pipeline.
     pub fn new(
         cfg: Arc<RollupConfig>,
-        sync_start: PipelineCursor,
+        sync_start: Arc<RwLock<PipelineCursor>>,
         caching_oracle: Arc<O>,
         blob_provider: B,
         chain_provider: OracleL1ChainProvider<O>,
@@ -92,7 +93,7 @@ where
             .l2_chain_provider(l2_chain_provider)
             .chain_provider(chain_provider)
             .builder(attributes)
-            .origin(sync_start.origin())
+            .origin(sync_start.read().origin())
             .build();
         Self { pipeline, caching_oracle }
     }

crates/proof-sdk/proof/src/l2/chain_provider.rs

@@ -8,12 +8,14 @@ use alloy_primitives::{Address, Bytes, B256};
 use alloy_rlp::Decodable;
 use async_trait::async_trait;
 use kona_derive::traits::L2ChainProvider;
+use kona_driver::PipelineCursor;
 use kona_executor::TrieDBProvider;
 use kona_mpt::{OrderedListWalker, TrieHinter, TrieNode, TrieProvider};
 use kona_preimage::{CommsClient, PreimageKey, PreimageKeyType};
 use maili_genesis::{RollupConfig, SystemConfig};
 use maili_protocol::{to_system_config, BatchValidationProvider, L2BlockInfo};
 use op_alloy_consensus::{OpBlock, OpTxEnvelope};
+use spin::RwLock;
 
 /// The oracle-backed L2 chain provider for the client program.
 #[derive(Debug, Clone)]
@@ -24,20 +26,36 @@ pub struct OracleL2ChainProvider<T: CommsClient> {
     rollup_config: RollupConfig,
     /// The preimage oracle client.
     oracle: Arc<T>,
+    /// The derivation pipeline cursor
+    cursor: Option<Arc<RwLock<PipelineCursor>>>,
 }
 
 impl<T: CommsClient> OracleL2ChainProvider<T> {
     /// Creates a new [OracleL2ChainProvider] with the given boot information and oracle client.
     pub const fn new(l2_head: B256, rollup_config: RollupConfig, oracle: Arc<T>) -> Self {
-        Self { l2_head, rollup_config, oracle }
+        Self { l2_head, rollup_config, oracle, cursor: None }
+    }
+
+    /// Updates the derivation pipeline cursor
+    pub fn set_cursor(&mut self, cursor: Arc<RwLock<PipelineCursor>>) {
+        self.cursor = Some(cursor);
+    }
+
+    /// Fetches the latest known safe head block hash according to the derivation pipeline cursor
+    /// or uses the initial l2_head value if no cursor is set.
+    pub async fn l2_safe_head(&self) -> Result<B256, OracleProviderError> {
+        self.cursor
+            .as_ref()
+            .map_or(Ok(self.l2_head), |cursor| Ok(cursor.read().l2_safe_head().block_info.hash))
     }
 }
 
 impl<T: CommsClient> OracleL2ChainProvider<T> {
     /// Returns a [Header] corresponding to the given L2 block number, by walking back from the
     /// L2 safe head.
     async fn header_by_number(&mut self, block_number: u64) -> Result<Header, OracleProviderError> {
-        let mut header = self.header_by_hash(self.l2_head)?;
+        // Fetch the starting block header.
+        let mut header = self.header_by_hash(self.l2_safe_head().await?)?;
 
         // Check if the block number is in range. If not, we can fail early.
         if block_number > header.number {
@@ -151,32 +169,20 @@ impl<T: CommsClient> TrieDBProvider for OracleL2ChainProvider<T> {
     fn bytecode_by_hash(&self, hash: B256) -> Result<Bytes, OracleProviderError> {
         // Fetch the bytecode preimage from the caching oracle.
         crate::block_on(async move {
-            self.oracle
-                .write(&HintType::L2Code.encode_with(&[hash.as_ref()]))
-                .await
-                .map_err(OracleProviderError::Preimage)?;
-
-            self.oracle
-                .get(PreimageKey::new(*hash, PreimageKeyType::Keccak256))
+            HintType::L2Code
+                .get_preimage(self.oracle.as_ref(), hash, PreimageKeyType::Keccak256)
                 .await
                 .map(Into::into)
-                .map_err(OracleProviderError::Preimage)
         })
     }
 
     fn header_by_hash(&self, hash: B256) -> Result<Header, OracleProviderError> {
         // Fetch the header from the caching oracle.
         crate::block_on(async move {
-            self.oracle
-                .write(&HintType::L2BlockHeader.encode_with(&[hash.as_ref()]))
-                .await
-                .map_err(OracleProviderError::Preimage)?;
+            let header_bytes = HintType::L2BlockHeader
+                .get_preimage(self.oracle.as_ref(), hash, PreimageKeyType::Keccak256)
+                .await?;
 
-            let header_bytes = self
-                .oracle
-                .get(PreimageKey::new(*hash, PreimageKeyType::Keccak256))
-                .await
-                .map_err(OracleProviderError::Preimage)?;
             Header::decode(&mut header_bytes.as_slice()).map_err(OracleProviderError::Rlp)
         })
     }