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::HashSet,
    fs::File,
    io::{BufReader, Read},
    path::PathBuf,
    sync::mpsc,
    thread,
    time::{Duration, Instant},
};

use alloy_primitives::{
    B256, U256, keccak256,
    map::{AddressMap, Entry},
};
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, DbDupCursorRW},
    models::{CompactU256, StorageBeforeTx, storage_sharded_key::StorageShardedKey},
    table::Decompress,
    tables,
    transaction::{DbTx, DbTxMut},
};
use reth_ethereum::{chainspec::EthChainSpec, tasks::Runtime};
use reth_etl::Collector;
use reth_primitives_traits::{Account, StorageEntry};
use reth_provider::{
    BlockNumReader, DBProvider, DatabaseProviderFactory, HashingWriter, RocksDBProviderFactory,
    StaticFileProviderFactory, StaticFileSegment, StorageChangeSetReader, StorageSettingsCache,
    TrieWriter,
};
use reth_trie::{IntermediateStateRootState, StateRootProgress};
use reth_trie_db::DatabaseStateRoot;
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;

/// ETL collector file size (200 MiB per temp file before spilling a new one).
const ETL_FILE_SIZE: usize = 200 * 1024 * 1024;

/// Maximum number of storage entries to hash per worker batch.
const WORKER_CHUNK_SIZE: usize = 100;

/// Bounded channel depth for the hashing worker thread.
const HASH_WORKER_QUEUE_DEPTH: usize = 256;

/// Initialize state from a binary dump file.
#[derive(Debug, Parser)]
pub 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 etl_dir = self
            .env
            .datadir
            .clone()
            .resolve_datadir(self.env.chain.chain())
            .data_dir()
            .join("etl");
        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_blocks = 0u64;

        // Track addresses and their account data for hashing
        let mut accounts_seen: AddressMap<Account> = AddressMap::default();
        let mut genesis_storage_keys = HashSet::new();

        // ETL collectors: accumulate entries sorted, spill to disk when full
        let mut hash_chunk: Vec<(alloy_primitives::Address, B256, CompactU256)> =
            Vec::with_capacity(WORKER_CHUNK_SIZE);
        let mut storage_changeset_collector: Collector<Vec<u8>, CompactU256> =
            Collector::new(ETL_FILE_SIZE, Some(etl_dir.clone()));
        let mut storage_history_collector: Collector<Vec<u8>, CompactU256> =
            Collector::new(ETL_FILE_SIZE, Some(etl_dir.clone()));

        for (index, entry) in provider_rw.storage_changeset(0)? {
            let raw_key = raw_storage_key(index.address(), entry.key);
            genesis_storage_keys.insert(raw_key.clone());
            storage_changeset_collector
                .insert(raw_key, CompactU256::from(entry.value))
                .wrap_err("storage changeset ETL insert of genesis storage failed")?;
        }

        // Single worker thread for keccak hashing: owns the hashed ETL collector, receives
        // batches over a bounded channel, and returns the collector when the sender drops.
        let (hash_tx, hash_rx) = mpsc::sync_channel::<
            Vec<(alloy_primitives::Address, B256, CompactU256)>,
        >(HASH_WORKER_QUEUE_DEPTH);
        let hashed_etl_dir = etl_dir;
        let hash_worker =
            thread::spawn(move || -> eyre::Result<Collector<Vec<u8>, CompactU256>> {
                let mut hashed_collector: Collector<Vec<u8>, CompactU256> =
                    Collector::new(ETL_FILE_SIZE, Some(hashed_etl_dir));
                while let Ok(chunk) = hash_rx.recv() {
                    let mut last_addr = alloy_primitives::Address::ZERO;
                    let mut hashed_addr = B256::ZERO;
                    for (address, slot, value) in chunk {
                        if address != last_addr {
                            last_addr = address;
                            hashed_addr = keccak256(address);
                        }
                        let mut hashed_key = Vec::with_capacity(64);
                        hashed_key.extend_from_slice(hashed_addr.as_slice());
                        hashed_key.extend_from_slice(keccak256(slot).as_slice());
                        hashed_collector
                            .insert(hashed_key, value)
                            .wrap_err("hashed ETL insert failed")?;
                    }
                }
                Ok(hashed_collector)
            });

        // Process blocks from binary file
        loop {
            // Read next block header; EOF means no more blocks.
            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 block"
            );

            // Ensure the account exists in the canonical state table (only on first encounter).
            // The binary dump is chunked: generate-state-bloat writes one block per token
            // per chunk, so the same token address appears in multiple blocks. This entry
            // is Vacant on the first chunk and Occupied on subsequent ones.
            // Preserving the genesis account is critical: TIP20 tokens have bytecode (0xEF)
            // set during genesis, and overwriting with Account::default() would clear the
            // code hash, making the token appear uninitialized.
            if let Entry::Vacant(e) = accounts_seen.entry(address) {
                let hashed_address = keccak256(address);
                let mut account_cursor = provider_rw
                    .tx_ref()
                    .cursor_read::<tables::HashedAccounts>()?;
                let account = match account_cursor.seek_exact(hashed_address)? {
                    Some((_, account)) => account,
                    None => {
                        return Err(eyre::eyre!(
                            "state bloat references account {address} that is missing from genesis"
                        ));
                    }
                };
                e.insert(account);
            }

            // Read entries into the hashed ETL collector.
            let mut entry_buf = [0u8; 64];
            let start = 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 compact_value = CompactU256::from(value);

                let raw_key = raw_storage_key(address, slot);
                if !genesis_storage_keys.contains(&raw_key) {
                    storage_changeset_collector
                        .insert(raw_key.clone(), CompactU256::from(U256::ZERO))
                        .wrap_err("storage changeset ETL insert failed")?;
                    storage_history_collector
                        .insert(raw_key, CompactU256::from(U256::ZERO))
                        .wrap_err("storage history ETL insert failed")?;
                }

                // Queue raw data for parallel hashing
                hash_chunk.push((address, slot, compact_value));

                // Send full batches to the hashing worker thread.
                if hash_chunk.len() >= WORKER_CHUNK_SIZE {
                    let chunk =
                        std::mem::replace(&mut hash_chunk, Vec::with_capacity(WORKER_CHUNK_SIZE));
                    hash_tx.send(chunk).wrap_err("hash worker disconnected")?;
                }

                total_entries += 1;

                log_collection_progress(&address, i, pair_count, start, &mut last_log);
            }

            total_blocks += 1;
        }

        // Send any remaining entries to the worker and join.
        if !hash_chunk.is_empty() {
            hash_tx
                .send(std::mem::take(&mut hash_chunk))
                .wrap_err("hash worker disconnected")?;
        }
        drop(hash_tx);
        let mut hashed_collector = hash_worker
            .join()
            .map_err(|_| eyre::eyre!("hash worker panicked"))??;

        info!(
            target: "tempo::cli",
            total_blocks,
            total_entries,
            "Entries collected, merging genesis hashed storage into ETL collector..."
        );

        // Merge existing genesis hashed storage into the collector.
        {
            let tx = provider_rw.tx_ref();
            let mut cursor = tx.cursor_read::<tables::HashedStorages>()?;
            let mut genesis_count = 0usize;
            let walker = cursor.walk(None)?;
            for row in walker {
                let (hashed_address, entry) = row?;
                let mut key = Vec::with_capacity(64);
                key.extend_from_slice(hashed_address.as_slice());
                key.extend_from_slice(entry.key.as_slice());
                hashed_collector
                    .insert(key, CompactU256::from(entry.value))
                    .wrap_err("ETL insert of genesis hashed storage failed")?;
                genesis_count += 1;
            }
            info!(
                target: "tempo::cli",
                genesis_count,
                "Genesis hashed storage entries merged into collector"
            );
        }

        let storage_settings = provider_rw.cached_storage_settings();
        ensure!(
            storage_settings.storage_v2,
            "init-from-binary-dump only supports storage v2 databases"
        );

        let storage_changeset_factory = provider_factory.clone();
        let storage_changeset_worker = thread::spawn(move || {
            write_storage_changesets(storage_changeset_factory, storage_changeset_collector)
        });

        let storage_history_factory = provider_factory;
        let storage_history_worker = thread::spawn(move || {
            write_storage_history(storage_history_factory, storage_history_collector)
        });

        // Load sorted entries from each ETL collector into its database table.
        // Strategy: iterate the sorted collector, deduplicate consecutive entries with
        // the same composite key, and bulk-insert via append_dup.
        // The table is cleared first so append_dup ordering is guaranteed.
        let total_hashes = hashed_collector.len();
        provider_rw.tx_ref().clear::<tables::HashedStorages>()?;
        let mut hashed_cursor = provider_rw
            .tx_ref()
            .cursor_dup_write::<tables::HashedStorages>()?;
        load_etl_to_cursor(
            &mut hashed_collector,
            total_hashes,
            "hashed storage",
            |k, v| {
                hashed_cursor.append_dup(
                    B256::from_slice(&k[..32]),
                    StorageEntry {
                        key: B256::from_slice(&k[32..]),
                        value: v,
                    },
                )
            },
        )?;
        drop(hashed_cursor);

        info!(
            target: "tempo::cli",
            total_hashes,
            "Storage written, writing hashed accounts..."
        );

        // Write hashed account entries using the real account metadata from plain state.
        // This preserves bytecode_hash for genesis accounts (e.g. TIP20 tokens with 0xEF code).
        provider_rw.insert_account_for_hashing(
            accounts_seen
                .iter()
                .map(|(addr, account)| (*addr, Some(*account))),
        )?;

        storage_changeset_worker
            .join()
            .map_err(|_| eyre::eyre!("storage changeset worker panicked"))??;
        storage_history_worker
            .join()
            .map_err(|_| eyre::eyre!("storage history worker panicked"))??;

        info!(
            target: "tempo::cli",
            addresses = accounts_seen.len(),
            "Hashed accounts written, computing state root and trie nodes..."
        );

        // Rebuild the merkle trie from scratch so the sparse trie cache on
        // block 1 doesn't hit stale genesis nodes and stall on a full rebuild.
        let trie_start = Instant::now();
        provider_rw.tx_ref().clear::<tables::AccountsTrie>()?;
        provider_rw.tx_ref().clear::<tables::StoragesTrie>()?;

        let mut resume: Option<IntermediateStateRootState> = None;
        let mut trie_writes = 0usize;

        // Incrementally compute the merkle root over all hashed accounts/storages.
        let state_root = {
            use reth_trie_db::{
                DatabaseHashedCursorFactory, DatabaseTrieCursorFactory, PackedKeyAdapter,
            };
            type DbStateRoot<'a, TX> = reth_trie::StateRoot<
                DatabaseTrieCursorFactory<&'a TX, PackedKeyAdapter>,
                DatabaseHashedCursorFactory<&'a TX>,
            >;

            // Compute state root in chunks, flushing trie nodes to disk between iterations.
            loop {
                match DbStateRoot::<_>::from_tx(provider_rw.tx_ref())
                    .with_intermediate_state(resume)
                    .root_with_progress()?
                {
                    StateRootProgress::Progress(state, _, updates) => {
                        trie_writes += provider_rw.write_trie_updates(updates)?;
                        info!(
                            target: "tempo::cli",
                            last_key = %state.account_root_state.last_hashed_key,
                            trie_writes,
                            elapsed = ?trie_start.elapsed(),
                            "Flushing trie updates"
                        );
                        resume = Some(*state);
                    }
                    StateRootProgress::Complete(root, _, updates) => {
                        trie_writes += provider_rw.write_trie_updates(updates)?;
                        break root;
                    }
                }
            }
        };

        info!(
            target: "tempo::cli",
            %state_root,
            trie_writes,
            elapsed = ?trie_start.elapsed(),
            "State root computed"
        );

        // Final commit
        provider_rw.commit()?;

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

        Ok(())
    }
}

fn write_storage_changesets<P>(
    provider: P,
    mut collector: Collector<Vec<u8>, CompactU256>,
) -> eyre::Result<()>
where
    P: StaticFileProviderFactory + Send + 'static,
{
    info!(target: "tempo::cli", "Writing storage changesets...");

    provider
        .static_file_provider()
        .delete_segment(StaticFileSegment::StorageChangeSets)?;

    let mut writer = provider.get_static_file_writer(0, StaticFileSegment::StorageChangeSets)?;
    writer.begin_storage_changeset(0)?;
    let total = collector.len();
    load_storage_etl(
        &mut collector,
        total,
        "storage changeset",
        |address, key, value| {
            writer.append_storage_changeset_entry(StorageBeforeTx {
                address,
                key,
                value,
            })?;
            Ok(())
        },
    )?;
    drop(writer);

    Ok(())
}

fn write_storage_history<P>(
    provider: P,
    mut collector: Collector<Vec<u8>, CompactU256>,
) -> eyre::Result<()>
where
    P: RocksDBProviderFactory + Send + 'static,
{
    info!(target: "tempo::cli", "Writing storage history...");

    let rocksdb = provider.rocksdb_provider();
    let mut batch = rocksdb.batch_with_auto_commit();
    let block_zero_history =
        tables::BlockNumberList::new([0]).expect("single block is always sorted");
    let total = collector.len();
    load_storage_etl(
        &mut collector,
        total,
        "storage history",
        |address, key, _| {
            let history_key = StorageShardedKey::last(address, key);
            if batch
                .get::<tables::StoragesHistory>(history_key.clone())?
                .is_none()
            {
                batch.put::<tables::StoragesHistory>(history_key, &block_zero_history)?;
            }
            Ok(())
        },
    )?;
    batch.commit()?;

    Ok(())
}

/// Composite ETL key for unhashed storage, sorted by address then slot.
fn raw_storage_key(address: alloy_primitives::Address, slot: B256) -> Vec<u8> {
    let mut key = Vec::with_capacity(52);
    key.extend_from_slice(address.as_slice());
    key.extend_from_slice(slot.as_slice());
    key
}

fn decode_raw_storage_key(key: &[u8]) -> (alloy_primitives::Address, B256) {
    (
        alloy_primitives::Address::from_slice(&key[..20]),
        B256::from_slice(&key[20..]),
    )
}

/// Iterate a raw storage ETL collector, deduplicate consecutive entries with the
/// same `(address, slot)` key, and call `write` for each unique entry.
fn load_storage_etl(
    collector: &mut Collector<Vec<u8>, CompactU256>,
    total: usize,
    label: &str,
    mut write: impl FnMut(alloy_primitives::Address, B256, U256) -> eyre::Result<()>,
) -> eyre::Result<()> {
    let total = total.max(1);
    let interval = (total / 10).max(1);
    let mut pending: Option<(Vec<u8>, Vec<u8>)> = None;
    for (index, item) in collector.iter()?.enumerate() {
        if index > 0 && index % interval == 0 {
            info!(
                target: "tempo::cli",
                progress = format_args!("{:.2}%", (index as f64 / total as f64) * 100.0),
                "Inserting {label}"
            );
        }

        let (key, value) = item.wrap_err("ETL iteration failed")?;
        if let Some((ref prev_key, ref prev_val)) = pending
            && *prev_key != key
        {
            let (address, storage_key) = decode_raw_storage_key(prev_key);
            write(
                address,
                storage_key,
                CompactU256::decompress_owned(prev_val.clone())?.into(),
            )?;
        }
        pending = Some((key, value));
    }

    if let Some((key, val)) = pending {
        let (address, storage_key) = decode_raw_storage_key(&key);
        write(
            address,
            storage_key,
            CompactU256::decompress_owned(val)?.into(),
        )?;
    }

    Ok(())
}

/// Iterate a sorted ETL collector, deduplicate consecutive entries with the same key
/// (keeping the last value), and call `append` for each unique entry.
fn load_etl_to_cursor(
    collector: &mut Collector<Vec<u8>, CompactU256>,
    total: usize,
    label: &str,
    mut append: impl FnMut(&[u8], U256) -> Result<(), reth_db_api::DatabaseError>,
) -> eyre::Result<()> {
    let interval = (total / 10).max(1);
    let mut pending: Option<(Vec<u8>, Vec<u8>)> = None;
    for (index, item) in collector.iter()?.enumerate() {
        if index > 0 && index % interval == 0 {
            info!(
                target: "tempo::cli",
                progress = format_args!("{:.2}%", (index as f64 / total as f64) * 100.0),
                "Inserting {label}"
            );
        }

        let (key, value) = item.wrap_err("ETL iteration failed")?;
        if let Some((ref prev_key, ref prev_val)) = pending
            && *prev_key != key
        {
            append(
                prev_key,
                CompactU256::decompress_owned(prev_val.clone())?.into(),
            )
            .wrap_err("cursor append failed")?;
        }
        pending = Some((key, value));
    }
    if let Some((key, val)) = pending {
        append(&key, CompactU256::decompress_owned(val)?.into())
            .wrap_err("cursor append failed")?;
    }
    Ok(())
}

/// Log collection progress every 5 seconds and on the final entry.
fn log_collection_progress(
    address: &alloy_primitives::Address,
    index: u64,
    total: u64,
    start: Instant,
    last_log: &mut Instant,
) {
    if last_log.elapsed() >= Duration::from_secs(5) || index + 1 == total {
        let pct = ((index + 1) as f64 / total as f64) * 100.0;
        let elapsed = start.elapsed();
        let pairs_per_sec = (index + 1) as f64 / elapsed.as_secs_f64();
        info!(
            target: "tempo::cli",
            %address,
            progress = format_args!("{}/{} ({pct:.0}%)", index + 1, total),
            elapsed = ?elapsed,
            pairs_per_sec = pairs_per_sec as u64,
            "Collecting storage"
        );
        *last_log = Instant::now();
    }
}