Implement tree states & hierarchical state DB

lcli/src/main.rs

@@ -15,6 +15,7 @@ mod new_testnet;
 mod parse_ssz;
 mod replace_state_pubkeys;
 mod skip_slots;
+mod state_diff;
 mod transition_blocks;
 
 use clap::{App, Arg, ArgMatches, SubCommand};
@@ -841,6 +842,22 @@ fn main() {
                         .help("Number of repeat runs, useful for benchmarking."),
                 )
         )
+        .subcommand(
+            SubCommand::with_name("state-diff")
+                .about("Compute a state diff for a pair of states")
+                .arg(
+                    Arg::with_name("state1")
+                        .value_name("STATE1")
+                        .takes_value(true)
+                        .help("Path to first SSZ state"),
+                )
+                .arg(
+                    Arg::with_name("state2")
+                        .value_name("STATE2")
+                        .takes_value(true)
+                        .help("Path to second SSZ state"),
+                )
+        )
         .get_matches();
 
     let result = matches
@@ -934,6 +951,8 @@ fn run<T: EthSpec>(
             .map_err(|e| format!("Failed to run indexed-attestations command: {}", e)),
         ("block-root", Some(matches)) => block_root::run::<T>(env, matches)
             .map_err(|e| format!("Failed to run block-root command: {}", e)),
+        ("state-diff", Some(matches)) => state_diff::run::<T>(env, matches)
+            .map_err(|e| format!("Failed to run state-diff command: {}", e)),
         (other, _) => Err(format!("Unknown subcommand {}. See --help.", other)),
     }
 }

lcli/src/replace_state_pubkeys.rs

@@ -42,7 +42,8 @@ pub fn run<T: EthSpec>(testnet_dir: PathBuf, matches: &ArgMatches) -> Result<(),
 
     let mut deposit_tree = DepositDataTree::create(&[], 0, DEPOSIT_TREE_DEPTH);
     let mut deposit_root = Hash256::zero();
-    for (index, validator) in state.validators_mut().iter_mut().enumerate() {
+    let validators = state.validators_mut();
+    for index in 0..validators.len() {
         let (secret, _) =
             recover_validator_secret_from_mnemonic(seed.as_bytes(), index as u32, KeyType::Voting)
                 .map_err(|e| format!("Unable to generate validator key: {:?}", e))?;
@@ -52,11 +53,14 @@ pub fn run<T: EthSpec>(testnet_dir: PathBuf, matches: &ArgMatches) -> Result<(),
 
         eprintln!("{}: {}", index, keypair.pk);
 
-        validator.pubkey = keypair.pk.into();
+        validators
+            .get_mut(index)
+            .unwrap()
+            .replace_pubkey(keypair.pk.into());
 
         // Update the deposit tree.
         let mut deposit_data = DepositData {
-            pubkey: validator.pubkey,
+            pubkey: *validators.get(index).unwrap().pubkey(),
             // Set this to a junk value since it's very time consuming to generate the withdrawal
             // keys and it's not useful for the time being.
             withdrawal_credentials: Hash256::zero(),

lcli/src/skip_slots.rs

@@ -55,7 +55,7 @@ use std::fs::File;
 use std::io::prelude::*;
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
-use types::{BeaconState, CloneConfig, EthSpec, Hash256};
+use types::{BeaconState, EthSpec, Hash256};
 
 const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
 
@@ -121,7 +121,7 @@ pub fn run<T: EthSpec>(env: Environment<T>, matches: &ArgMatches) -> Result<(),
     };
 
     for i in 0..runs {
-        let mut state = state.clone_with(CloneConfig::all());
+        let mut state = state.clone();
 
         let start = Instant::now();
 

lcli/src/state_diff.rs

@@ -0,0 +1,44 @@
+use crate::transition_blocks::load_from_ssz_with;
+use clap::ArgMatches;
+use clap_utils::{parse_optional, parse_required};
+use environment::Environment;
+use eth2::{types::BlockId, BeaconNodeHttpClient, SensitiveUrl, Timeouts};
+use std::path::PathBuf;
+use std::time::{Duration, Instant};
+use store::hdiff::{HDiff, HDiffBuffer};
+use types::{BeaconState, EthSpec, FullPayload, SignedBeaconBlock};
+
+pub fn run<T: EthSpec>(env: Environment<T>, matches: &ArgMatches) -> Result<(), String> {
+    let state1_path: PathBuf = parse_required(matches, "state1")?;
+    let state2_path: PathBuf = parse_required(matches, "state2")?;
+    let spec = &T::default_spec();
+
+    let state1 = load_from_ssz_with(&state1_path, spec, BeaconState::<T>::from_ssz_bytes)?;
+    let state2 = load_from_ssz_with(&state2_path, spec, BeaconState::<T>::from_ssz_bytes)?;
+
+    let buffer1 = HDiffBuffer::from_state(state1.clone());
+    let buffer2 = HDiffBuffer::from_state(state2.clone());
+
+    let t = std::time::Instant::now();
+    let diff = HDiff::compute(&buffer1, &buffer2).unwrap();
+    let elapsed = t.elapsed();
+
+    println!("Diff size");
+    println!("- state: {} bytes", diff.state_diff_len());
+    println!("- balances: {} bytes", diff.balances_diff_len());
+    println!("Computation time: {}ms", elapsed.as_millis());
+
+    // Re-apply.
+    let mut recon_buffer = HDiffBuffer::from_state(state1);
+
+    let t = std::time::Instant::now();
+    diff.apply(&mut recon_buffer).unwrap();
+
+    println!("Diff application time: {}ms", t.elapsed().as_millis());
+
+    let recon = recon_buffer.into_state(spec).unwrap();
+
+    assert_eq!(state2, recon);
+
+    Ok(())
+}

lcli/src/transition_blocks.rs

@@ -83,7 +83,7 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use store::HotColdDB;
-use types::{BeaconState, ChainSpec, CloneConfig, EthSpec, Hash256, SignedBeaconBlock};
+use types::{BeaconState, ChainSpec, EthSpec, Hash256, SignedBeaconBlock};
 
 const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
 
@@ -195,7 +195,10 @@ pub fn run<T: EthSpec>(env: Environment<T>, matches: &ArgMatches) -> Result<(),
     let store = Arc::new(store);
 
     debug!("Building pubkey cache (might take some time)");
-    let validator_pubkey_cache = ValidatorPubkeyCache::new(&pre_state, store)
+    let validator_pubkey_cache = store.immutable_validators.clone();
+    validator_pubkey_cache
+        .write()
+        .import_new_pubkeys(&pre_state)
         .map_err(|e| format!("Failed to create pubkey cache: {:?}", e))?;
 
     /*
@@ -226,8 +229,9 @@ pub fn run<T: EthSpec>(env: Environment<T>, matches: &ArgMatches) -> Result<(),
      */
 
     let mut output_post_state = None;
+    let mut saved_ctxt = None;
     for i in 0..runs {
-        let pre_state = pre_state.clone_with(CloneConfig::all());
+        let pre_state = pre_state.clone();
         let block = block.clone();
 
         let start = Instant::now();
@@ -238,7 +242,8 @@ pub fn run<T: EthSpec>(env: Environment<T>, matches: &ArgMatches) -> Result<(),
             block,
             state_root_opt,
             &config,
-            &validator_pubkey_cache,
+            &*validator_pubkey_cache.read(),
+            &mut saved_ctxt,
             spec,
         )?;
 
@@ -288,9 +293,12 @@ pub fn run<T: EthSpec>(env: Environment<T>, matches: &ArgMatches) -> Result<(),
             .map_err(|e| format!("Unable to write to output file: {:?}", e))?;
     }
 
+    drop(pre_state);
+
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 fn do_transition<T: EthSpec>(
     mut pre_state: BeaconState<T>,
     block_root: Hash256,
@@ -298,6 +306,7 @@ fn do_transition<T: EthSpec>(
     mut state_root_opt: Option<Hash256>,
     config: &Config,
     validator_pubkey_cache: &ValidatorPubkeyCache<EphemeralHarnessType<T>>,
+    saved_ctxt: &mut Option<ConsensusContext<T>>,
     spec: &ChainSpec,
 ) -> Result<BeaconState<T>, String> {
     if !config.exclude_cache_builds {
@@ -339,9 +348,22 @@ fn do_transition<T: EthSpec>(
         .map_err(|e| format!("Unable to build caches: {:?}", e))?;
     debug!("Build all caches (again): {:?}", t.elapsed());
 
-    let mut ctxt = ConsensusContext::new(pre_state.slot())
-        .set_current_block_root(block_root)
-        .set_proposer_index(block.message().proposer_index());
+    let mut ctxt = if let Some(ctxt) = saved_ctxt {
+        ctxt.clone()
+    } else {
+        let mut ctxt = ConsensusContext::new(pre_state.slot())
+            .set_current_block_root(block_root)
+            .set_proposer_index(block.message().proposer_index());
+
+        if config.exclude_cache_builds {
+            ctxt = ctxt.set_epoch_cache(
+                EpochCache::new(&pre_state, spec)
+                    .map_err(|e| format!("unable to build epoch cache: {e:?}"))?,
+            );
+            *saved_ctxt = Some(ctxt.clone());
+        }
+        ctxt
+    };
 
     if !config.no_signature_verification {
         let get_pubkey = move |validator_index| {