arb_storage/

state_ops.rs

use std::collections::HashMap;

use alloy_primitives::{Address, Bytes, U256, address, keccak256};
use arb_storage_errors::{DatabaseError, DatabaseErrorInfo, StorageError};
use revm::Database;

fn db_read_error<E: core::fmt::Display>(err: E) -> StorageError {
    DatabaseError::Read(DatabaseErrorInfo::new(err.to_string())).into()
}

/// ArbOS state address — the fictional account that stores all ArbOS state.
pub const ARBOS_STATE_ADDRESS: Address = address!("A4B05FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF");

/// Filtered transactions state address — a separate account for tracking filtered tx hashes.
pub const FILTERED_TX_STATE_ADDRESS: Address = address!("a4b0500000000000000000000000000000000001");

/// Ensures the account exists in the cache. If the account doesn't exist
/// (database returned None), creates it with default values (nonce=0, balance=0).
fn ensure_cache_account<D: Database>(state: &mut revm::database::State<D>, addr: Address) {
    use revm_database::AccountStatus;

    let _ = state.load_cache_account(addr);

    if let Some(cached) = state.cache.accounts.get_mut(&addr)
        && cached.account.is_none()
    {
        cached.account = Some(revm_database::PlainAccount {
            info: revm_state::AccountInfo {
                balance: U256::ZERO,
                nonce: 0,
                code_hash: keccak256([]),
                code: None,
                account_id: None,
            },
            storage: Default::default(),
        });
        cached.status = AccountStatus::InMemoryChange;
    }
}

/// Reads a storage slot from an arbitrary account, checking cache -> bundle -> database.
pub fn read_storage_at<D: Database>(
    state: &mut revm::database::State<D>,
    account: Address,
    slot: U256,
) -> Result<U256, StorageError> {
    if let Some(cached_acc) = state.cache.accounts.get(&account)
        && let Some(ref account) = cached_acc.account
        && let Some(&value) = account.storage.get(&slot)
    {
        return Ok(value);
    }

    if let Some(acc) = state.bundle_state.state.get(&account)
        && let Some(slot_entry) = acc.storage.get(&slot)
    {
        return Ok(slot_entry.present_value);
    }

    state.database.storage(account, slot).map_err(db_read_error)
}

/// Writes a storage slot to the ArbOS account using the transition mechanism.
///
/// This ensures changes survive merge_transitions() and are properly journaled.
/// Skips no-op writes where value == current value.
pub fn write_arbos_storage<D: Database>(
    state: &mut revm::database::State<D>,
    slot: U256,
    value: U256,
) -> Result<(), StorageError> {
    write_storage_at(state, ARBOS_STATE_ADDRESS, slot, value)
}

/// Writes a storage slot to an arbitrary account using the transition mechanism.
pub fn write_storage_at<D: Database>(
    state: &mut revm::database::State<D>,
    account: Address,
    slot: U256,
    value: U256,
) -> Result<(), StorageError> {
    use revm_database::states::StorageSlot;

    ensure_cache_account(state, account);

    let current_value = {
        state
            .cache
            .accounts
            .get(&account)
            .and_then(|ca| ca.account.as_ref())
            .and_then(|a| a.storage.get(&slot).copied())
    }
    .or_else(|| {
        state
            .bundle_state
            .state
            .get(&account)
            .and_then(|a| a.storage.get(&slot))
            .map(|s| s.present_value)
    });

    let original_value = state
        .database
        .storage(account, slot)
        .map_err(db_read_error)?;

    let prev_value = current_value.unwrap_or(original_value);

    if value == prev_value {
        return Ok(());
    }

    // Storage-only write: AccountInfo is unchanged, so one clone suffices.
    let (info, previous_info, previous_status, current_status) = {
        let cached_acc = match state.cache.accounts.get_mut(&account) {
            Some(acc) => acc,
            None => return Ok(()),
        };

        let previous_status = cached_acc.status;
        let (info, previous_info, had_no_nonce_and_code) = match cached_acc.account.as_ref() {
            Some(a) => {
                let info = a.info.clone();
                let had_no_nonce_and_code = info.has_no_code_and_nonce();
                (Some(info.clone()), Some(info), had_no_nonce_and_code)
            }
            None => (None, None, false),
        };

        if let Some(ref mut account) = cached_acc.account {
            account.storage.insert(slot, value);
        }

        cached_acc.status = cached_acc.status.on_changed(had_no_nonce_and_code);
        let current_status = cached_acc.status;
        (info, previous_info, previous_status, current_status)
    };

    if account == ARBOS_STATE_ADDRESS {
        tracing::trace!(
            target: "arb::storage",
            ?slot,
            ?value,
            ?prev_value,
            ?original_value,
            "write_storage_at applying transition"
        );
    }
    let mut storage_changes: revm_database::StorageWithOriginalValues = HashMap::default();
    storage_changes.insert(slot, StorageSlot::new_changed(original_value, value));

    let transition = revm::database::TransitionAccount {
        info,
        status: current_status,
        previous_info,
        previous_status,
        storage: storage_changes,
        storage_was_destroyed: false,
    };

    state.apply_transition([(account, transition)]);
    Ok(())
}

/// Reads the balance of an account from the state.
pub fn get_account_balance<D: Database>(
    state: &mut revm::database::State<D>,
    addr: Address,
) -> U256 {
    if let Some(cached_acc) = state.cache.accounts.get(&addr)
        && let Some(ref account) = cached_acc.account
    {
        return account.info.balance;
    }

    state
        .database
        .basic(addr)
        .ok()
        .flatten()
        .map(|info| info.balance)
        .unwrap_or(U256::ZERO)
}

/// Sets the nonce of an account, loading it into cache if needed.
pub fn set_account_nonce<D: Database>(
    state: &mut revm::database::State<D>,
    addr: Address,
    nonce: u64,
) {
    ensure_cache_account(state, addr);

    let (previous_info, previous_status, current_info, current_status) = {
        let cached_acc = match state.cache.accounts.get_mut(&addr) {
            Some(acc) => acc,
            None => return,
        };
        let previous_status = cached_acc.status;
        let previous_info = cached_acc.account.as_ref().map(|a| a.info.clone());

        if let Some(ref mut account) = cached_acc.account {
            account.info.nonce = nonce;
        }

        let had_no_nonce_and_code = previous_info
            .as_ref()
            .map(|info| info.has_no_code_and_nonce())
            .unwrap_or_default();
        cached_acc.status = cached_acc.status.on_changed(had_no_nonce_and_code);

        let current_info = cached_acc.account.as_ref().map(|a| a.info.clone());
        let current_status = cached_acc.status;
        (previous_info, previous_status, current_info, current_status)
    };

    let transition = revm::database::TransitionAccount {
        info: current_info,
        status: current_status,
        previous_info,
        previous_status,
        storage: HashMap::default(),
        storage_was_destroyed: false,
    };
    state.apply_transition(vec![(addr, transition)]);
}

/// Sets the code of an account, loading it into cache if needed.
pub fn set_account_code<D: Database>(
    state: &mut revm::database::State<D>,
    addr: Address,
    code: Bytes,
) {
    use revm_state::Bytecode;

    ensure_cache_account(state, addr);
    let code_hash = keccak256(&code);
    let bytecode = Bytecode::new_raw(code);

    let (previous_info, previous_status, current_info, current_status) = {
        let cached_acc = match state.cache.accounts.get_mut(&addr) {
            Some(acc) => acc,
            None => return,
        };
        let previous_status = cached_acc.status;
        let previous_info = cached_acc.account.as_ref().map(|a| a.info.clone());

        if let Some(ref mut account) = cached_acc.account {
            account.info.code_hash = code_hash;
            account.info.code = Some(bytecode);
        }

        let had_no_nonce_and_code = previous_info
            .as_ref()
            .map(|info| info.has_no_code_and_nonce())
            .unwrap_or_default();
        cached_acc.status = cached_acc.status.on_changed(had_no_nonce_and_code);

        let current_info = cached_acc.account.as_ref().map(|a| a.info.clone());
        let current_status = cached_acc.status;
        (previous_info, previous_status, current_info, current_status)
    };

    let transition = revm::database::TransitionAccount {
        info: current_info,
        status: current_status,
        previous_info,
        previous_status,
        storage: HashMap::default(),
        storage_was_destroyed: false,
    };
    state.apply_transition(vec![(addr, transition)]);
}

#[cfg(test)]
mod tests {
    use revm_database::{StateBuilder, states::bundle_state::BundleRetention};

    use super::*;

    /// In-memory database that returns empty for everything.
    #[derive(Default)]
    struct EmptyDb;

    impl Database for EmptyDb {
        type Error = std::convert::Infallible;
        fn basic(
            &mut self,
            _address: Address,
        ) -> Result<Option<revm_state::AccountInfo>, Self::Error> {
            Ok(None)
        }
        fn code_by_hash(
            &mut self,
            _code_hash: alloy_primitives::B256,
        ) -> Result<revm_state::Bytecode, Self::Error> {
            Ok(revm_state::Bytecode::default())
        }
        fn storage(&mut self, _address: Address, _index: U256) -> Result<U256, Self::Error> {
            Ok(U256::ZERO)
        }
        fn block_hash(&mut self, _number: u64) -> Result<alloy_primitives::B256, Self::Error> {
            Ok(alloy_primitives::B256::ZERO)
        }
    }

    fn make_state() -> revm::database::State<EmptyDb> {
        StateBuilder::new()
            .with_database(EmptyDb)
            .with_bundle_update()
            .build()
    }

    #[test]
    fn test_write_storage_at_creates_transition() {
        let mut state = make_state();
        let slot = U256::from(42);
        let value = U256::from(12345);

        write_storage_at(&mut state, ARBOS_STATE_ADDRESS, slot, value).unwrap();

        // Verify value is in cache.
        let cached = state.cache.accounts.get(&ARBOS_STATE_ADDRESS).unwrap();
        let stored = cached
            .account
            .as_ref()
            .unwrap()
            .storage
            .get(&slot)
            .copied()
            .unwrap();
        assert_eq!(stored, value, "Value should be in cache");

        // Merge transitions into bundle.
        state.merge_transitions(BundleRetention::Reverts);
        let bundle = state.take_bundle();

        // Verify value is in bundle.
        let bundle_acct = bundle
            .state
            .get(&ARBOS_STATE_ADDRESS)
            .expect("ArbOS account should be in bundle after merge");
        let bundle_slot = bundle_acct
            .storage
            .get(&slot)
            .expect("Slot should be in bundle storage");
        assert_eq!(
            bundle_slot.present_value, value,
            "Bundle present_value should match"
        );
    }

    #[test]
    fn test_write_zero_value_is_noop_for_new_slot() {
        let mut state = make_state();
        let slot = U256::from(42);

        // Writing 0 to a slot that doesn't exist (DB returns 0) should be a no-op.
        write_storage_at(&mut state, ARBOS_STATE_ADDRESS, slot, U256::ZERO).expect("write zero");

        // After merge, the slot should NOT be in the bundle.
        state.merge_transitions(BundleRetention::Reverts);
        let bundle = state.take_bundle();

        // Account might or might not be in bundle, but the slot should not.
        if let Some(acct) = bundle.state.get(&ARBOS_STATE_ADDRESS) {
            assert!(
                !acct.storage.contains_key(&slot),
                "Slot written with zero should not appear in bundle"
            );
        }
    }

    #[test]
    fn test_write_survives_multiple_transitions() {
        let mut state = make_state();
        let slot_a = U256::from(10);
        let slot_b = U256::from(20);

        write_storage_at(&mut state, ARBOS_STATE_ADDRESS, slot_a, U256::from(100))
            .expect("write A");
        write_storage_at(&mut state, ARBOS_STATE_ADDRESS, slot_b, U256::from(200))
            .expect("write B");

        // Merge and check both survive.
        state.merge_transitions(BundleRetention::Reverts);
        let bundle = state.take_bundle();

        let acct = bundle
            .state
            .get(&ARBOS_STATE_ADDRESS)
            .expect("ArbOS account should be in bundle");
        assert_eq!(
            acct.storage.get(&slot_a).unwrap().present_value,
            U256::from(100),
            "Slot A should survive merge"
        );
        assert_eq!(
            acct.storage.get(&slot_b).unwrap().present_value,
            U256::from(200),
            "Slot B should survive merge"
        );
    }

    #[test]
    fn test_read_after_write_returns_written_value() {
        let mut state = make_state();
        let slot = U256::from(42);
        let value = U256::from(99999);

        write_storage_at(&mut state, ARBOS_STATE_ADDRESS, slot, value).expect("write");

        let read_val = read_storage_at(&mut state, ARBOS_STATE_ADDRESS, slot).expect("read");
        assert_eq!(read_val, value, "Read should return written value");
    }

    /// Simulates the real block execution flow:
    /// 1. StartBlock internal tx writes slot A (baseFee) via write_storage_at
    /// 2. EVM commit for internal tx (empty state)
    /// 3. EVM commit for user tx (modifies different accounts)
    /// 4. Post-commit hook writes slot B (gasBacklog) via write_storage_at
    /// 5. Merge transitions Both slots should survive in the bundle.
    #[test]
    fn test_write_survives_evm_commit_flow() {
        let mut state = make_state();
        let slot_basefee = U256::from(10);
        let slot_backlog = U256::from(20);

        write_storage_at(
            &mut state,
            ARBOS_STATE_ADDRESS,
            slot_basefee,
            U256::from(100_000_000),
        )
        .expect("baseFee write");

        // Step 2: EVM commit for internal tx (empty state).
        use revm_database::DatabaseCommit;
        let empty_state: alloy_primitives::map::HashMap<Address, revm_state::Account> =
            Default::default();
        state.commit(empty_state);

        // Step 3: EVM commit for user tx (modifies a different account).
        let sender = address!("1111111111111111111111111111111111111111");
        let mut user_changes: alloy_primitives::map::HashMap<Address, revm_state::Account> =
            Default::default();
        // Load sender into cache first so commit doesn't panic.
        let _ = state.load_cache_account(sender);
        let mut sender_acct = revm_state::Account::default();
        sender_acct.info.balance = U256::from(1_000_000);
        sender_acct.info.nonce = 1;
        sender_acct.mark_touch();
        user_changes.insert(sender, sender_acct);
        state.commit(user_changes);

        write_storage_at(
            &mut state,
            ARBOS_STATE_ADDRESS,
            slot_backlog,
            U256::from(540_000),
        )
        .expect("gasBacklog write");

        // Step 5: Merge transitions.
        state.merge_transitions(BundleRetention::Reverts);
        let bundle = state.take_bundle();

        // Both slots should be in the bundle.
        let acct = bundle
            .state
            .get(&ARBOS_STATE_ADDRESS)
            .expect("ArbOS account should be in bundle");
        assert_eq!(
            acct.storage.get(&slot_basefee).unwrap().present_value,
            U256::from(100_000_000),
            "baseFee slot should survive"
        );
        assert_eq!(
            acct.storage.get(&slot_backlog).unwrap().present_value,
            U256::from(540_000),
            "gasBacklog slot should survive"
        );
    }
}