feat: BWIP

core/lib/dal/migrations/20240619143458_add_vm_run_data_blob_url_column.down.sql

@@ -0,0 +1,2 @@
+ALTER TABLE proof_generation_details DROP COLUMN IF EXISTS vm_run_data_blob_url;
+DROP TABLE IF EXISTS vm_runner_protective_reads;

core/lib/dal/migrations/20240619143458_add_vm_run_data_blob_url_column.up.sql

@@ -0,0 +1,10 @@
+ALTER TABLE proof_generation_details
+ ADD COLUMN IF NOT EXISTS vm_run_data_blob_url DEFAULT NULL;
+
+CREATE TABLE IF NOT EXISTS vm_runner_bwip
+(
+    l1_batch_number       BIGINT    NOT NULL PRIMARY KEY,
+    created_at            TIMESTAMP NOT NULL,
+    updated_at            TIMESTAMP NOT NULL,
+    time_taken            TIME
+);

core/lib/dal/src/vm_runner_dal.rs

@@ -84,4 +84,80 @@ impl VmRunnerDal<'_, '_> {
         .await?;
         Ok(())
     }
+
+    pub async fn get_bwip_latest_processed_batch(
+        &mut self,
+        default_batch: L1BatchNumber,
+    ) -> DalResult<L1BatchNumber> {
+        let row = sqlx::query!(
+            r#"
+            SELECT
+                COALESCE(MAX(l1_batch_number), $1) AS "last_processed_l1_batch!"
+            FROM
+                vm_runner_bwip
+            "#,
+            default_batch.0 as i32
+        )
+        .instrument("get_bwip_latest_processed_batch")
+        .report_latency()
+        .fetch_one(self.storage)
+        .await?;
+        Ok(L1BatchNumber(row.last_processed_l1_batch as u32))
+    }
+
+    pub async fn get_bwip_last_ready_batch(
+        &mut self,
+        default_batch: L1BatchNumber,
+        window_size: u32,
+    ) -> DalResult<L1BatchNumber> {
+        let row = sqlx::query!(
+            r#"
+            WITH
+                available_batches AS (
+                    SELECT
+                        MAX(number) AS "last_batch"
+                    FROM
+                        l1_batches
+                ),
+                processed_batches AS (
+                    SELECT
+                        COALESCE(MAX(l1_batch_number), $1) + $2 AS "last_ready_batch"
+                    FROM
+                        vm_runner_bwip
+                )
+            SELECT
+                LEAST(last_batch, last_ready_batch) AS "last_ready_batch!"
+            FROM
+                available_batches
+                FULL JOIN processed_batches ON TRUE
+            "#,
+            default_batch.0 as i32,
+            window_size as i32
+        )
+        .instrument("get_bwip_last_ready_batch")
+        .report_latency()
+        .fetch_one(self.storage)
+        .await?;
+        Ok(L1BatchNumber(row.last_ready_batch as u32))
+    }
+
+    pub async fn mark_bwip_batch_as_completed(
+        &mut self,
+        l1_batch_number: L1BatchNumber,
+    ) -> DalResult<()> {
+        sqlx::query!(
+            r#"
+            INSERT INTO
+                vm_runner_bwip (l1_batch_number, created_at, updated_at)
+            VALUES
+                ($1, NOW(), NOW())
+            "#,
+            i64::from(l1_batch_number.0),
+        )
+        .instrument("mark_bwip_batch_as_completed")
+        .report_latency()
+        .execute(self.storage)
+        .await?;
+        Ok(())
+    }
 }

core/lib/object_store/src/objects.rs

@@ -86,7 +86,6 @@ impl StoredObject for SnapshotFactoryDependencies {
             .map_err(From::from)
     }
 }
-
 impl StoredObject for SnapshotStorageLogsChunk {
     const BUCKET: Bucket = Bucket::StorageSnapshot;
     type Key<'a> = SnapshotStorageLogsStorageKey;

core/lib/prover_interface/Cargo.toml

@@ -12,6 +12,7 @@ categories.workspace = true
 [dependencies]
 zksync_object_store.workspace = true
 zksync_types.workspace = true
+zksync_state.workspace = true
 
 # We can use the newest api to send proofs to L1.
 circuit_sequencer_api_1_5_0.workspace = true

core/lib/prover_interface/src/inputs.rs

@@ -1,9 +1,13 @@
-use std::{convert::TryInto, fmt::Debug};
+use std::{collections::HashMap, convert::TryInto, fmt::Debug};
 
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, Bytes};
 use zksync_object_store::{serialize_using_bincode, Bucket, StoredObject};
-use zksync_types::{L1BatchNumber, H256, U256};
+pub use zksync_state::WitnessStorage;
+use zksync_types::{
+    block::L1BatchHeader, commitment::L1BatchWithMetadata, witness_block_state::WitnessBlockState,
+    L1BatchNumber, ProtocolVersionId, H256, U256,
+};
 
 const HASH_LEN: usize = H256::len_bytes();
 
@@ -132,16 +136,47 @@ impl PrepareBasicCircuitsJob {
     }
 }
 
-/// Enriched `PrepareBasicCircuitsJob`. All the other fields are taken from the `l1_batches` table.
-#[derive(Debug, Clone)]
-pub struct BasicCircuitWitnessGeneratorInput {
-    pub block_number: L1BatchNumber,
-    pub previous_block_hash: H256,
-    pub previous_block_timestamp: u64,
-    pub block_timestamp: u64,
-    pub used_bytecodes_hashes: Vec<U256>,
+pub struct VMRunWitnessInputData {
+    pub l1_batch_header: L1BatchHeader,
+    pub previous_batch_with_metadata: L1BatchWithMetadata,
+    pub used_bytecodes: HashMap<U256, Vec<[u8; 32]>>,
     pub initial_heap_content: Vec<(usize, U256)>,
-    pub merkle_paths_input: PrepareBasicCircuitsJob,
+    pub protocol_version: ProtocolVersionId,
+
+    pub bootloader_code: Vec<[u8; 32]>,
+    pub default_account_code_hash: U256,
+    pub storage_refunds: Vec<u32>,
+    pub pubdata_costs: Option<Vec<i32>>,
+    pub witness_block_state: WitnessBlockState,
+}
+
+impl StoredObject for VMRunWitnessInputData {
+    const BUCKET: Bucket = Bucket::WitnessInput;
+
+    type Key<'a> = L1BatchNumber;
+
+    fn encode_key(key: Self::Key<'_>) -> String {
+        format!("vm_run_data_{key}.bin")
+    }
+
+    serialize_using_bincode!();
+}
+
+pub struct WitnessInputData {
+    pub vm_run_data: VMRunWitnessInputData,
+    pub merkle_paths: PrepareBasicCircuitsJob,
+}
+
+impl StoredObject for WitnessInputData {
+    const BUCKET: Bucket = Bucket::WitnessInput;
+
+    type Key<'a> = L1BatchNumber;
+
+    fn encode_key(key: Self::Key<'_>) -> String {
+        format!("witness_inputs_{key}.bin")
+    }
+
+    serialize_using_bincode!();
 }
 
 #[cfg(test)]

core/lib/state/src/lib.rs

@@ -17,32 +17,34 @@ use zksync_types::{
     H256,
 };
 
-mod cache;
-mod catchup;
-mod in_memory;
-mod postgres;
-mod rocksdb;
-mod shadow_storage;
-mod storage_factory;
-mod storage_view;
-#[cfg(test)]
-mod test_utils;
-
 pub use self::{
     cache::sequential_cache::SequentialCache,
     catchup::{AsyncCatchupTask, RocksdbCell},
-    in_memory::InMemoryStorage,
     // Note, that `test_infra` of the bootloader tests relies on this value to be exposed
+    in_memory::InMemoryStorage,
     in_memory::IN_MEMORY_STORAGE_DEFAULT_NETWORK_ID,
     postgres::{PostgresStorage, PostgresStorageCaches, PostgresStorageCachesTask},
     rocksdb::{
         RocksdbStorage, RocksdbStorageBuilder, RocksdbStorageOptions, StateKeeperColumnFamily,
     },
     shadow_storage::ShadowStorage,
     storage_factory::{BatchDiff, PgOrRocksdbStorage, ReadStorageFactory, RocksdbWithMemory},
-    storage_view::{StorageView, StorageViewMetrics},
+    storage_view::{StorageView, StorageViewCache, StorageViewMetrics},
+    witness::WitnessStorage,
 };
 
+mod cache;
+mod catchup;
+mod in_memory;
+mod postgres;
+mod rocksdb;
+mod shadow_storage;
+mod storage_factory;
+mod storage_view;
+#[cfg(test)]
+mod test_utils;
+mod witness;
+
 /// Functionality to read from the VM storage.
 pub trait ReadStorage: fmt::Debug {
     /// Read value of the key.

core/lib/state/src/storage_view.rs

@@ -45,14 +45,33 @@ pub struct StorageView<S> {
     storage_handle: S,
     // Used for caching and to get the list/count of modified keys
     modified_storage_keys: HashMap<StorageKey, StorageValue>,
+    cache: StorageViewCache,
+    metrics: StorageViewMetrics,
+}
+
+#[derive(Debug, Clone)]
+pub struct StorageViewCache {
     // Used purely for caching
     read_storage_keys: HashMap<StorageKey, StorageValue>,
     // Cache for `contains_key()` checks. The cache is only valid within one L1 batch execution.
-    initial_writes_cache: HashMap<StorageKey, bool>,
-    metrics: StorageViewMetrics,
+    initial_writes: HashMap<StorageKey, bool>,
+}
+
+impl StorageViewCache {
+    pub fn read_storage_keys(&self) -> HashMap<StorageKey, StorageValue> {
+        self.read_storage_keys.clone()
+    }
+
+    pub fn initial_writes(&self) -> HashMap<StorageKey, bool> {
+        self.initial_writes.clone()
+    }
 }
 
 impl<S> StorageView<S> {
+    pub fn cache(&self) -> StorageViewCache {
+        self.cache.clone()
+    }
+
     /// Returns the modified storage keys
     pub fn modified_storage_keys(&self) -> &HashMap<StorageKey, StorageValue> {
         &self.modified_storage_keys
@@ -90,8 +109,10 @@ impl<S: ReadStorage + fmt::Debug> StorageView<S> {
         Self {
             storage_handle,
             modified_storage_keys: HashMap::new(),
-            read_storage_keys: HashMap::new(),
-            initial_writes_cache: HashMap::new(),
+            cache: StorageViewCache {
+                read_storage_keys: HashMap::new(),
+                initial_writes: HashMap::new(),
+            },
             metrics: StorageViewMetrics::default(),
         }
     }
@@ -102,10 +123,10 @@ impl<S: ReadStorage + fmt::Debug> StorageView<S> {
         let cached_value = self
             .modified_storage_keys
             .get(key)
-            .or_else(|| self.read_storage_keys.get(key));
+            .or_else(|| self.cache.read_storage_keys.get(key));
         cached_value.copied().unwrap_or_else(|| {
             let value = self.storage_handle.read_value(key);
-            self.read_storage_keys.insert(*key, value);
+            self.cache.read_storage_keys.insert(*key, value);
             self.metrics.time_spent_on_storage_missed += started_at.elapsed();
             self.metrics.storage_invocations_missed += 1;
             value
@@ -114,8 +135,8 @@ impl<S: ReadStorage + fmt::Debug> StorageView<S> {
 
     fn cache_size(&self) -> usize {
         self.modified_storage_keys.len() * mem::size_of::<(StorageKey, StorageValue)>()
-            + self.initial_writes_cache.len() * mem::size_of::<(StorageKey, bool)>()
-            + self.read_storage_keys.len() * mem::size_of::<(StorageKey, StorageValue)>()
+            + self.cache.initial_writes.len() * mem::size_of::<(StorageKey, bool)>()
+            + self.cache.read_storage_keys.len() * mem::size_of::<(StorageKey, StorageValue)>()
     }
 
     /// Returns the current metrics.
@@ -153,11 +174,11 @@ impl<S: ReadStorage + fmt::Debug> ReadStorage for StorageView<S> {
     /// Only keys contained in the underlying storage will return `false`. If a key was
     /// inserted using [`Self::set_value()`], it will still return `true`.
     fn is_write_initial(&mut self, key: &StorageKey) -> bool {
-        if let Some(&is_write_initial) = self.initial_writes_cache.get(key) {
+        if let Some(&is_write_initial) = self.cache.initial_writes.get(key) {
             is_write_initial
         } else {
             let is_write_initial = self.storage_handle.is_write_initial(key);
-            self.initial_writes_cache.insert(*key, is_write_initial);
+            self.cache.initial_writes.insert(*key, is_write_initial);
             is_write_initial
         }
     }
@@ -173,7 +194,7 @@ impl<S: ReadStorage + fmt::Debug> ReadStorage for StorageView<S> {
 
 impl<S: ReadStorage + fmt::Debug> WriteStorage for StorageView<S> {
     fn read_storage_keys(&self) -> &HashMap<StorageKey, StorageValue> {
-        &self.read_storage_keys
+        &self.cache.read_storage_keys
     }
 
     fn set_value(&mut self, key: StorageKey, value: StorageValue) -> StorageValue {