commitment.rs

//! Deterministic state-change helpers for the example chain.
//!
//! `StateChanges` uses `BTreeMap` so the encoded form is canonical and deterministic.

use std::collections::BTreeMap;

use alloy_evm::revm::primitives::{Address, B256, U256};
use bytes::{Buf, BufMut};
use commonware_codec::{EncodeSize, Error as CodecError, RangeCfg, Read, ReadExt, Write};

#[derive(Clone, Copy, Debug)]
/// Limits used when decoding deterministic state changes.
pub struct StateChangesCfg {
    /// Maximum number of touched accounts allowed in a delta.
    pub max_accounts: usize,
    /// Maximum number of storage slots that can be decoded per account.
    pub max_storage_slots: usize,
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Canonical representation of a touched account's post-transaction state.
pub struct AccountChange {
    /// True if the account was accessed during execution.
    pub touched: bool,
    /// True if the account was newly created during execution.
    pub created: bool,
    /// True if the account was self-destructed during execution.
    pub selfdestructed: bool,
    /// Final nonce after execution.
    pub nonce: u64,
    /// Final balance after execution.
    pub balance: U256,
    /// Hash of the account bytecode after execution.
    pub code_hash: B256,
    /// Storage slot updates keyed by slot.
    pub storage: BTreeMap<U256, U256>,
}

#[derive(Clone, Debug, PartialEq, Eq, Default)]
/// Canonical per-transaction state delta.
pub struct StateChanges {
    /// Map of account changes keyed by address.
    pub accounts: BTreeMap<Address, AccountChange>,
}

impl StateChanges {
    /// Returns true when no accounts were touched.
    pub fn is_empty(&self) -> bool {
        self.accounts.is_empty()
    }
}

impl Write for AccountChange {
    fn write(&self, buf: &mut impl BufMut) {
        self.touched.write(buf);
        self.created.write(buf);
        self.selfdestructed.write(buf);
        self.nonce.write(buf);
        write_u256(&self.balance, buf);
        write_b256(&self.code_hash, buf);

        self.storage.len().write(buf);
        for (slot, value) in self.storage.iter() {
            write_u256(slot, buf);
            write_u256(value, buf);
        }
    }
}

impl EncodeSize for AccountChange {
    fn encode_size(&self) -> usize {
        self.touched.encode_size()
            + self.created.encode_size()
            + self.selfdestructed.encode_size()
            + self.nonce.encode_size()
            + 32
            + 32
            + self.storage.len().encode_size()
            + self.storage.len() * (32 + 32)
    }
}

impl Read for AccountChange {
    type Cfg = StateChangesCfg;

    fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, CodecError> {
        let touched = bool::read(buf)?;
        let created = bool::read(buf)?;
        let selfdestructed = bool::read(buf)?;
        let nonce = u64::read(buf)?;
        let balance = read_u256(buf)?;
        let code_hash = read_b256(buf)?;

        let slots = usize::read_cfg(buf, &RangeCfg::new(0..=cfg.max_storage_slots))?;
        let mut storage = BTreeMap::new();
        for _ in 0..slots {
            let slot = read_u256(buf)?;
            let value = read_u256(buf)?;
            storage.insert(slot, value);
        }

        Ok(Self { touched, created, selfdestructed, nonce, balance, code_hash, storage })
    }
}

impl Write for StateChanges {
    fn write(&self, buf: &mut impl BufMut) {
        self.accounts.len().write(buf);
        for (address, change) in self.accounts.iter() {
            write_address(address, buf);
            change.write(buf);
        }
    }
}

impl EncodeSize for StateChanges {
    fn encode_size(&self) -> usize {
        self.accounts.len().encode_size()
            + self.accounts.values().map(|change| 20 + change.encode_size()).sum::<usize>()
    }
}

impl Read for StateChanges {
    type Cfg = StateChangesCfg;

    fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, CodecError> {
        let accounts = usize::read_cfg(buf, &RangeCfg::new(0..=cfg.max_accounts))?;
        let mut map = BTreeMap::new();
        for _ in 0..accounts {
            let address = read_address(buf)?;
            let change = AccountChange::read_cfg(buf, cfg)?;
            map.insert(address, change);
        }
        Ok(Self { accounts: map })
    }
}

fn write_address(value: &Address, buf: &mut impl BufMut) {
    buf.put_slice(value.as_slice());
}

fn read_address(buf: &mut impl Buf) -> Result<Address, CodecError> {
    if buf.remaining() < 20 {
        return Err(CodecError::EndOfBuffer);
    }
    let mut out = [0u8; 20];
    buf.copy_to_slice(&mut out);
    Ok(Address::from(out))
}

fn write_b256(value: &B256, buf: &mut impl BufMut) {
    buf.put_slice(value.as_slice());
}

fn read_b256(buf: &mut impl Buf) -> Result<B256, CodecError> {
    if buf.remaining() < 32 {
        return Err(CodecError::EndOfBuffer);
    }
    let mut out = [0u8; 32];
    buf.copy_to_slice(&mut out);
    Ok(B256::from(out))
}

fn write_u256(value: &U256, buf: &mut impl BufMut) {
    buf.put_slice(&value.to_be_bytes::<32>());
}

fn read_u256(buf: &mut impl Buf) -> Result<U256, CodecError> {
    if buf.remaining() < 32 {
        return Err(CodecError::EndOfBuffer);
    }
    let mut out = [0u8; 32];
    buf.copy_to_slice(&mut out);
    Ok(U256::from_be_bytes(out))
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use alloy_evm::revm::primitives::{Address, B256, U256};
    use commonware_codec::{Decode as _, Encode as _};

    use super::*;

    fn cfg() -> StateChangesCfg {
        StateChangesCfg { max_accounts: 16, max_storage_slots: 64 }
    }

    fn sample_changes_order_a() -> StateChanges {
        let mut changes = StateChanges::default();

        let mut storage1 = BTreeMap::new();
        storage1.insert(U256::from(2u64), U256::from(200u64));
        storage1.insert(U256::from(1u64), U256::from(100u64));
        changes.accounts.insert(
            Address::from([0x11u8; 20]),
            AccountChange {
                touched: true,
                created: false,
                selfdestructed: false,
                nonce: 7,
                balance: U256::from(1234u64),
                code_hash: B256::from([0xAAu8; 32]),
                storage: storage1,
            },
        );

        let mut storage2 = BTreeMap::new();
        storage2.insert(U256::from(5u64), U256::from(42u64));
        changes.accounts.insert(
            Address::from([0x22u8; 20]),
            AccountChange {
                touched: true,
                created: true,
                selfdestructed: false,
                nonce: 1,
                balance: U256::from(999u64),
                code_hash: B256::from([0xBBu8; 32]),
                storage: storage2,
            },
        );

        changes
    }

    #[test]
    fn test_state_changes_roundtrip() {
        let changes = sample_changes_order_a();
        let encoded = changes.encode();
        let decoded = StateChanges::decode_cfg(encoded, &cfg()).expect("decode changes");
        assert_eq!(changes, decoded);
    }

    #[test]
    fn test_empty_state_changes_roundtrip() {
        let changes = StateChanges::default();
        assert!(changes.is_empty());
        let encoded = changes.encode();
        let decoded = StateChanges::decode_cfg(encoded, &cfg()).expect("decode empty");
        assert_eq!(changes, decoded);
        assert!(decoded.is_empty());
    }

    #[test]
    fn test_account_change_roundtrip() {
        let mut storage = BTreeMap::new();
        storage.insert(U256::from(1u64), U256::from(100u64));
        let change = AccountChange {
            touched: true,
            created: true,
            selfdestructed: false,
            nonce: 42,
            balance: U256::from(1_000_000u64),
            code_hash: B256::repeat_byte(0xAB),
            storage,
        };

        let encoded = change.encode();
        let decoded = AccountChange::decode_cfg(encoded, &cfg()).expect("decode change");
        assert_eq!(change, decoded);
    }

    #[test]
    fn test_account_change_empty_storage() {
        let change = AccountChange {
            touched: false,
            created: false,
            selfdestructed: true,
            nonce: 0,
            balance: U256::ZERO,
            code_hash: B256::ZERO,
            storage: BTreeMap::new(),
        };

        let encoded = change.encode();
        let decoded = AccountChange::decode_cfg(encoded, &cfg()).expect("decode");
        assert_eq!(change, decoded);
    }

    #[test]
    fn test_state_changes_is_empty() {
        let mut changes = StateChanges::default();
        assert!(changes.is_empty());

        changes.accounts.insert(
            Address::ZERO,
            AccountChange {
                touched: true,
                created: false,
                selfdestructed: false,
                nonce: 1,
                balance: U256::from(100u64),
                code_hash: B256::ZERO,
                storage: BTreeMap::new(),
            },
        );
        assert!(!changes.is_empty());
    }

    #[test]
    fn test_state_changes_cfg_debug() {
        let cfg = StateChangesCfg { max_accounts: 10, max_storage_slots: 20 };
        let debug_str = format!("{:?}", cfg);
        assert!(debug_str.contains("max_accounts: 10"));
        assert!(debug_str.contains("max_storage_slots: 20"));
    }
}