feat: add update of saved total supply on first mint after migration

Author: Raid5594

basic_bootloader/src/bootloader/run_single_interaction.rs

@@ -1,11 +1,16 @@
 use crate::bootloader::errors::BootloaderInterfaceError;
 use crate::bootloader::runner::{run_till_completion, RunnerMemoryBuffers};
+use crypto::sha3::{Digest, Keccak256};
 use errors::BootloaderSubsystemError;
+use system_hooks::addresses_constants::{
+    L2_ASSET_TRACKER_ADDRESS, L2_BASE_TOKEN_HOLDER_ADDRESS, L2_CHAIN_ASSET_HANDLER_ADDRESS,
+};
 use zk_ee::common_structs::system_hooks::HooksStorage;
 use zk_ee::system::errors::subsystem::SubsystemError;
 use zk_ee::system::errors::{runtime::RuntimeError, system::SystemError};
 use zk_ee::system::CallModifier;
 use zk_ee::system::{EthereumLikeTypes, System};
+use zk_ee::utils::Bytes32;
 use zk_ee::{interface_error, internal_error, wrap_error};
 
 use super::*;
@@ -53,6 +58,139 @@ where
                 }
             })?;
 
+        Self::update_saved_total_supply(system, resources)?;
+
+        Ok(())
+    }
+
+    /// Computes keccak256(abi.encode(key, base_slot)) for Solidity mapping slot derivation.
+    fn solidity_mapping_slot(key: &[u8; 32], base_slot: u64) -> Bytes32 {
+        let mut hasher = Keccak256::new();
+        hasher.update(key);
+        hasher.update(&U256::from(base_slot).to_be_bytes::<32>());
+        Bytes32::from_array(hasher.finalize().into())
+    }
+
+    /// Replicates the L2AssetTracker's `_getOrSaveTotalSupply` logic for the base token.
+    ///
+    /// On each mint, checks whether `savedTotalSupply[migrationNumber][baseTokenAssetId]`
+    /// has been recorded in the L2AssetTracker contract. If not, computes the current
+    /// total supply from the L2BaseTokenZKOS contract and writes it.
+    ///
+    /// This is needed because zksync-os mints tokens natively (bypassing the Solidity
+    /// L2BaseToken.mint() which would normally call handleFinalizeBaseTokenBridgingOnL2
+    /// on the L2AssetTracker).
+    fn update_saved_total_supply(
+        system: &mut System<S>,
+        resources: &mut S::Resources,
+    ) -> Result<(), BootloaderSubsystemError>
+    where
+        S::IO: IOSubsystemExt,
+    {
+        let ee = ExecutionEnvironmentType::EVM;
+
+        // 1. Read BASE_TOKEN_ASSET_ID from L2AssetTracker (slot 155)
+        let base_token_asset_id = system.io.storage_read::<false>(
+            ee,
+            resources,
+            &L2_ASSET_TRACKER_ADDRESS,
+            &Bytes32::from_u256_be(&U256::from(155)),
+        )?;
+        if base_token_asset_id.is_zero() {
+            // Not initialized yet (before genesis upgrade), skip
+            return Ok(());
+        }
+
+        // 2. Read migrationNumber[chainId] from L2ChainAssetHandler (base slot 207)
+        let chain_id = system.get_chain_id();
+        let chain_id_bytes = U256::from(chain_id).to_be_bytes::<32>();
+        let migration_number_slot =
+            Self::solidity_mapping_slot(&chain_id_bytes, 207);
+        let migration_number = system.io.storage_read::<false>(
+            ee,
+            resources,
+            &L2_CHAIN_ASSET_HANDLER_ADDRESS,
+            &migration_number_slot,
+        )?;
+
+        // 3. Compute savedTotalSupply[migrationNumber][baseTokenAssetId] slot
+        //    level1 = keccak256(abi.encode(migrationNumber, 156))
+        //    level2 = keccak256(abi.encode(baseTokenAssetId, level1))
+        let migration_number_bytes = migration_number.as_u8_array();
+        let level1_slot =
+            Self::solidity_mapping_slot(&migration_number_bytes, 156);
+        let level1_bytes = level1_slot.as_u8_array();
+        let asset_id_bytes = base_token_asset_id.as_u8_array();
+        let mut hasher = Keccak256::new();
+        hasher.update(&asset_id_bytes);
+        hasher.update(&level1_bytes);
+        let struct_base_slot = Bytes32::from_array(hasher.finalize().into());
+
+        // 4. Read isSaved (bool at struct_base_slot + 0)
+        let is_saved_value = system.io.storage_read::<false>(
+            ee,
+            resources,
+            &L2_ASSET_TRACKER_ADDRESS,
+            &struct_base_slot,
+        )?;
+        if !is_saved_value.is_zero() {
+            // Already saved, nothing to do
+            return Ok(());
+        }
+
+        // 5. Compute totalSupply = _zkosPreV31TotalSupply + (INITIAL_BASE_TOKEN_HOLDER_BALANCE - holderBalance)
+        //    _zkosPreV31TotalSupply is at slot 2 of L2BaseToken (0x800a)
+        let pre_v31_supply_bytes = system.io.storage_read::<false>(
+            ee,
+            resources,
+            &system_hooks::addresses_constants::L2_BASE_TOKEN_ADDRESS,
+            &Bytes32::from_u256_be(&U256::from(2)),
+        )?;
+        let pre_v31_supply = pre_v31_supply_bytes.into_u256_be();
+
+        // INITIAL_BASE_TOKEN_HOLDER_BALANCE = 2^127 - 1
+        let initial_holder_balance: U256 = (U256::from(1) << 127) - U256::from(1);
+
+        // Read L2_BASE_TOKEN_HOLDER native balance
+        let holder_balance_u256 = system
+            .io
+            .get_nominal_token_balance(ee, resources, &L2_BASE_TOKEN_HOLDER_ADDRESS)?;
+
+        let total_supply = pre_v31_supply
+            .checked_add(
+                initial_holder_balance
+                    .checked_sub(holder_balance_u256)
+                    .unwrap_or(U256::ZERO),
+            )
+            .unwrap_or(U256::ZERO);
+
+        // 6. Write savedTotalSupply[migrationNumber][assetId] = {isSaved: true, amount: totalSupply}
+        //    Struct slot 0: isSaved (bool, value = 1)
+        //    Struct slot 1: amount (uint256)
+        system.io.storage_write::<false>(
+            ee,
+            resources,
+            &L2_ASSET_TRACKER_ADDRESS,
+            &struct_base_slot,
+            &Bytes32::from_u256_be(&U256::from(1)), // isSaved = true
+        )?;
+
+        // Compute struct_base_slot + 1 for the amount field
+        let amount_slot_u256 = struct_base_slot.into_u256_be() + U256::from(1);
+        let amount_slot = Bytes32::from_u256_be(&amount_slot_u256);
+        system.io.storage_write::<false>(
+            ee,
+            resources,
+            &L2_ASSET_TRACKER_ADDRESS,
+            &amount_slot,
+            &Bytes32::from_u256_be(&total_supply),
+        )?;
+
+        #[cfg(not(target_arch = "riscv32"))]
+        let _ = system.get_logger().write_fmt(format_args!(
+            "Saved total supply {total_supply:?} for base token in L2AssetTracker\n"
+        ));
+
         Ok(())
     }
 

system_hooks/src/addresses_constants.rs

@@ -42,3 +42,11 @@ pub const COMPRESSOR_ADDRESS: B160 = B160::from_limbs([0x800e, 0, 0]);
 pub const COMPLEX_UPGRADER_ADDRESS: B160 = B160::from_limbs([0x800f, 0, 0]);
 pub const KECCAK_SYSTEM_CONTRACT_ADDRESS: B160 = B160::from_limbs([0x8010, 0, 0]);
 pub const PUBDATA_CHUNK_PUBLISHER_ADDRESS: B160 = B160::from_limbs([0x8011, 0, 0]);
+
+// Built-in user space contracts (offset 0x10000)
+/// L2 Chain Asset Handler address (0x1000a)
+pub const L2_CHAIN_ASSET_HANDLER_ADDRESS: B160 = B160::from_limbs([0x1000a, 0, 0]);
+/// L2 Asset Tracker address (0x1000f)
+pub const L2_ASSET_TRACKER_ADDRESS: B160 = B160::from_limbs([0x1000f, 0, 0]);
+/// L2 Base Token Holder address (0x10011)
+pub const L2_BASE_TOKEN_HOLDER_ADDRESS: B160 = B160::from_limbs([0x10011, 0, 0]);