tempo/

init_state.rs

//! Initialize state from a binary dump file.
//!
//! This command loads TIP20 storage slots from a binary file and applies them
//! to the genesis state. The binary format is produced by `tempo-xtask generate-state-bloat`.

use std::{
    collections::HashMap,
    fs::File,
    io::{BufReader, Read},
    path::PathBuf,
};

use alloy_primitives::{B256, U256, map::HashSet};
use clap::Parser;
use eyre::{Context as _, ensure};
use reth_chainspec::EthereumHardforks;
use reth_cli_commands::common::{AccessRights, CliNodeTypes, EnvironmentArgs};
use reth_db_api::{
    cursor::{DbCursorRO, DbCursorRW},
    tables,
    transaction::DbTxMut,
};
use reth_ethereum::{chainspec::EthChainSpec, tasks::Runtime};
use reth_primitives_traits::{Account, StorageEntry};
use reth_provider::{BlockNumReader, DatabaseProviderFactory, HashingWriter};
use reth_storage_api::DBProvider;
use tempo_chainspec::spec::TempoChainSpecParser;
use tracing::info;

/// Magic bytes for the state bloat binary format (8 bytes)
const MAGIC: &[u8; 8] = b"TEMPOSB\x00";

/// Expected format version
const VERSION: u16 = 1;

/// Initialize state from a binary dump file.
#[derive(Debug, Parser)]
pub(crate) struct InitFromBinaryDump<C: reth_cli::chainspec::ChainSpecParser = TempoChainSpecParser>
{
    #[command(flatten)]
    env: EnvironmentArgs<C>,

    /// Path to the binary state dump file.
    ///
    /// The file should be generated by `tempo-xtask generate-state-bloat`.
    #[arg(value_name = "BINARY_DUMP_FILE")]
    state: PathBuf,
}

impl<C: reth_cli::chainspec::ChainSpecParser<ChainSpec: EthChainSpec + EthereumHardforks>>
    InitFromBinaryDump<C>
{
    /// Execute the init-from-binary-dump command.
    pub(crate) async fn execute<N>(self, runtime: Runtime) -> eyre::Result<()>
    where
        N: CliNodeTypes<ChainSpec = C::ChainSpec>,
    {
        info!(target: "tempo::cli", "Tempo init-from-binary-dump starting");

        let environment = self.env.init::<N>(AccessRights::RW, runtime)?;
        let provider_factory = environment.provider_factory;

        let provider_rw = provider_factory.database_provider_rw()?;

        // Verify we're at genesis (block 0)
        let last_block = provider_rw.last_block_number()?;
        ensure!(
            last_block == 0,
            "init-from-binary-dump must be run on a freshly initialized database at block 0, \
             but found block {last_block}"
        );

        info!(target: "tempo::cli", path = %self.state.display(), "Loading binary state dump");

        let file = File::open(&self.state)
            .wrap_err_with(|| format!("failed to open {}", self.state.display()))?;
        let mut reader = BufReader::with_capacity(64 * 1024 * 1024, file);

        let mut total_entries = 0u64;
        let mut total_tokens = 0u64;

        // Collect storage entries per address for hashing
        let mut storage_for_hashing: HashMap<alloy_primitives::Address, Vec<StorageEntry>> =
            HashMap::new();

        // Track addresses for account hashing (we need to create empty accounts)
        let mut addresses_seen: HashSet<alloy_primitives::Address> = HashSet::default();

        // Process blocks from binary file
        loop {
            // Try to read header
            let mut header_buf = [0u8; 40];
            match reader.read_exact(&mut header_buf) {
                Ok(()) => {}
                Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break,
                Err(e) => return Err(e).wrap_err("failed to read block header"),
            }

            // Validate magic
            ensure!(
                &header_buf[..8] == MAGIC,
                "invalid magic bytes in block header"
            );

            // Validate version
            let version = u16::from_be_bytes([header_buf[8], header_buf[9]]);
            ensure!(
                version == VERSION,
                "unsupported binary format version {version}, expected {VERSION}"
            );

            // Skip flags (2 bytes at offset 10)

            // Read address (20 bytes at offset 12)
            let mut address_bytes = [0u8; 20];
            address_bytes.copy_from_slice(&header_buf[12..32]);
            let address = alloy_primitives::Address::from(address_bytes);

            // Read pair count (8 bytes at offset 32)
            let pair_count = u64::from_be_bytes(header_buf[32..40].try_into().unwrap());

            info!(
                target: "tempo::cli",
                %address,
                pair_count,
                "Processing token storage"
            );

            addresses_seen.insert(address);

            // Get cursors for plain state tables
            let tx = provider_rw.tx_ref();
            let mut storage_cursor = tx.cursor_dup_write::<tables::PlainStorageState>()?;
            let mut account_cursor = tx.cursor_write::<tables::PlainAccountState>()?;

            // Insert empty account for this address (required for storage to be included in trie)
            // Only insert if the account doesn't already exist from genesis
            if account_cursor.seek_exact(address)?.is_none() {
                account_cursor.upsert(address, &Account::default())?;
            }

            // Collect storage for hashing
            let storage_entries = storage_for_hashing.entry(address).or_default();

            // Read and insert entries
            let mut entry_buf = [0u8; 64];
            let start = std::time::Instant::now();
            let mut last_log = start;
            for i in 0..pair_count {
                reader
                    .read_exact(&mut entry_buf)
                    .wrap_err("failed to read storage entry")?;

                let slot = B256::from_slice(&entry_buf[..32]);
                let value = U256::from_be_bytes::<32>(entry_buf[32..64].try_into().unwrap());

                // Skip zero values (they represent deletion)
                if value.is_zero() {
                    continue;
                }

                let entry = StorageEntry { key: slot, value };

                // Insert into plain storage state
                storage_cursor.upsert(address, &entry)?;

                // Collect for hashed storage
                storage_entries.push(entry);

                total_entries += 1;

                let now = std::time::Instant::now();
                if now.duration_since(last_log) >= std::time::Duration::from_secs(5)
                    || i + 1 == pair_count
                {
                    let pct = ((i + 1) as f64 / pair_count as f64) * 100.0;
                    let elapsed = start.elapsed();
                    let pairs_per_sec = (i + 1) as f64 / elapsed.as_secs_f64();
                    info!(
                        target: "tempo::cli",
                        %address,
                        progress = format_args!("{}/{} ({pct:.0}%)", i + 1, pair_count),
                        elapsed = ?elapsed,
                        pairs_per_sec = pairs_per_sec as u64,
                        "Inserting storage"
                    );
                    last_log = now;
                }
            }

            total_tokens += 1;
        }

        info!(
            target: "tempo::cli",
            total_tokens,
            total_entries,
            "Plain storage state written, now writing hashed state..."
        );

        // Write hashed account entries for addresses that have storage.
        // We create empty accounts (zero balance, zero nonce, no code) for precompile addresses.
        // This is required for the storage to be included in the state trie computation.
        let empty_account = Account::default();
        provider_rw.insert_account_for_hashing(
            addresses_seen
                .iter()
                .map(|addr| (*addr, Some(empty_account))),
        )?;

        info!(
            target: "tempo::cli",
            addresses = addresses_seen.len(),
            "Hashed accounts written"
        );

        // Write hashed storage entries for trie computation
        provider_rw.insert_storage_for_hashing(
            storage_for_hashing
                .into_iter()
                .map(|(addr, entries)| (addr, entries.into_iter())),
        )?;

        info!(target: "tempo::cli", "Hashed storage written");

        // Commit the transaction
        provider_rw.commit()?;

        info!(
            target: "tempo::cli",
            total_tokens,
            total_entries,
            "Binary state dump loaded successfully"
        );

        // Note: The state root will need to be recomputed.
        // The node will compute the actual state root on startup.
        info!(
            target: "tempo::cli",
            "State loaded. The node will compute the state root on startup."
        );

        Ok(())
    }
}