immutable validator database factoring (#2297)

* initial immutable validator database factoring

* remove changes from chain_dag: this abstraction properly belongs in beacon_chain_db

* add merging mutable/immutable validator portions; individually test database roundtripping of immutable validators and states-sans-immutable-validators

* update test summaries

* use stew/assign2 instead of Nim assignment

* add reading/writing of immutable validators in chaindag

* remove unused import

* replace chunked k/v store of immutable validators with per-row SQL table storage

* use List instead of HashList

* un-stub some ncli_db code so that it uses

* switch HashArray to array; move BeaconStateNoImmutableValidators from datatypes to beacon_chain_db

* begin only-mutable-part state storage

* uncomment some assigns

* work around https://github.com/nim-lang/Nim/issues/17253

* fix most of the issues/oversights; local sim runs again

* fix test suite by adding missing beaconstate field to copy function

* have ncli bench also store immutable validators

* extract some immutable-validator-specific code from the beacon chain db module

* add more rigorous database state roundtripping, with changing validator sets

* adjust ncli_db to use new schema

* simplify putState/getState by moving all immutable validator accounting into beacon state DB

* remove redundant test case and move code to immutable-beacon-chain module

* more efficient, but still brute-force, mutable+immutable validator merging

* reuse BeaconState in getState

* ensure HashList/HashArray caches are cleared when reusing getState buffers; add ncli_db and a unit test to verify this

* HashList.clear() -> HashList.clearCache()

* only copy incrementally necessary immutable validators

* increase strictness of test cases and fix/work around resulting HashList cache invalidation issues

* remove explanatory scaffolding

* allow for storage of full (with all validators) states for backwards/forwards-compatibility

* adjust DbSeq type usage

* store full, with-validators, state every 64 epochs to enable reverting versions

* reduce memory allocation and intermediate objects in state storage codepath

* eliminate allocation/copying through intermediate BeaconStateNoImmutableValidators objects

* skip benchmarking initial genesis-validator-heavy state store

* always store new-style state and sometimes old-style state

* document intent behind BeaconState/Validator type-punnery

* more accurate failure message on SQLite in-memory database initialization failure

AllTests-mainnet.md

@@ -23,11 +23,14 @@ OK: 1/1 Fail: 0/1 Skip: 0/1
 + empty database [Preset: mainnet]                                                           OK
 + find ancestors [Preset: mainnet]                                                           OK
 + sanity check blocks [Preset: mainnet]                                                      OK
++ sanity check full states [Preset: mainnet]                                                 OK
++ sanity check full states, reusing buffers [Preset: mainnet]                                OK
 + sanity check genesis roundtrip [Preset: mainnet]                                           OK
 + sanity check state diff roundtrip [Preset: mainnet]                                        OK
 + sanity check states [Preset: mainnet]                                                      OK
++ sanity check states, reusing buffers [Preset: mainnet]                                     OK
 ```
-OK: 6/6 Fail: 0/6 Skip: 0/6
+OK: 9/9 Fail: 0/9 Skip: 0/9
 ## Beacon node
 ```diff
 + Compile                                                                                    OK
@@ -279,4 +282,4 @@ OK: 1/1 Fail: 0/1 Skip: 0/1
 OK: 1/1 Fail: 0/1 Skip: 0/1
 
 ---TOTAL---
-OK: 148/157 Fail: 0/157 Skip: 9/157
+OK: 151/160 Fail: 0/160 Skip: 9/160

beacon_chain/beacon_chain_db.nim

@@ -1,11 +1,18 @@
+# beacon_chain
+# Copyright (c) 2018-2021 Status Research & Development GmbH
+# Licensed and distributed under either of
+#   * MIT license (license terms in the root directory or at https://opensource.org/licenses/MIT).
+#   * Apache v2 license (license terms in the root directory or at https://www.apache.org/licenses/LICENSE-2.0).
+# at your option. This file may not be copied, modified, or distributed except according to those terms.
+
 {.push raises: [Defect].}
 
 import
   typetraits, tables,
-  stew/[endians2, io2, objects, results],
+  stew/[assign2, endians2, io2, objects, results],
   serialization, chronicles, snappy,
   eth/db/[kvstore, kvstore_sqlite3],
-  ./networking/network_metadata,
+  ./networking/network_metadata, ./beacon_chain_db_immutable,
   ./spec/[crypto, datatypes, digest, state_transition],
   ./ssz/[ssz_serialization, merkleization],
   ./eth1/merkle_minimal,
@@ -22,6 +29,7 @@ type
     keyspace: int
 
   DepositsSeq = DbSeq[DepositData]
+  ImmutableValidatorsSeq = DbSeq[ImmutableValidatorData]
 
   DepositsMerkleizer* = SszMerkleizer[depositContractLimit]
 
@@ -45,6 +53,12 @@ type
     backend: KvStoreRef
     preset*: RuntimePreset
     genesisDeposits*: DepositsSeq
+
+    # ImmutableValidatorsSeq only stores the total count; it's a proxy for SQL
+    # queries.
+    immutableValidators*: ImmutableValidatorsSeq
+    immutableValidatorsMem*: seq[ImmutableValidatorData]
+
     checkpoint*: proc() {.gcsafe.}
 
   Keyspaces* = enum
@@ -96,6 +110,7 @@ type
       ## a simple append-diff representation helps significantly. Various roots
       ## are stored in a mod-increment pattern across fixed-sized arrays, which
       ## addresses most of the rest of the BeaconState sizes.
+    kHashToStateOnlyMutableValidators
 
   BeaconBlockSummary* = object
     slot*: Slot
@@ -122,6 +137,10 @@ func subkey[N: static int](kind: DbKeyKind, key: array[N, byte]):
 func subkey(kind: type BeaconState, key: Eth2Digest): auto =
   subkey(kHashToState, key.data)
 
+func subkey(
+    kind: type BeaconStateNoImmutableValidators, key: Eth2Digest): auto =
+  subkey(kHashToStateOnlyMutableValidators, key.data)
+
 func subkey(kind: type SignedBeaconBlock, key: Eth2Digest): auto =
   subkey(kHashToBlock, key.data)
 
@@ -213,35 +232,42 @@ proc contains*[K, V](m: DbMap[K, V], key: K): bool =
 template insert*[K, V](t: var Table[K, V], key: K, value: V) =
   add(t, key, value)
 
+proc loadImmutableValidators(db: BeaconChainDB): seq[ImmutableValidatorData] =
+  # TODO not called, but build fails otherwise
+  for i in 0 ..< db.immutableValidators.len:
+    result.add db.immutableValidators.get(i)
+
 proc init*(T: type BeaconChainDB,
            preset: RuntimePreset,
            dir: string,
            inMemory = false): BeaconChainDB =
+  var sqliteStore =
     if inMemory:
-    # TODO
+      SqStoreRef.init("", "test", inMemory = true).expect("working database (out of memory?)")
-    # To support testing, the inMemory store should offer the complete
-    # functionalityof the database-backed one (i.e. tracking of deposits
-    # and validators)
-    T(backend: kvStore MemStoreRef.init(),
-      preset: preset)
     else:
       let s = secureCreatePath(dir)
       doAssert s.isOk # TODO(zah) Handle this in a better way
 
-    let sqliteStore = SqStoreRef.init(
+      SqStoreRef.init(
-      dir, "nbc", Keyspaces, manualCheckpoint = true).expect("working database (disk broken/full?)")
+        dir, "nbc", Keyspaces,
+        manualCheckpoint = true).expect("working database (disk broken/full?)")
 
   # Remove the deposits table we used before we switched
   # to storing only deposit contract checkpoints
   if sqliteStore.exec("DROP TABLE IF EXISTS deposits;").isErr:
     debug "Failed to drop the deposits table"
 
-    var genesisDepositsSeq =
+  var
+    genesisDepositsSeq =
       DbSeq[DepositData].init(sqliteStore, "genesis_deposits")
+    immutableValidatorsSeq =
+      DbSeq[ImmutableValidatorData].init(sqliteStore, "immutable_validators")
 
   T(backend: kvStore sqliteStore,
     preset: preset,
     genesisDeposits: genesisDepositsSeq,
+    immutableValidators: immutableValidatorsSeq,
+    immutableValidatorsMem: loadImmutableValidators(immutableValidatorsSeq),
     checkpoint: proc() = sqliteStore.checkpoint()
     )
 
@@ -324,15 +350,42 @@ proc putBlock*(db: BeaconChainDB, value: SigVerifiedSignedBeaconBlock) =
   db.put(subkey(SignedBeaconBlock, value.root), value)
   db.put(subkey(BeaconBlockSummary, value.root), value.message.toBeaconBlockSummary())
 
-proc putState*(db: BeaconChainDB, key: Eth2Digest, value: BeaconState) =
+proc updateImmutableValidators(
+    db: BeaconChainDB, immutableValidators: var seq[ImmutableValidatorData],
+    validators: auto) =
+  let
+    numValidators = validators.lenu64
+    origNumImmutableValidators = immutableValidators.lenu64
+
+  doAssert immutableValidators.len == db.immutableValidators.len
+
+  if numValidators <= origNumImmutableValidators:
+    return
+
+  for validatorIndex in origNumImmutableValidators ..< numValidators:
+    # This precedes state storage
+    let immutableValidator =
+      getImmutableValidatorData(validators[validatorIndex])
+    db.immutableValidators.add immutableValidator
+    immutableValidators.add immutableValidator
+
+proc putState*(db: BeaconChainDB, key: Eth2Digest, value: var BeaconState) =
   # TODO prune old states - this is less easy than it seems as we never know
   #      when or if a particular state will become finalized.
+  updateImmutableValidators(db, db.immutableValidatorsMem, value.validators)
+  db.put(
+    subkey(BeaconStateNoImmutableValidators, key),
+    cast[ref BeaconStateNoImmutableValidators](addr value)[])
 
-  db.put(subkey(type value, key), value)
+proc putState*(db: BeaconChainDB, value: var BeaconState) =
-
-proc putState*(db: BeaconChainDB, value: BeaconState) =
   db.putState(hash_tree_root(value), value)
 
+proc putStateFull*(db: BeaconChainDB, key: Eth2Digest, value: BeaconState) =
+  db.put(subkey(BeaconState, key), value)
+
+proc putStateFull*(db: BeaconChainDB, value: BeaconState) =
+  db.putStateFull(hash_tree_root(value), value)
+
 proc putStateRoot*(db: BeaconChainDB, root: Eth2Digest, slot: Slot,
     value: Eth2Digest) =
   db.put(subkey(root, slot), value)
@@ -346,6 +399,8 @@ proc delBlock*(db: BeaconChainDB, key: Eth2Digest) =
 
 proc delState*(db: BeaconChainDB, key: Eth2Digest) =
   db.backend.del(subkey(BeaconState, key)).expect("working database (disk broken/full?)")
+  db.backend.del(subkey(BeaconStateNoImmutableValidators, key)).expect(
+    "working database (disk broken/full?)")
 
 proc delStateRoot*(db: BeaconChainDB, root: Eth2Digest, slot: Slot) =
   db.backend.del(subkey(root, slot)).expect("working database (disk broken/full?)")
@@ -389,6 +444,51 @@ proc getBlockSummary*(db: BeaconChainDB, key: Eth2Digest): Opt[BeaconBlockSummar
   if db.get(subkey(BeaconBlockSummary, key), result.get) != GetResult.found:
     result.err()
 
+proc getStateOnlyMutableValidators(
+    db: BeaconChainDB, key: Eth2Digest, output: var BeaconState,
+    rollback: RollbackProc): bool =
+  ## Load state into `output` - BeaconState is large so we want to avoid
+  ## re-allocating it if possible
+  ## Return `true` iff the entry was found in the database and `output` was
+  ## overwritten.
+  ## Rollback will be called only if output was partially written - if it was
+  ## not found at all, rollback will not be called
+  # TODO rollback is needed to deal with bug - use `noRollback` to ignore:
+  #      https://github.com/nim-lang/Nim/issues/14126
+  # TODO RVO is inefficient for large objects:
+  #      https://github.com/nim-lang/Nim/issues/13879
+
+  case db.get(
+    subkey(
+      BeaconStateNoImmutableValidators, key),
+      cast[ref BeaconStateNoImmutableValidators](addr output)[])
+  of GetResult.found:
+    let numValidators = output.validators.len
+    doAssert db.immutableValidatorsMem.len >= numValidators
+
+    output.validators.hashes.setLen(0)
+    for item in output.validators.indices.mitems():
+      item = 0
+
+    for i in 0 ..< numValidators:
+      let
+        # Bypass hash cache invalidation
+        dstValidator = addr output.validators.data[i]
+        srcValidator = addr db.immutableValidatorsMem[i]
+
+      assign(dstValidator.pubkey, srcValidator.pubkey)
+      assign(dstValidator.withdrawal_credentials,
+        srcValidator.withdrawal_credentials)
+
+    output.validators.growHashes()
+
+    true
+  of GetResult.notFound:
+    false
+  of GetResult.corrupted:
+    rollback(output)
+    false
+
 proc getState*(
     db: BeaconChainDB, key: Eth2Digest, output: var BeaconState,
     rollback: RollbackProc): bool =
@@ -402,6 +502,9 @@ proc getState*(
   #      https://github.com/nim-lang/Nim/issues/14126
   # TODO RVO is inefficient for large objects:
   #      https://github.com/nim-lang/Nim/issues/13879
+  if getStateOnlyMutableValidators(db, key, output, rollback):
+    return true
+
   case db.get(subkey(BeaconState, key), output)
   of GetResult.found:
     true
@@ -454,6 +557,8 @@ proc containsBlock*(db: BeaconChainDB, key: Eth2Digest): bool =
   db.backend.contains(subkey(SignedBeaconBlock, key)).expect("working database (disk broken/full?)")
 
 proc containsState*(db: BeaconChainDB, key: Eth2Digest): bool =
+  db.backend.contains(subkey(BeaconStateNoImmutableValidators, key)).expect(
+      "working database  (disk broken/full?)") or
     db.backend.contains(subkey(BeaconState, key)).expect("working database (disk broken/full?)")
 
 proc containsStateDiff*(db: BeaconChainDB, key: Eth2Digest): bool =

beacon_chain/beacon_chain_db_immutable.nim

@@ -0,0 +1,84 @@
+# beacon_chain
+# Copyright (c) 2021 Status Research & Development GmbH
+# Licensed and distributed under either of
+#   * MIT license (license terms in the root directory or at https://opensource.org/licenses/MIT).
+#   * Apache v2 license (license terms in the root directory or at https://www.apache.org/licenses/LICENSE-2.0).
+# at your option. This file may not be copied, modified, or distributed except according to those terms.
+
+{.push raises: [Defect].}
+
+import
+  tables,
+  stew/[assign2, endians2, io2, objects, results],
+  serialization, chronicles,
+  eth/db/[kvstore, kvstore_sqlite3],
+  ./spec/[crypto, datatypes, digest],
+  ./ssz/[ssz_serialization, merkleization],
+  filepath
+
+type 
+  # https://github.com/ethereum/eth2.0-specs/blob/v1.0.1/specs/phase0/beacon-chain.md#beaconstate
+  # Memory-representation-equivalent to a v1.0.1 BeaconState for in-place SSZ reading and writing
+  BeaconStateNoImmutableValidators* = object
+    # Versioning
+    genesis_time*: uint64
+    genesis_validators_root*: Eth2Digest
+    slot*: Slot
+    fork*: Fork
+
+    # History
+    latest_block_header*: BeaconBlockHeader ##\
+    ## `latest_block_header.state_root == ZERO_HASH` temporarily
+
+    block_roots*: HashArray[Limit SLOTS_PER_HISTORICAL_ROOT, Eth2Digest] ##\
+    ## Needed to process attestations, older to newer
+
+    state_roots*: HashArray[Limit SLOTS_PER_HISTORICAL_ROOT, Eth2Digest]
+    historical_roots*: HashList[Eth2Digest, Limit HISTORICAL_ROOTS_LIMIT]
+
+    # Eth1
+    eth1_data*: Eth1Data
+    eth1_data_votes*:
+      HashList[Eth1Data, Limit(EPOCHS_PER_ETH1_VOTING_PERIOD * SLOTS_PER_EPOCH)]
+    eth1_deposit_index*: uint64
+
+    # Registry
+    validators*: HashList[ValidatorStatus, Limit VALIDATOR_REGISTRY_LIMIT]
+    balances*: HashList[uint64, Limit VALIDATOR_REGISTRY_LIMIT]
+
+    # Randomness
+    randao_mixes*: HashArray[Limit EPOCHS_PER_HISTORICAL_VECTOR, Eth2Digest]
+
+    # Slashings
+    slashings*: HashArray[Limit EPOCHS_PER_SLASHINGS_VECTOR, uint64] ##\
+    ## Per-epoch sums of slashed effective balances
+
+    # Attestations
+    previous_epoch_attestations*:
+      HashList[PendingAttestation, Limit(MAX_ATTESTATIONS * SLOTS_PER_EPOCH)]
+    current_epoch_attestations*:
+      HashList[PendingAttestation, Limit(MAX_ATTESTATIONS * SLOTS_PER_EPOCH)]
+
+    # Finality
+    justification_bits*: uint8 ##\
+    ## Bit set for every recent justified epoch
+    ## Model a Bitvector[4] as a one-byte uint, which should remain consistent
+    ## with ssz/hashing.
+
+    previous_justified_checkpoint*: Checkpoint ##\
+    ## Previous epoch snapshot
+
+    current_justified_checkpoint*: Checkpoint
+    finalized_checkpoint*: Checkpoint
+
+static:
+  # Each of these pairs of types has ABI-compatible memory representations, so
+  # that the SSZ serialization can read and write directly from an object with
+  # only mutable portions of BeaconState into a full BeaconState without using
+  # any extra copies.
+  doAssert sizeof(Validator) == sizeof(ValidatorStatus)
+  doAssert sizeof(BeaconState) == sizeof(BeaconStateNoImmutableValidators)
+
+proc loadImmutableValidators*(dbSeq: var auto): seq[ImmutableValidatorData] =
+  for i in 0 ..< dbSeq.len:
+    result.add dbSeq.get(i)

beacon_chain/consensus_object_pools/blockchain_dag.nim

@@ -533,7 +533,7 @@ proc getState(dag: ChainDAGRef, state: var StateData, bs: BlockSlot): bool =
 
   false
 
-proc putState*(dag: ChainDAGRef, state: StateData) =
+proc putState*(dag: ChainDAGRef, state: var StateData) =
   # Store a state and its root
   logScope:
     blck = shortLog(state.blck)
@@ -551,6 +551,11 @@ proc putState*(dag: ChainDAGRef, state: StateData) =
   # transaction to prevent database inconsistencies, but the state loading code
   # is resilient against one or the other going missing
   dag.db.putState(state.data.root, state.data.data)
+
+  # Allow backwards-compatible version rollback with bounded recovery cost
+  if state.data.data.slot.epoch mod 64 == 0:
+    dag.db.putStateFull(state.data.root, state.data.data)
+
   dag.db.putStateRoot(state.blck.root, state.data.data.slot, state.data.root)
 
 func getRef*(dag: ChainDAGRef, root: Eth2Digest): BlockRef =
@@ -1041,7 +1046,7 @@ proc isInitialized*(T: type ChainDAGRef, db: BeaconChainDB): bool =
 
 proc preInit*(
     T: type ChainDAGRef, db: BeaconChainDB,
-    genesisState, tailState: BeaconState, tailBlock: SignedBeaconBlock) =
+    genesisState, tailState: var BeaconState, tailBlock: SignedBeaconBlock) =
   # write a genesis state, the way the ChainDAGRef expects it to be stored in
   # database
   # TODO probably should just init a block pool with the freshly written

beacon_chain/spec/datatypes/base.nim

@@ -581,6 +581,13 @@ type
   # https://github.com/ethereum/eth2.0-specs/blob/v1.0.1/specs/phase0/beacon-chain.md#validator
   ValidatorStatus* = object
     # This is a validator without the expensive, immutable, append-only parts
+    # serialized. They're represented in memory to allow in-place SSZ reading
+    # and writing compatibly with the full Validator object.
+
+    pubkey* {.dontserialize.}: ValidatorPubKey
+
+    withdrawal_credentials* {.dontserialize.}: Eth2Digest ##\
+    ## Commitment to pubkey for withdrawals and transfers
 
     effective_balance*: uint64 ##\
     ## Balance at stake

beacon_chain/statediff.nim

@@ -78,20 +78,19 @@ func replaceOrAddDecodeEth1Votes[T, U](
   for item in votes1:
     votes0.add item
 
-func getMutableValidatorStatus(validator: Validator): ValidatorStatus =
-  ValidatorStatus(
-      effective_balance: validator.effective_balance,
-      slashed: validator.slashed,
-      activation_eligibility_epoch: validator.activation_eligibility_epoch,
-      activation_epoch: validator.activation_epoch,
-      exit_epoch: validator.exit_epoch,
-      withdrawable_epoch: validator.withdrawable_epoch)
-
 func getMutableValidatorStatuses(state: BeaconState):
     List[ValidatorStatus, Limit VALIDATOR_REGISTRY_LIMIT] =
-  # use mapIt + .init(foo)?
+  result.setLen(state.validators.len)
-  for validator in state.validators:
+  for i in 0 ..< state.validators.len:
-    result.add getMutableValidatorStatus(validator)
+    let validator = unsafeAddr state.validators.data[i]
+    assign(result[i].effective_balance, validator.effective_balance)
+    assign(result[i].slashed, validator.slashed)
+    assign(
+      result[i].activation_eligibility_epoch,
+      validator.activation_eligibility_epoch)
+    assign(result[i].activation_epoch, validator.activation_epoch)
+    assign(result[i].exit_epoch, validator.exit_epoch)
+    assign(result[i].withdrawable_epoch, validator.withdrawable_epoch)
 
 func diffStates*(state0, state1: BeaconState): BeaconStateDiff =
   doAssert state1.slot > state0.slot

ncli/ncli_db.nim

@@ -169,7 +169,11 @@ proc cmdBench(conf: DbConf, runtimePreset: RuntimePreset) =
       withTimer(timers[tDbStore]):
         dbBenchmark.putBlock(b)
 
-    if conf.storeStates and state[].data.slot.isEpoch:
+    if state[].data.slot.isEpoch and conf.storeStates:
+      if state[].data.slot.epoch < 2:
+        dbBenchmark.putState(state[].root, state[].data)
+        dbBenchmark.checkpoint()
+      else:
         withTimer(timers[tDbStore]):
           dbBenchmark.putState(state[].root, state[].data)
           dbBenchmark.checkpoint()

tests/test_beacon_chain_db.nim

@@ -1,5 +1,5 @@
 # Nimbus
-# Copyright (c) 2018 Status Research & Development GmbH
+# Copyright (c) 2018-2021 Status Research & Development GmbH
 # Licensed under either of
 #  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or https://www.apache.org/licenses/LICENSE-2.0)
 #  * MIT license ([LICENSE-MIT](LICENSE-MIT) or https://opensource.org/licenses/MIT)
@@ -7,14 +7,17 @@
 
 {.used.}
 
-import  options, unittest, sequtils,
+import  algorithm, options, sequtils, unittest,
-  ../beacon_chain/[beacon_chain_db, extras, interop],
+  ../beacon_chain/[beacon_chain_db, extras, interop, ssz],
-  ../beacon_chain/ssz,
   ../beacon_chain/spec/[
     beaconstate, datatypes, digest, crypto, state_transition, presets],
+  ../beacon_chain/consensus_object_pools/blockchain_dag,
   eth/db/kvstore,
   # test utilies
-  ./testutil, ./testblockutil
+  ./testutil, ./testblockutil, ./teststateutil
+
+when isMainModule:
+  import chronicles # or some random compile error happens...
 
 proc getStateRef(db: BeaconChainDB, root: Eth2Digest): NilableBeaconStateRef =
   # load beaconstate the way the block pool does it - into an existing instance
@@ -71,18 +74,100 @@ suiteReport "Beacon chain DB" & preset():
 
   wrappedTimedTest "sanity check states" & preset():
     var
-      db = BeaconChainDB.init(defaultRuntimePreset, "", inMemory = true)
+      db = makeTestDB(SLOTS_PER_EPOCH)
+      dag = init(ChainDAGRef, defaultRuntimePreset, db)
+      testStates = getTestStates(dag.headState.data)
 
-    let
+    # Ensure transitions beyond just adding validators and increasing slots
-      state = BeaconStateRef()
+    sort(testStates) do (x, y: ref HashedBeaconState) -> int:
-      root = hash_tree_root(state[])
+      cmp($x.root, $y.root)
 
-    db.putState(state[])
+    for state in testStates:
+      db.putState(state[].data)
+      let root = hash_tree_root(state[].data)
 
       check:
         db.containsState(root)
         hash_tree_root(db.getStateRef(root)[]) == root
 
+      db.delState(root)
+      check: not db.containsState(root)
+
+    db.close()
+
+  wrappedTimedTest "sanity check states, reusing buffers" & preset():
+    var
+      db = makeTestDB(SLOTS_PER_EPOCH)
+      dag = init(ChainDAGRef, defaultRuntimePreset, db)
+
+    let stateBuffer = BeaconStateRef()
+    var testStates = getTestStates(dag.headState.data)
+
+    # Ensure transitions beyond just adding validators and increasing slots
+    sort(testStates) do (x, y: ref HashedBeaconState) -> int:
+      cmp($x.root, $y.root)
+
+    for state in testStates:
+      db.putState(state[].data)
+      let root = hash_tree_root(state[].data)
+
+      check:
+        db.getState(root, stateBuffer[], noRollback)
+        db.containsState(root)
+        hash_tree_root(stateBuffer[]) == root
+
+      db.delState(root)
+      check: not db.containsState(root)
+
+    db.close()
+
+  wrappedTimedTest "sanity check full states" & preset():
+    var
+      db = makeTestDB(SLOTS_PER_EPOCH)
+      dag = init(ChainDAGRef, defaultRuntimePreset, db)
+      testStates = getTestStates(dag.headState.data)
+
+    # Ensure transitions beyond just adding validators and increasing slots
+    sort(testStates) do (x, y: ref HashedBeaconState) -> int:
+      cmp($x.root, $y.root)
+
+    for state in testStates:
+      db.putStateFull(state[].data)
+      let root = hash_tree_root(state[].data)
+
+      check:
+        db.containsState(root)
+        hash_tree_root(db.getStateRef(root)[]) == root
+
+      db.delState(root)
+      check: not db.containsState(root)
+
+    db.close()
+
+  wrappedTimedTest "sanity check full states, reusing buffers" & preset():
+    var
+      db = makeTestDB(SLOTS_PER_EPOCH)
+      dag = init(ChainDAGRef, defaultRuntimePreset, db)
+
+    let stateBuffer = BeaconStateRef()
+    var testStates = getTestStates(dag.headState.data)
+
+    # Ensure transitions beyond just adding validators and increasing slots
+    sort(testStates) do (x, y: ref HashedBeaconState) -> int:
+      cmp($x.root, $y.root)
+
+    for state in testStates:
+      db.putStateFull(state[].data)
+      let root = hash_tree_root(state[].data)
+
+      check:
+        db.getState(root, stateBuffer[], noRollback)
+        db.containsState(root)
+        hash_tree_root(stateBuffer[]) == root
+
+      db.delState(root)
+      check: not db.containsState(root)
+
     db.close()
 
   wrappedTimedTest "find ancestors" & preset():

tests/test_statediff.nim

@@ -9,75 +9,15 @@
 
 import
   options, sequtils, unittest,
-  ./testutil,
+  ./testutil, ./teststateutil,
-  ./helpers/math_helpers,
+  ../beacon_chain/spec/[datatypes, digest, helpers, presets],
-  ./mocking/mock_deposits,
-  ../beacon_chain/spec/[beaconstate, datatypes, digest, helpers,
-    state_transition, presets],
   ../beacon_chain/[beacon_node_types, statediff],
   ../beacon_chain/ssz,
-  ../beacon_chain/consensus_object_pools/[blockchain_dag, block_quarantine, block_clearance]
+  ../beacon_chain/consensus_object_pools/[blockchain_dag, block_quarantine]
 
 when isMainModule:
   import chronicles # or some random compile error happens...
 
-proc valid_deposit(state: var BeaconState) =
-  # TODO copy/pasted from foo; refactor
-  const deposit_amount = MAX_EFFECTIVE_BALANCE
-  let validator_index = state.validators.len
-  let deposit = mockUpdateStateForNewDeposit(
-                  state,
-                  uint64 validator_index,
-                  deposit_amount,
-                  flags = {}
-                )
-
-  let pre_val_count = state.validators.len
-  let pre_balance = if validator_index < pre_val_count:
-                      state.balances[validator_index]
-                    else:
-                      0
-  check:
-    process_deposit(defaultRuntimePreset(), state, deposit, {}).isOk
-    state.validators.len == pre_val_count + 1
-    state.balances.len == pre_val_count + 1
-    state.balances[validator_index] == pre_balance + deposit.data.amount
-    state.validators[validator_index].effective_balance ==
-      round_multiple_down(
-        min(MAX_EFFECTIVE_BALANCE, state.balances[validator_index]),
-        EFFECTIVE_BALANCE_INCREMENT
-      )
-
-proc getTestStates(initialState: HashedBeaconState):
-    seq[ref HashedBeaconState] =
-  # Randomly generated slot numbers, with a jump to around
-  # SLOTS_PER_HISTORICAL_ROOT to force wraparound of those
-  # slot-based mod/increment fields.
-  const stateEpochs = [
-    0, 1,
-
-    # Around minimal wraparound SLOTS_PER_HISTORICAL_ROOT wraparound
-    5, 6, 7, 8, 9,
-
-    39, 40, 97, 98, 99, 113, 114, 115, 116, 130, 131, 145, 146, 192, 193,
-    232, 233, 237, 238,
-
-    # Approaching and passing SLOTS_PER_HISTORICAL_ROOT wraparound
-    254, 255, 256, 257, 258]
-
-  var
-    tmpState = assignClone(initialState)
-    cache = StateCache()
-
-  for i, epoch in stateEpochs:
-    let slot = epoch.Epoch.compute_start_slot_at_epoch
-    if tmpState.data.slot < slot:
-      doAssert process_slots(tmpState[], slot, cache)
-    if i mod 3 == 0:
-      valid_deposit(tmpState.data)
-    doAssert tmpState.data.slot == slot
-    result.add assignClone(tmpState[])
-
 template wrappedTimedTest(name: string, body: untyped) =
   # `check` macro takes a copy of whatever it's checking, on the stack!
   # This leads to stack overflow

tests/teststateutil.nim

@@ -0,0 +1,71 @@
+# Nimbus
+# Copyright (c) 2021 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or https://www.apache.org/licenses/LICENSE-2.0)
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT) or https://opensource.org/licenses/MIT)
+# at your option. This file may not be copied, modified, or distributed except according to those terms.
+
+{.push raises: [Defect].}
+
+import
+  options, stew/endians2,
+  ./mocking/mock_deposits,
+  ./helpers/math_helpers,
+  ../beacon_chain/ssz/merkleization,
+  ../beacon_chain/spec/[beaconstate, crypto, datatypes, presets,
+                        helpers, state_transition]
+
+proc valid_deposit(state: var BeaconState) =
+  const deposit_amount = MAX_EFFECTIVE_BALANCE
+  let validator_index = state.validators.len
+  let deposit = mockUpdateStateForNewDeposit(
+                  state,
+                  uint64 validator_index,
+                  deposit_amount,
+                  flags = {}
+                )
+
+  let pre_val_count = state.validators.len
+  let pre_balance = if validator_index < pre_val_count:
+                      state.balances[validator_index]
+                    else:
+                      0
+  doAssert process_deposit(defaultRuntimePreset(), state, deposit, {}).isOk
+  doAssert state.validators.len == pre_val_count + 1
+  doAssert state.balances.len == pre_val_count + 1
+  doAssert state.balances[validator_index] == pre_balance + deposit.data.amount
+  doAssert state.validators[validator_index].effective_balance ==
+    round_multiple_down(
+      min(MAX_EFFECTIVE_BALANCE, state.balances[validator_index]),
+      EFFECTIVE_BALANCE_INCREMENT
+    )
+
+proc getTestStates*(initialState: HashedBeaconState):
+    seq[ref HashedBeaconState] =
+  # Randomly generated slot numbers, with a jump to around
+  # SLOTS_PER_HISTORICAL_ROOT to force wraparound of those
+  # slot-based mod/increment fields.
+  const stateEpochs = [
+    0, 1,
+
+    # Around minimal wraparound SLOTS_PER_HISTORICAL_ROOT wraparound
+    5, 6, 7, 8, 9,
+
+    39, 40, 97, 98, 99, 113, 114, 115, 116, 130, 131, 145, 146, 192, 193,
+    232, 233, 237, 238,
+
+    # Approaching and passing SLOTS_PER_HISTORICAL_ROOT wraparound
+    254, 255, 256, 257, 258]
+
+  var
+    tmpState = assignClone(initialState)
+    cache = StateCache()
+
+  for i, epoch in stateEpochs:
+    let slot = epoch.Epoch.compute_start_slot_at_epoch
+    if tmpState.data.slot < slot:
+      doAssert process_slots(tmpState[], slot, cache)
+    if i mod 3 == 0:
+      valid_deposit(tmpState.data)
+    doAssert tmpState.data.slot == slot
+    result.add assignClone(tmpState[])

tests/testutil.nim

@@ -91,7 +91,7 @@ template timedTest*(name, body) =
   # TODO noto thread-safe as-is
   testTimes.add (f, name)
 
-proc makeTestDB*(tailState: BeaconState, tailBlock: SignedBeaconBlock): BeaconChainDB =
+proc makeTestDB*(tailState: var BeaconState, tailBlock: SignedBeaconBlock): BeaconChainDB =
   result = BeaconChainDB.init(defaultRuntimePreset, "", inMemory = true)
   ChainDAGRef.preInit(result, tailState, tailState, tailBlock)