zk os integration: use nonces from storage, not from transactions table

core/lib/dal/Cargo.toml

@@ -26,6 +26,8 @@ zksync_consensus_utils.workspace = true
 zksync_protobuf.workspace = true
 zksync_db_connection.workspace = true
 zksync_l1_contract_interface.workspace = true
+zksync_zkos_vm_runner.workspace = true
+
 
 itertools.workspace = true
 thiserror.workspace = true

core/lib/dal/src/storage_web3_dal.rs

@@ -5,13 +5,8 @@ use zksync_db_connection::{
     error::DalResult,
     instrument::{InstrumentExt, Instrumented},
 };
-use zksync_types::{
-    get_code_key, get_nonce_key, h256_to_u256,
-    utils::{decompose_full_nonce, storage_key_for_standard_token_balance},
-    AccountTreeId, Address, L1BatchNumber, L2BlockNumber, Nonce, StorageKey,
-    FAILED_CONTRACT_DEPLOYMENT_BYTECODE_HASH, H256, U256,
-};
-
+use zksync_types::{get_code_key, get_nonce_key, h256_to_u256, utils::{decompose_full_nonce, storage_key_for_standard_token_balance}, AccountTreeId, Address, L1BatchNumber, L2BlockNumber, Nonce, StorageKey, FAILED_CONTRACT_DEPLOYMENT_BYTECODE_HASH, H256, U256, h256_to_address};
+use zksync_zkos_vm_runner::{zkos_nonce_flat_key};
 use crate::{models::storage_block::ResolvedL1BatchForL2Block, Core, CoreDal};
 
 /// Raw bytecode information returned by [`StorageWeb3Dal::get_contract_code_unchecked()`].
@@ -32,9 +27,9 @@ impl StorageWeb3Dal<'_, '_> {
         address: Address,
         block_number: L2BlockNumber,
     ) -> DalResult<U256> {
-        let nonce_key = get_nonce_key(&address);
+        let nonce_key = zkos_nonce_flat_key(address);
         let nonce_value = self
-            .get_historical_value_unchecked(nonce_key.hashed_key(), block_number)
+            .get_historical_value_unchecked(nonce_key, block_number)
             .await?;
         let full_nonce = h256_to_u256(nonce_value);
         Ok(decompose_full_nonce(full_nonce).0)
@@ -45,6 +40,7 @@ impl StorageWeb3Dal<'_, '_> {
         &mut self,
         addresses: &[Address],
     ) -> DalResult<HashMap<Address, Nonce>> {
+        //note: this is only used on EN (tx_sender/proxy.rs), so we don't adopt it
         let nonce_keys: HashMap<_, _> = addresses
             .iter()
             .map(|address| (get_nonce_key(address).hashed_key(), *address))

core/lib/types/Cargo.toml

@@ -18,6 +18,7 @@ zksync_contracts.workspace = true
 zksync_mini_merkle_tree.workspace = true
 zksync_protobuf.workspace = true
 zksync_crypto_primitives.workspace = true
+zk_os_system_hooks.workspace = true
 
 anyhow.workspace = true
 chrono = { workspace = true, features = ["serde"] }

core/lib/types/src/storage/mod.rs

@@ -3,6 +3,7 @@ use core::fmt::Debug;
 use blake2::{Blake2s256, Digest};
 pub use log::*;
 use serde::{Deserialize, Serialize};
+use zk_os_system_hooks::addresses_constants::ACCOUNT_PARTIAL_DATA_STORAGE_ADDRESS;
 use zksync_basic_types::{web3::keccak256, L2ChainId};
 pub use zksync_system_constants::*;
 

core/node/api_server/src/web3/namespaces/eth.rs

@@ -402,15 +402,46 @@ impl EthNamespace {
         address: Address,
         block_id: Option<BlockId>,
     ) -> Result<U256, Web3Error> {
-        //todo: instead of using transactions table, read nonce directly from storage
+        let block_id = block_id.unwrap_or(BlockId::Number(BlockNumber::Pending));
+        self.current_method().set_block_id(block_id);
+
         let mut connection = self.state.acquire_connection().await?;
-        let nonce: Option<U256> = connection
-            .transactions_web3_dal()
-            .zkos_max_nonce_by_initiator_account(address)
+
+        let block_number = self.state.resolve_block(&mut connection, block_id).await?;
+        self.set_block_diff(block_number);
+        let full_nonce = connection
+            .storage_web3_dal()
+            .get_address_historical_nonce(address, block_number)
             .await
             .map_err(DalError::generalize)?;
-        tracing::info!("account_nonce for {:?}: {:?}", address, nonce);
-        Ok((nonce.map(|n| n + 1).unwrap_or(U256::zero())).into())
+
+        // TODO (SMA-1612): currently account nonce is returning always, but later we will
+        //  return account nonce for account abstraction and deployment nonce for non account abstraction.
+        //  Strip off deployer nonce part.
+        let (mut account_nonce, _) = decompose_full_nonce(full_nonce);
+
+        if matches!(block_id, BlockId::Number(BlockNumber::Pending)) {
+            let account_nonce_u64 = u64::try_from(account_nonce)
+                .map_err(|err| anyhow::anyhow!("nonce conversion failed: {err}"))?;
+            account_nonce = if let Some(account_nonce) = self
+                .state
+                .tx_sink()
+                .lookup_pending_nonce(address, account_nonce_u64 as u32)
+                .await?
+            {
+                account_nonce.0.into()
+            } else {
+                // No nonce hint in the sink: get pending nonces from the mempool
+                connection
+                    .transactions_web3_dal()
+                    .next_nonce_by_initiator_account(address, account_nonce_u64)
+                    .await
+                    .map_err(DalError::generalize)?
+            };
+        }
+
+        assert_eq!(nonce.map(|n| n + 1).unwrap_or(U256::zero()), account_nonce);
+        Ok(nonce.map(|n| n + 1).unwrap_or(U256::zero()))
     }
 
     pub async fn get_transaction_impl(

core/node/state_keeper/Cargo.toml

@@ -29,6 +29,7 @@ zksync_node_test_utils.workspace = true
 zksync_vm_executor.workspace = true
 zksync_system_constants.workspace = true
 zksync_base_token_adjuster.workspace = true
+zksync_zkos_vm_runner.workspace = true
 
 anyhow.workspace = true
 async-trait.workspace = true

core/node/state_keeper/src/mempool_actor.rs

@@ -13,7 +13,7 @@ use zksync_node_fee_model::BatchFeeModelInputProvider;
 #[cfg(test)]
 use zksync_types::H256;
 use zksync_types::{get_nonce_key, vm::VmVersion, Address, Nonce, Transaction, U256};
-
+use zksync_zkos_vm_runner::zkos_nonce_flat_key;
 use super::{metrics::KEEPER_METRICS, types::MempoolGuard};
 
 /// Creates a mempool filter for L2 transactions based on the current L1 gas price.
@@ -156,36 +156,32 @@ async fn get_transaction_nonces(
     storage: &mut Connection<'_, Core>,
     transactions: &[&Transaction],
 ) -> anyhow::Result<HashMap<Address, Nonce>> {
-    //todo: revert changes once nonces stored in storage properly
-
-    // let (nonce_keys, address_by_nonce_key): (Vec<_>, HashMap<_, _>) = transactions
-    //     .iter()
-    //     .map(|tx| {
-    //         let address = tx.initiator_account();
-    //         let nonce_key = get_nonce_key(&address).hashed_key();
-    //         (nonce_key, (nonce_key, address))
-    //     })
-    //     .unzip();
-
-    let mut result: HashMap<Address, Nonce> = HashMap::new();
-    for tx in transactions {
-        let address = tx.initiator_account();
-        let maybe_nonce =  storage
-            .transactions_web3_dal()
-            .zkos_max_nonce_by_initiator_account(address)
-            .await
-            .context("failed getting nonce from storage")?;
-        match maybe_nonce {
-            None => {
-                continue;
-            }
-            Some(nonce) => {
-                result.insert(address, Nonce(nonce.as_u32()));
-            }
-        }
-    }
-
-    Ok(result)
+    let (nonce_keys, address_by_nonce_key): (Vec<_>, HashMap<_, _>) = transactions
+        .iter()
+        .map(|tx| {
+            let address = tx.initiator_account();
+            let nonce_key = zkos_nonce_flat_key(address);
+            (nonce_key, (nonce_key, address))
+        })
+        .unzip();
+
+    let nonce_values = storage
+        .storage_web3_dal()
+        .get_values(&nonce_keys)
+        .await
+        .context("failed getting nonces from storage")?;
+
+    let storage_result: HashMap<Address, Nonce> = nonce_values
+        .into_iter()
+        .map(|(nonce_key, nonce_value)| {
+            // `unwrap()` is safe by construction.
+            let be_u32_bytes: [u8; 4] = nonce_value[28..].try_into().unwrap();
+            let nonce = Nonce(u32::from_be_bytes(be_u32_bytes));
+            (address_by_nonce_key[&nonce_key], nonce)
+        })
+        .collect();
+
+    Ok(storage_result)
 }
 
 #[cfg(test)]

core/node/zkos_vm_runner/Cargo.toml

@@ -16,5 +16,6 @@ zksync_types.workspace = true
 zk_os_forward_system.workspace = true
 zk_ee.workspace = true
 zk_os_basic_system.workspace = true
+zk_os_system_hooks.workspace = true
 
 ruint.workspace = true

core/node/zkos_vm_runner/src/lib.rs

@@ -1 +1,17 @@
+use zk_ee::common_structs::derive_flat_storage_key;
+use zk_ee::utils::Bytes32;
+use zk_os_system_hooks::addresses_constants::ACCOUNT_PARTIAL_DATA_STORAGE_ADDRESS;
+use zksync_types::{Address, address_to_h256, H256, h256_to_address};
+use crate::zkos_conversions::{bytes32_to_h256, h256_to_bytes32};
+
 pub mod zkos_conversions;
+
+pub fn zkos_nonce_flat_key(address: Address) -> H256 {
+    let nonce_holder = ACCOUNT_PARTIAL_DATA_STORAGE_ADDRESS;
+    let key = h256_to_bytes32(address_to_h256(&address));
+    bytes32_to_h256(derive_flat_storage_key(
+        &nonce_holder,
+        &key
+    ))
+
+}