fix: Update prover to use the correct storage oracle (#446)

## What ❔

Making the FRI prover use the correct storage oracle

## Why ❔

<!-- Why are these changes done? What goal do they contribute to? What
are the principles behind them? -->
<!-- Example: PR templates ensure PR reviewers, observers, and future
iterators are in context about the evolution of repos. -->

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zk fmt` and `zk lint`.

prover/witness_generator/src/basic_circuits.rs

@@ -22,7 +22,10 @@ use zksync_prover_fri_types::circuit_definitions::zkevm_circuits::scheduler::blo
 use zksync_prover_fri_types::circuit_definitions::zkevm_circuits::scheduler::input::SchedulerCircuitInstanceWitness;
 use zksync_prover_fri_types::{AuxOutputWitnessWrapper, get_current_pod_name};
 
-use multivm::vm_latest::{constants::MAX_CYCLES_FOR_TX, HistoryDisabled, StorageOracle};
+use crate::storage_oracle::StorageOracle;
+use multivm::vm_latest::{
+    constants::MAX_CYCLES_FOR_TX, HistoryDisabled, StorageOracle as VmStorageOracle,
+};
 use zksync_config::configs::FriWitnessGeneratorConfig;
 use zksync_dal::fri_witness_generator_dal::FriWitnessJobStatus;
 use zksync_dal::ConnectionPool;
@@ -537,6 +540,13 @@ async fn generate_witness(
         .map(|hash| u256_to_h256(*hash))
         .collect();
 
+    let storage_refunds = connection
+        .blocks_dal()
+        .get_storage_refunds(input.block_number)
+        .await
+        .unwrap()
+        .unwrap();
+
     let mut used_bytecodes = connection.storage_dal().get_factory_deps(&hashes).await;
     if input.used_bytecodes_hashes.contains(&account_code_hash) {
         used_bytecodes.insert(account_code_hash, account_bytecode);
@@ -608,10 +618,14 @@ async fn generate_witness(
         let connection = rt_handle
             .block_on(connection_pool.access_storage())
             .unwrap();
+
         let storage = PostgresStorage::new(rt_handle, connection, last_miniblock_number, true);
         let storage_view = StorageView::new(storage).to_rc_ptr();
-        let storage_oracle: StorageOracle<StorageView<PostgresStorage<'_>>, HistoryDisabled> =
-            StorageOracle::new(storage_view.clone());
+
+        let vm_storage_oracle: VmStorageOracle<StorageView<PostgresStorage<'_>>, HistoryDisabled> =
+            VmStorageOracle::new(storage_view.clone());
+        let storage_oracle = StorageOracle::new(vm_storage_oracle, storage_refunds);
+
         zkevm_test_harness::external_calls::run_with_fixed_params(
             Address::zero(),
             BOOTLOADER_ADDRESS,

prover/witness_generator/src/lib.rs

@@ -5,6 +5,7 @@ pub mod leaf_aggregation;
 pub mod node_aggregation;
 pub mod precalculated_merkle_paths_provider;
 pub mod scheduler;
+mod storage_oracle;
 pub mod utils;
 
 #[cfg(test)]

prover/witness_generator/src/main.rs

@@ -29,6 +29,7 @@ mod leaf_aggregation;
 mod node_aggregation;
 mod precalculated_merkle_paths_provider;
 mod scheduler;
+mod storage_oracle;
 mod utils;
 
 #[derive(Debug, StructOpt)]

prover/witness_generator/src/storage_oracle.rs

@@ -0,0 +1,46 @@
+use zksync_types::zkevm_test_harness::zk_evm::abstractions::{
+    RefundType, RefundedAmounts, Storage,
+};
+use zksync_types::{LogQuery, Timestamp};
+
+#[derive(Debug)]
+pub struct StorageOracle<T> {
+    inn: T,
+    storage_refunds: std::vec::IntoIter<u32>,
+}
+
+impl<T> StorageOracle<T> {
+    pub fn new(inn: T, storage_refunds: Vec<u32>) -> Self {
+        Self {
+            inn,
+            storage_refunds: storage_refunds.into_iter(),
+        }
+    }
+}
+
+impl<T: Storage> Storage for StorageOracle<T> {
+    fn estimate_refunds_for_write(
+        &mut self,
+        _monotonic_cycle_counter: u32,
+        _partial_query: &LogQuery,
+    ) -> RefundType {
+        let pubdata_bytes = self.storage_refunds.next().expect("Missing refund");
+        RefundType::RepeatedWrite(RefundedAmounts {
+            pubdata_bytes,
+            ergs: 0,
+        })
+    }
+
+    fn execute_partial_query(&mut self, monotonic_cycle_counter: u32, query: LogQuery) -> LogQuery {
+        self.inn
+            .execute_partial_query(monotonic_cycle_counter, query)
+    }
+
+    fn finish_frame(&mut self, timestamp: Timestamp, panicked: bool) {
+        self.inn.finish_frame(timestamp, panicked)
+    }
+
+    fn start_frame(&mut self, timestamp: Timestamp) {
+        self.inn.start_frame(timestamp)
+    }
+}