From 09671fdc0f61c5773eb3de42f893761a450e43d2 Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Mon, 27 Apr 2020 12:18:00 +0200
Subject: [PATCH] avoid memory allocations and copies when loading states

* rolls back some of the ref changes
* adds utility to calculate stack sizes
---
 Makefile                           |   3 +-
 beacon_chain/attestation_pool.nim  |   8 +-
 beacon_chain/beacon_chain_db.nim   |  43 ++++++++--
 beacon_chain/beacon_node.nim       |  50 ++++++------
 beacon_chain/beacon_node_types.nim |  14 ----
 beacon_chain/block_pool.nim        | 127 +++++++++++++++++------------
 beacon_chain/spec/datatypes.nim    |   7 +-
 beacon_chain/state_transition.nim  |  16 ++--
 ncli/ncli_hash_tree_root.nim       |  21 ++---
 ncli/ncli_pretty.nim               |  21 ++---
 ncli/ncli_transition.nim           |   9 +-
 nfuzz/libnfuzz.nim                 |  32 ++++----
 research/stackSizes.nim            |  24 ++++++
 tests/test_attestation_pool.nim    |  58 ++++++-------
 tests/test_beacon_chain_db.nim     |  23 +++---
 tests/test_block_pool.nim          |  22 ++---
 tests/testutil.nim                 |   4 +-
 17 files changed, 274 insertions(+), 208 deletions(-)
 create mode 100644 research/stackSizes.nim

diff --git a/Makefile b/Makefile
index 50f3132feb..7f20d0e216 100644
--- a/Makefile
+++ b/Makefile
@@ -22,7 +22,8 @@ TOOLS := \
 	ncli_hash_tree_root \
 	ncli_pretty \
 	ncli_transition \
-	process_dashboard
+	process_dashboard \
+	stackSizes
 	# bench_bls_sig_agggregation TODO reenable after bls v0.10.1 changes
 TOOLS_DIRS := \
 	beacon_chain \
diff --git a/beacon_chain/attestation_pool.nim b/beacon_chain/attestation_pool.nim
index 858e9a0a5c..38ab31457e 100644
--- a/beacon_chain/attestation_pool.nim
+++ b/beacon_chain/attestation_pool.nim
@@ -180,7 +180,7 @@ proc addResolved(pool: var AttestationPool, blck: BlockRef, attestation: Attesta
     pool.blockPool, pool.blockPool.tmpState,
     BlockSlot(blck: blck, slot: attestation.data.slot))
 
-  template state(): BeaconState = pool.blockPool.tmpState.data.data[]
+  template state(): BeaconState = pool.blockPool.tmpState.data.data
 
   if not validate(state, attestation):
     notice "Invalid attestation",
@@ -456,7 +456,7 @@ proc selectHead*(pool: AttestationPool): BlockRef =
     justifiedHead = pool.blockPool.latestJustifiedBlock()
 
   let newHead =
-    lmdGhost(pool, pool.blockPool.justifiedState.data.data[], justifiedHead.blck)
+    lmdGhost(pool, pool.blockPool.justifiedState.data.data, justifiedHead.blck)
 
   newHead
 
@@ -529,9 +529,9 @@ proc isValidAttestation*(
   # as it supports aggregated attestations (which this can't be)
   var cache = get_empty_per_epoch_cache()
   if not is_valid_indexed_attestation(
-      pool.blockPool.headState.data.data[],
+      pool.blockPool.headState.data.data,
       get_indexed_attestation(
-        pool.blockPool.headState.data.data[], attestation, cache), {}):
+        pool.blockPool.headState.data.data, attestation, cache), {}):
     debug "isValidAttestation: signature verification failed"
     return false
 
diff --git a/beacon_chain/beacon_chain_db.nim b/beacon_chain/beacon_chain_db.nim
index 0e300e5146..c9ad13c90d 100644
--- a/beacon_chain/beacon_chain_db.nim
+++ b/beacon_chain/beacon_chain_db.nim
@@ -1,10 +1,12 @@
 {.push raises: [Defect].}
 
 import
-  options, typetraits, stew/[results, endians2],
+  typetraits, stew/[results, endians2],
   serialization, chronicles,
   spec/[datatypes, digest, crypto],
-  kvstore, ssz
+  kvstore, ssz, state_transition
+
+export results
 
 type
   BeaconChainDB* = ref object
@@ -45,7 +47,7 @@ func subkey[N: static int](kind: DbKeyKind, key: array[N, byte]):
   result[0] = byte ord(kind)
   result[1 .. ^1] = key
 
-func subkey(kind: type BeaconStateRef, key: Eth2Digest): auto =
+func subkey(kind: type BeaconState, key: Eth2Digest): auto =
   subkey(kHashToState, key.data)
 
 func subkey(kind: type SignedBeaconBlock, key: Eth2Digest): auto =
@@ -86,13 +88,13 @@ proc get(db: BeaconChainDB, key: openArray[byte], T: typedesc): Opt[T] =
 proc putBlock*(db: BeaconChainDB, key: Eth2Digest, value: SignedBeaconBlock) =
   db.put(subkey(type value, key), value)
 
-proc putState*(db: BeaconChainDB, key: Eth2Digest, value: BeaconStateRef) =
+proc putState*(db: BeaconChainDB, key: Eth2Digest, value: 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.
 
   db.put(subkey(type value, key), value)
 
-proc putState*(db: BeaconChainDB, value: BeaconStateRef) =
+proc putState*(db: BeaconChainDB, value: BeaconState) =
   db.putState(hash_tree_root(value), value)
 
 proc putStateRoot*(db: BeaconChainDB, root: Eth2Digest, slot: Slot,
@@ -108,7 +110,7 @@ proc delBlock*(db: BeaconChainDB, key: Eth2Digest) =
     "working database")
 
 proc delState*(db: BeaconChainDB, key: Eth2Digest) =
-  db.backend.del(subkey(BeaconStateRef, key)).expect("working database")
+  db.backend.del(subkey(BeaconState, key)).expect("working database")
 
 proc delStateRoot*(db: BeaconChainDB, root: Eth2Digest, slot: Slot) =
   db.backend.del(subkey(root, slot)).expect("working database")
@@ -122,8 +124,31 @@ proc putTailBlock*(db: BeaconChainDB, key: Eth2Digest) =
 proc getBlock*(db: BeaconChainDB, key: Eth2Digest): Opt[SignedBeaconBlock] =
   db.get(subkey(SignedBeaconBlock, key), SignedBeaconBlock)
 
-proc getState*(db: BeaconChainDB, key: Eth2Digest): Opt[BeaconStateRef] =
-  db.get(subkey(BeaconStateRef, key), BeaconStateRef)
+proc getState*(
+    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.
+  # 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
+  # TODO address is needed because there's no way to express lifetimes in nim
+  #      we'll use unsafeAddr to find the code later
+  let outputAddr = unsafeAddr output # callback is local
+  proc decode(data: openArray[byte]) =
+    try:
+      # TODO can't write to output directly..
+      outputAddr[] = SSZ.decode(data, BeaconState)
+    except SerializationError as e:
+      # If the data can't be deserialized, it could be because it's from a
+      # version of the software that uses a different SSZ encoding
+      warn "Unable to deserialize data, old database?", err = e.msg
+      rollback(outputAddr[])
+
+  db.backend.get(subkey(BeaconState, key), decode).expect("working database")
 
 proc getStateRoot*(db: BeaconChainDB,
                    root: Eth2Digest,
@@ -140,7 +165,7 @@ proc containsBlock*(db: BeaconChainDB, key: Eth2Digest): bool =
   db.backend.contains(subkey(SignedBeaconBlock, key)).expect("working database")
 
 proc containsState*(db: BeaconChainDB, key: Eth2Digest): bool =
-  db.backend.contains(subkey(BeaconStateRef, key)).expect("working database")
+  db.backend.contains(subkey(BeaconState, key)).expect("working database")
 
 iterator getAncestors*(db: BeaconChainDB, root: Eth2Digest):
     tuple[root: Eth2Digest, blck: SignedBeaconBlock] =
diff --git a/beacon_chain/beacon_node.nim b/beacon_chain/beacon_node.nim
index 3fdb0d15c6..aa27c6efab 100644
--- a/beacon_chain/beacon_node.nim
+++ b/beacon_chain/beacon_node.nim
@@ -23,7 +23,7 @@ import
   attestation_pool, block_pool, eth2_network, eth2_discovery,
   beacon_node_types, mainchain_monitor, version, ssz, ssz/dynamic_navigator,
   sync_protocol, request_manager, validator_keygen, interop, statusbar,
-  attestation_aggregation, sync_manager
+  attestation_aggregation, sync_manager, state_transition
 
 const
   genesisFile = "genesis.ssz"
@@ -122,11 +122,15 @@ proc getStateFromSnapshot(conf: BeaconNodeConf): NilableBeaconStateRef =
       quit 1
 
   try:
-    result = SSZ.decode(snapshotContents, BeaconStateRef)
+    let res = BeaconStateRef()
+    res[] = SSZ.decode(snapshotContents, BeaconState)
+    result = res
   except SerializationError:
     error "Failed to import genesis file", path = genesisPath
     quit 1
 
+  echo sizeof(result[])
+
   info "Loaded genesis state", path = genesisPath
 
   if writeGenesisFile:
@@ -192,7 +196,7 @@ proc init*(T: type BeaconNode, conf: BeaconNodeConf): Future[BeaconNode] {.async
       let tailBlock = get_initial_beacon_block(genesisState[])
 
       try:
-        BlockPool.preInit(db, genesisState, tailBlock)
+        BlockPool.preInit(db, genesisState[], tailBlock)
         doAssert BlockPool.isInitialized(db), "preInit should have initialized db"
       except CatchableError as e:
         error "Failed to initialize database", err = e.msg
@@ -219,7 +223,7 @@ proc init*(T: type BeaconNode, conf: BeaconNodeConf): Future[BeaconNode] {.async
     nil
 
   let
-    enrForkId = enrForkIdFromState(blockPool.headState.data.data[])
+    enrForkId = enrForkIdFromState(blockPool.headState.data.data)
     topicBeaconBlocks = getBeaconBlocksTopic(enrForkId.forkDigest)
     topicAggregateAndProofs = getAggregateAndProofsTopic(enrForkId.forkDigest)
     network = await createEth2Node(conf, enrForkId)
@@ -235,7 +239,7 @@ proc init*(T: type BeaconNode, conf: BeaconNodeConf): Future[BeaconNode] {.async
     blockPool: blockPool,
     attestationPool: AttestationPool.init(blockPool),
     mainchainMonitor: mainchainMonitor,
-    beaconClock: BeaconClock.init(blockPool.headState.data.data[]),
+    beaconClock: BeaconClock.init(blockPool.headState.data.data),
     rpcServer: rpcServer,
     forkDigest: enrForkId.forkDigest,
     topicBeaconBlocks: topicBeaconBlocks,
@@ -410,15 +414,15 @@ proc proposeBlock(node: BeaconNode,
         (get_eth1data_stub(state.eth1_deposit_index, slot.compute_epoch_at_slot()),
          newSeq[Deposit]())
       else:
-        node.mainchainMonitor.getBlockProposalData(state[])
+        node.mainchainMonitor.getBlockProposalData(state)
 
     let message = makeBeaconBlock(
-      state[],
+      state,
       head.root,
       validator.genRandaoReveal(state.fork, state.genesis_validators_root, slot),
       eth1data,
       Eth2Digest(),
-      node.attestationPool.getAttestationsForBlock(state[]),
+      node.attestationPool.getAttestationsForBlock(state),
       deposits)
 
     if not message.isSome():
@@ -576,16 +580,16 @@ proc handleAttestations(node: BeaconNode, head: BlockRef, slot: Slot) =
   #      version here that calculates the committee for a single slot only
   node.blockPool.withState(node.blockPool.tmpState, attestationHead):
     var cache = get_empty_per_epoch_cache()
-    let committees_per_slot = get_committee_count_at_slot(state[], slot)
+    let committees_per_slot = get_committee_count_at_slot(state, slot)
 
     for committee_index in 0'u64..<committees_per_slot:
       let committee = get_beacon_committee(
-        state[], slot, committee_index.CommitteeIndex, cache)
+        state, slot, committee_index.CommitteeIndex, cache)
 
       for index_in_committee, validatorIdx in committee:
-        let validator = node.getAttachedValidator(state[], validatorIdx)
+        let validator = node.getAttachedValidator(state, validatorIdx)
         if validator != nil:
-          let ad = makeAttestationData(state[], slot, committee_index, blck.root)
+          let ad = makeAttestationData(state, slot, committee_index, blck.root)
           attestations.add((ad, committee.len, index_in_committee, validator))
 
     for a in attestations:
@@ -649,21 +653,21 @@ proc broadcastAggregatedAttestations(
   let bs = BlockSlot(blck: aggregationHead, slot: aggregationSlot)
   node.blockPool.withState(node.blockPool.tmpState, bs):
     let
-      committees_per_slot = get_committee_count_at_slot(state[], aggregationSlot)
+      committees_per_slot = get_committee_count_at_slot(state, aggregationSlot)
     var cache = get_empty_per_epoch_cache()
     for committee_index in 0'u64..<committees_per_slot:
       let committee = get_beacon_committee(
-        state[], aggregationSlot, committee_index.CommitteeIndex, cache)
+        state, aggregationSlot, committee_index.CommitteeIndex, cache)
 
       for index_in_committee, validatorIdx in committee:
-        let validator = node.getAttachedValidator(state[], validatorIdx)
+        let validator = node.getAttachedValidator(state, validatorIdx)
         if validator != nil:
           # This is slightly strange/inverted control flow, since really it's
           # going to happen once per slot, but this is the best way to get at
           # the validator index and private key pair. TODO verify it only has
           # one isSome() with test.
           let aggregateAndProof =
-            aggregate_attestations(node.attestationPool, state[],
+            aggregate_attestations(node.attestationPool, state,
               committee_index.CommitteeIndex,
               # TODO https://github.com/status-im/nim-beacon-chain/issues/545
               # this assumes in-process private keys
@@ -1060,13 +1064,13 @@ proc installBeaconApiHandlers(rpcServer: RpcServer, node: BeaconNode) =
     requireOneOf(slot, root)
     if slot.isSome:
       let blk = node.blockPool.head.blck.atSlot(slot.get)
-      var tmpState = emptyStateData()
-      node.blockPool.withState(tmpState, blk):
+      node.blockPool.withState(node.blockPool.tmpState, blk):
         return jsonResult(state)
     else:
-      let state = node.db.getState(root.get)
-      if state.isSome:
-        return jsonResult(state.get)
+      let tmp = BeaconStateRef() # TODO use tmpState - but load the entire StateData!
+      let state = node.db.getState(root.get, tmp[], noRollback)
+      if state:
+        return jsonResult(tmp[])
       else:
         return StringOfJson("null")
 
@@ -1193,7 +1197,7 @@ proc start(node: BeaconNode) =
     bs = BlockSlot(blck: head.blck, slot: head.blck.slot)
 
   node.blockPool.withState(node.blockPool.tmpState, bs):
-    node.addLocalValidators(state[])
+    node.addLocalValidators(state)
 
   node.run()
 
@@ -1282,7 +1286,7 @@ when hasPrompt:
           # TODO slow linear scan!
           for idx, b in node.blockPool.headState.data.data.balances:
             if node.getAttachedValidator(
-                node.blockPool.headState.data.data[], ValidatorIndex(idx)) != nil:
+                node.blockPool.headState.data.data, ValidatorIndex(idx)) != nil:
               balance += b
           formatGwei(balance)
 
diff --git a/beacon_chain/beacon_node_types.nim b/beacon_chain/beacon_node_types.nim
index 4382f13ac2..9b94f6e147 100644
--- a/beacon_chain/beacon_node_types.nim
+++ b/beacon_chain/beacon_node_types.nim
@@ -225,20 +225,6 @@ type
     root*: Eth2Digest
     historySlots*: uint64
 
-func emptyStateData*: StateData =
-  StateData(
-    data: HashedBeaconState(
-      # Please note that this initialization is needed in order
-      # to allocate memory for the BeaconState:
-      data: BeaconStateRef(),
-      root: default(Eth2Digest)
-    ),
-    blck: default(BlockRef))
-
-func clone*(other: StateData): StateData =
-  StateData(data: clone(other.data),
-            blck: other.blck)
-
 proc shortLog*(v: AttachedValidator): string = shortLog(v.pubKey)
 
 chronicles.formatIt BlockSlot:
diff --git a/beacon_chain/block_pool.nim b/beacon_chain/block_pool.nim
index dfee9cd4ed..5bace1bb94 100644
--- a/beacon_chain/block_pool.nim
+++ b/beacon_chain/block_pool.nim
@@ -33,7 +33,7 @@ template withState*(
   updateStateData(pool, cache, blockSlot)
 
   template hashedState(): HashedBeaconState {.inject, used.} = cache.data
-  template state(): BeaconStateRef {.inject, used.} = cache.data.data
+  template state(): BeaconState {.inject, used.} = cache.data.data
   template blck(): BlockRef {.inject, used.} = cache.blck
   template root(): Eth2Digest {.inject, used.} = cache.data.root
 
@@ -208,9 +208,10 @@ proc init*(T: type BlockPool, db: BeaconChainDB): BlockPool =
   # we're loading the head block, the corresponding state does not necessarily
   # exist in the database - we'll load this latest state we know about and use
   # that as finalization point.
-  let stateOpt = db.getState(latestStateRoot.get().stateRoot)
-  doAssert stateOpt.isSome, "failed to obtain latest state. database corrupt?"
-  let tmpState = stateOpt.get
+  let tmpState = BeaconStateRef()
+  let stateOpt = db.getState(
+    latestStateRoot.get().stateRoot, tmpState[], noRollback)
+  doAssert stateOpt, "failed to obtain latest state. database corrupt?"
 
   let
     finalizedSlot =
@@ -232,12 +233,13 @@ proc init*(T: type BlockPool, db: BeaconChainDB): BlockPool =
 
   let headState = StateData(
     data: HashedBeaconState(
-      data: tmpState, root: latestStateRoot.get().stateRoot),
+      data: tmpState[], root: latestStateRoot.get().stateRoot),
     blck: latestStateRoot.get().blckRef)
 
-  let justifiedState = db.getState(justifiedStateRoot)
-  doAssert justifiedState.isSome,
-           "failed to obtain latest justified state. database corrupt?"
+  let justifiedState = BeaconStateRef() # TODO avoid scratch space
+  let found = db.getState(justifiedStateRoot, justifiedState[], noRollback)
+  # TODO turn into regular error, this can happen
+  doAssert found, "failed to obtain latest justified state. database corrupt?"
 
   # For the initialization of `tmpState` below.
   # Please note that it's initialized few lines below
@@ -262,7 +264,7 @@ proc init*(T: type BlockPool, db: BeaconChainDB): BlockPool =
     heads: @[head],
     headState: headState,
     justifiedState: StateData(
-      data: HashedBeaconState(data: justifiedState.get, root: justifiedStateRoot),
+      data: HashedBeaconState(data: justifiedState[], root: justifiedStateRoot),
       blck: justifiedHead.blck),
     tmpState: default(StateData)
   )
@@ -270,7 +272,7 @@ proc init*(T: type BlockPool, db: BeaconChainDB): BlockPool =
 
   res.updateStateData(res.headState, BlockSlot(blck: head.blck,
                                                slot: head.blck.slot))
-  res.tmpState = clone(res.headState)
+  res.tmpState = res.headState
   res
 
 proc addResolvedBlock(
@@ -419,7 +421,7 @@ proc add*(
     # Careful, tmpState.data has been updated but not blck - we need to create
     # the BlockRef first!
     pool.tmpState.blck = pool.addResolvedBlock(
-      pool.tmpState.data.data[], blockRoot, signedBlock, parent)
+      pool.tmpState.data.data, blockRoot, signedBlock, parent)
 
     return pool.tmpState.blck
 
@@ -571,6 +573,36 @@ func checkMissing*(pool: var BlockPool): seq[FetchRecord] =
     if v.tries.popcount() == 1:
       result.add(FetchRecord(root: k, historySlots: v.slots))
 
+proc getState(
+    pool: BlockPool, db: BeaconChainDB, stateRoot: Eth2Digest, blck: BlockRef,
+    output: var StateData): bool =
+  let outputAddr = unsafeAddr output # local scope
+  proc rollback(v: var BeaconState) =
+    if outputAddr == (unsafeAddr pool.headState):
+      # TODO seeing the headState in the rollback shouldn't happen - we load
+      #      head states only when updating the head position, and by that time
+      #      the database will have gone through enough sanity checks that
+      #      SSZ exceptions shouldn't happen, which is when rollback happens.
+      #      Nonetheless, this is an ugly workaround that needs to go away
+      doAssert false, "Cannot alias headState"
+
+    outputAddr[] = pool.headState
+
+  if not db.getState(stateRoot, output.data.data, rollback):
+    return false
+
+  output.blck = blck
+  output.data.root = stateRoot
+
+  true
+
+proc getState(pool: BlockPool, db: BeaconChainDB, bs: BlockSlot, output: var StateData): bool =
+  let stateRoot = db.getStateRoot(bs.blck.root, bs.slot)
+  if not stateRoot.isSome:
+    return false
+
+  pool.getState(db, stateRoot.get(), bs.blck, output)
+
 proc putState(pool: BlockPool, state: HashedBeaconState, blck: BlockRef) =
   # TODO we save state at every epoch start but never remove them - we also
   #      potentially save multiple states per slot if reorgs happen, meaning
@@ -640,12 +672,12 @@ proc skipAndUpdateState(
 
   pool.skipAndUpdateState(state.data, blck.refs, blck.data.message.slot - 1)
 
-  var statePtr = unsafeAddr state # safe because `restore` is locally scoped
-  proc restore(v: var HashedBeaconState) =
+  var statePtr = unsafeAddr state # safe because `rollback` is locally scoped
+  proc rollback(v: var HashedBeaconState) =
     doAssert (addr(statePtr.data) == addr v)
     statePtr[] = pool.headState
 
-  let ok  = state_transition(state.data, blck.data, flags, restore)
+  let ok  = state_transition(state.data, blck.data, flags, rollback)
   if ok:
     pool.putState(state.data, blck.refs)
 
@@ -711,39 +743,35 @@ proc rewindState(pool: BlockPool, state: var StateData, bs: BlockSlot):
   let
     ancestor = ancestors.pop()
     root = stateRoot.get()
-    ancestorState =
-      if pool.db.containsState(root):
-        pool.db.getState(root)
-      elif pool.cachedStates[0].containsState(root):
-        pool.cachedStates[0].getState(root)
-      else:
-        pool.cachedStates[1].getState(root)
 
-  if ancestorState.isNone():
-    # TODO this should only happen if the database is corrupt - we walked the
-    #      list of parent blocks and couldn't find a corresponding state in the
-    #      database, which should never happen (at least we should have the
-    #      tail state in there!)
-    error "Couldn't find ancestor state or block parent missing!",
-      blockRoot = shortLog(bs.blck.root),
-      blockSlot = shortLog(bs.blck.slot),
-      slot = shortLog(bs.slot),
-      cat = "crash"
-    doAssert false, "Oh noes, we passed big bang!"
+  if pool.cachedStates[0].containsState(root):
+    doAssert pool.getState(pool.cachedStates[0], root, ancestor.refs, state)
+  elif pool.cachedStates[1].containsState(root):
+    doAssert pool.getState(pool.cachedStates[1], root, ancestor.refs, state)
+  else:
+    let found = pool.getState(pool.db, root, ancestor.refs, state)
+
+    if not found:
+      # TODO this should only happen if the database is corrupt - we walked the
+      #      list of parent blocks and couldn't find a corresponding state in the
+      #      database, which should never happen (at least we should have the
+      #      tail state in there!)
+      error "Couldn't find ancestor state or block parent missing!",
+        blockRoot = shortLog(bs.blck.root),
+        blockSlot = shortLog(bs.blck.slot),
+        slot = shortLog(bs.slot),
+        cat = "crash"
+      doAssert false, "Oh noes, we passed big bang!"
 
   trace "Replaying state transitions",
     stateSlot = shortLog(state.data.data.slot),
     ancestorStateRoot = shortLog(ancestor.data.message.state_root),
-    ancestorStateSlot = shortLog(ancestorState.get().slot),
+    ancestorStateSlot = shortLog(state.data.data.slot),
     slot = shortLog(bs.slot),
     blockRoot = shortLog(bs.blck.root),
     ancestors = ancestors.len,
     cat = "replay_state"
 
-  state.data.data[] = ancestorState.get()[]
-  state.data.root = stateRoot.get()
-  state.blck = ancestor.refs
-
   ancestors
 
 proc getStateDataCached(pool: BlockPool, state: var StateData, bs: BlockSlot): bool =
@@ -758,12 +786,10 @@ proc getStateDataCached(pool: BlockPool, state: var StateData, bs: BlockSlot): b
 
       let
         root = tmp.get()
-        ancestorState = db.getState(root)
 
-      doAssert ancestorState.isSome()
-      state.data.data = ancestorState.get()
-      state.data.root = root
-      state.blck = pool.get(bs.blck).refs
+      let found = pool.getState(db, root, bs.blck, state)
+      doAssert found, "we just checked for it above and it shouldn't be corrupt"
+
       return true
 
   false
@@ -811,12 +837,9 @@ proc updateStateData*(pool: BlockPool, state: var StateData, bs: BlockSlot) =
 proc loadTailState*(pool: BlockPool): StateData =
   ## Load the state associated with the current tail in the pool
   let stateRoot = pool.db.getBlock(pool.tail.root).get().message.state_root
-  StateData(
-    data: HashedBeaconState(
-      data: pool.db.getState(stateRoot).get(),
-      root: stateRoot),
-    blck: pool.tail
-  )
+  let found = pool.getState(pool.db, stateRoot, pool.tail, result)
+  # TODO turn into regular error, this can happen
+  doAssert found, "Failed to load tail state, database corrupt?"
 
 proc delState(pool: BlockPool, bs: BlockSlot) =
   # Delete state state and mapping for a particular block+slot
@@ -992,7 +1015,7 @@ proc isInitialized*(T: type BlockPool, db: BeaconChainDB): bool =
   return true
 
 proc preInit*(
-    T: type BlockPool, db: BeaconChainDB, state: BeaconStateRef,
+    T: type BlockPool, db: BeaconChainDB, state: BeaconState,
     signedBlock: SignedBeaconBlock) =
   # write a genesis state, the way the BlockPool expects it to be stored in
   # database
@@ -1021,14 +1044,14 @@ proc getProposer*(pool: BlockPool, head: BlockRef, slot: Slot): Option[Validator
     var cache = get_empty_per_epoch_cache()
 
     # https://github.com/ethereum/eth2.0-specs/blob/v0.11.1/specs/phase0/validator.md#validator-assignments
-    let proposerIdx = get_beacon_proposer_index(state[], cache)
+    let proposerIdx = get_beacon_proposer_index(state, cache)
     if proposerIdx.isNone:
       warn "Missing proposer index",
         slot=slot,
         epoch=slot.compute_epoch_at_slot,
         num_validators=state.validators.len,
         active_validators=
-          get_active_validator_indices(state[], slot.compute_epoch_at_slot),
+          get_active_validator_indices(state, slot.compute_epoch_at_slot),
         balances=state.balances
       return
 
@@ -1130,7 +1153,7 @@ proc isValidBeaconBlock*(pool: var BlockPool,
   pool.withState(pool.tmpState, bs):
     let
       blockRoot = hash_tree_root(signed_beacon_block.message)
-      domain = get_domain(pool.headState.data.data[], DOMAIN_BEACON_PROPOSER,
+      domain = get_domain(pool.headState.data.data, DOMAIN_BEACON_PROPOSER,
         compute_epoch_at_slot(signed_beacon_block.message.slot))
       signing_root = compute_signing_root(blockRoot, domain)
       proposer_index = signed_beacon_block.message.proposer_index
diff --git a/beacon_chain/spec/datatypes.nim b/beacon_chain/spec/datatypes.nim
index 6b372edc14..c8113aaf7c 100644
--- a/beacon_chain/spec/datatypes.nim
+++ b/beacon_chain/spec/datatypes.nim
@@ -391,7 +391,7 @@ type
 
   # TODO to be replaced with some magic hash caching
   HashedBeaconState* = object
-    data*: BeaconStateRef
+    data*: BeaconState
     root*: Eth2Digest # hash_tree_root(data)
 
   StateCache* = object
@@ -579,11 +579,6 @@ func newClone*[T](x: ref T): ref T not nil =
   new result
   result[] = x[]
 
-func clone*(other: HashedBeaconState): HashedBeaconState =
-  HashedBeaconState(
-    data: newClone(other.data),
-    root: other.root)
-
 template init*(T: type BitList, len: int): auto = T init(BitSeq, len)
 template len*(x: BitList): auto = len(BitSeq(x))
 template bytes*(x: BitList): auto = bytes(BitSeq(x))
diff --git a/beacon_chain/state_transition.nim b/beacon_chain/state_transition.nim
index 08dfac061b..d66e50e3df 100644
--- a/beacon_chain/state_transition.nim
+++ b/beacon_chain/state_transition.nim
@@ -139,7 +139,7 @@ proc verifyStateRoot(state: BeaconState, blck: BeaconBlock): bool =
     true
 
 type
-  RollbackProc* = proc(v: var BeaconState) {.gcsafe.}
+  RollbackProc* = proc(v: var BeaconState) {.gcsafe, raises: [Defect].}
 
 proc noRollback*(state: var BeaconState) =
   trace "Skipping rollback of broken state"
@@ -247,14 +247,14 @@ proc process_slots*(state: var HashedBeaconState, slot: Slot) =
     if is_epoch_transition:
       # Note: Genesis epoch = 0, no need to test if before Genesis
       try:
-        beacon_previous_validators.set(get_epoch_validator_count(state.data[]))
+        beacon_previous_validators.set(get_epoch_validator_count(state.data))
       except Exception as e: # TODO https://github.com/status-im/nim-metrics/pull/22
         trace "Couldn't update metrics", msg = e.msg
-      process_epoch(state.data[])
+      process_epoch(state.data)
     state.data.slot += 1
     if is_epoch_transition:
       try:
-        beacon_current_validators.set(get_epoch_validator_count(state.data[]))
+        beacon_current_validators.set(get_epoch_validator_count(state.data))
       except Exception as e: # TODO https://github.com/status-im/nim-metrics/pull/22
         trace "Couldn't update metrics", msg = e.msg
     state.root = hash_tree_root(state.data)
@@ -272,18 +272,18 @@ proc state_transition*(
   process_slots(state, signedBlock.message.slot)
 
   if skipBLSValidation in flags or
-      verify_block_signature(state.data[], signedBlock):
+      verify_block_signature(state.data, signedBlock):
 
     var per_epoch_cache = get_empty_per_epoch_cache()
-    if processBlock(state.data[], signedBlock.message, flags, per_epoch_cache):
-      if skipStateRootValidation in flags or verifyStateRoot(state.data[], signedBlock.message):
+    if processBlock(state.data, signedBlock.message, flags, per_epoch_cache):
+      if skipStateRootValidation in flags or verifyStateRoot(state.data, signedBlock.message):
         # State root is what it should be - we're done!
 
         # TODO when creating a new block, state_root is not yet set.. comparing
         #      with zero hash here is a bit fragile however, but this whole thing
         #      should go away with proper hash caching
         state.root =
-          if signedBlock.message.state_root == Eth2Digest(): hash_tree_root(state.data[])
+          if signedBlock.message.state_root == Eth2Digest(): hash_tree_root(state.data)
           else: signedBlock.message.state_root
 
         return true
diff --git a/ncli/ncli_hash_tree_root.nim b/ncli/ncli_hash_tree_root.nim
index 424201961e..9818534bc2 100644
--- a/ncli/ncli_hash_tree_root.nim
+++ b/ncli/ncli_hash_tree_root.nim
@@ -8,16 +8,17 @@ import
 cli do(kind: string, file: string):
 
   template printit(t: untyped) {.dirty.} =
-      let v =
-        if cmpIgnoreCase(ext, ".ssz") == 0:
-          SSZ.loadFile(file, t)
-        elif cmpIgnoreCase(ext, ".json") == 0:
-          JSON.loadFile(file, t)
-        else:
-          echo "Unknown file type: ", ext
-          quit 1
+    let v = new t
+    v[] =
+      if cmpIgnoreCase(ext, ".ssz") == 0:
+        SSZ.loadFile(file, t)
+      elif cmpIgnoreCase(ext, ".json") == 0:
+        JSON.loadFile(file, t)
+      else:
+        echo "Unknown file type: ", ext
+        quit 1
 
-      echo hash_tree_root(v).data.toHex(true)
+    echo hash_tree_root(v).data.toHex(true)
 
   let ext = splitFile(file).ext
 
@@ -30,6 +31,6 @@ cli do(kind: string, file: string):
   of "deposit": printit(Deposit)
   of "deposit_data": printit(DepositData)
   of "eth1_data": printit(Eth1Data)
-  of "state": printit(BeaconStateRef)
+  of "state": printit(BeaconState)
   of "proposer_slashing": printit(ProposerSlashing)
   of "voluntary_exit": printit(VoluntaryExit)
diff --git a/ncli/ncli_pretty.nim b/ncli/ncli_pretty.nim
index f4cccd7254..9de9ddf5cb 100644
--- a/ncli/ncli_pretty.nim
+++ b/ncli/ncli_pretty.nim
@@ -7,16 +7,17 @@ import
 cli do(kind: string, file: string):
 
   template printit(t: untyped) {.dirty.} =
-      let v =
-        if cmpIgnoreCase(ext, ".ssz") == 0:
-          SSZ.loadFile(file, t)
-        elif cmpIgnoreCase(ext, ".json") == 0:
-          JSON.loadFile(file, t)
-        else:
-          echo "Unknown file type: ", ext
-          quit 1
+    let v = new t
+    v[] =
+      if cmpIgnoreCase(ext, ".ssz") == 0:
+        SSZ.loadFile(file, t)
+      elif cmpIgnoreCase(ext, ".json") == 0:
+        JSON.loadFile(file, t)
+      else:
+        echo "Unknown file type: ", ext
+        quit 1
 
-      echo JSON.encode(v, pretty = true)
+    echo JSON.encode(v[], pretty = true)
 
   let ext = splitFile(file).ext
 
@@ -29,7 +30,7 @@ cli do(kind: string, file: string):
   of "deposit": printit(Deposit)
   of "deposit_data": printit(DepositData)
   of "eth1_data": printit(Eth1Data)
-  of "state": printit(BeaconStateRef)
+  of "state": printit(BeaconState)
   of "proposer_slashing": printit(ProposerSlashing)
   of "voluntary_exit": printit(VoluntaryExit)
   else: echo "Unknown kind"
diff --git a/ncli/ncli_transition.nim b/ncli/ncli_transition.nim
index 3574c2b884..235d248b32 100644
--- a/ncli/ncli_transition.nim
+++ b/ncli/ncli_transition.nim
@@ -4,13 +4,16 @@ import
   ../beacon_chain/[extras, state_transition, ssz]
 
 cli do(pre: string, blck: string, post: string, verifyStateRoot = false):
+
   let
-    stateX = SSZ.loadFile(pre, BeaconStateRef)
+    stateY = (ref HashedBeaconState)()
     blckX = SSZ.loadFile(blck, SignedBeaconBlock)
     flags = if verifyStateRoot: {skipStateRootValidation} else: {}
 
-  var stateY = HashedBeaconState(data: stateX, root: hash_tree_root(stateX))
-  if not state_transition(stateY, blckX, flags, noRollback):
+  stateY.data = SSZ.loadFile(pre, BeaconState)
+  stateY.root = hash_tree_root(stateY.data)
+
+  if not state_transition(stateY[], blckX, flags, noRollback):
     error "State transition failed"
   else:
     SSZ.saveFile(post, stateY.data)
diff --git a/nfuzz/libnfuzz.nim b/nfuzz/libnfuzz.nim
index 56a13a5a14..1f8c1ae142 100644
--- a/nfuzz/libnfuzz.nim
+++ b/nfuzz/libnfuzz.nim
@@ -10,23 +10,23 @@ import
 
 type
   AttestationInput = object
-    state: BeaconStateRef
+    state: BeaconState
     attestation: Attestation
   AttesterSlashingInput = object
-    state: BeaconStateRef
+    state: BeaconState
     attesterSlashing: AttesterSlashing
   BlockInput = object
-    state: BeaconStateRef
+    state: BeaconState
     beaconBlock: SignedBeaconBlock
   BlockHeaderInput = BlockInput
   DepositInput = object
-    state: BeaconStateRef
+    state: BeaconState
     deposit: Deposit
   ProposerSlashingInput = object
-    state: BeaconStateRef
+    state: BeaconState
     proposerSlashing: ProposerSlashing
   VoluntaryExitInput = object
-    state: BeaconStateRef
+    state: BeaconState
     exit: SignedVoluntaryExit
   # This and AssertionError are raised to indicate programming bugs
   # A wrapper to allow exception tracking to identify unexpected exceptions
@@ -61,7 +61,9 @@ template decodeAndProcess(typ, process: untyped): bool =
 
   var
     cache {.used, inject.} = get_empty_per_epoch_cache()
-    data {.inject.} =
+    data {.inject.} = new typ
+
+  data[] =
       try:
         SSZ.decode(input, typ)
       except MalformedSszError as e:
@@ -89,44 +91,44 @@ template decodeAndProcess(typ, process: untyped): bool =
       raise newException(FuzzCrashError, "Unexpected Exception in state transition", e)
 
   if processOk:
-    copyState(data.state[], output, output_size)
+    copyState(data.state, output, output_size)
   else:
     false
 
 proc nfuzz_attestation(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(AttestationInput):
-    process_attestation(data.state[], data.attestation, flags, cache)
+    process_attestation(data.state, data.attestation, flags, cache)
 
 proc nfuzz_attester_slashing(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(AttesterSlashingInput):
-    process_attester_slashing(data.state[], data.attesterSlashing, flags, cache)
+    process_attester_slashing(data.state, data.attesterSlashing, flags, cache)
 
 proc nfuzz_block(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(BlockInput):
-    state_transition(data.state[], data.beaconBlock, flags, noRollback)
+    state_transition(data.state, data.beaconBlock, flags, noRollback)
 
 proc nfuzz_block_header(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(BlockHeaderInput):
-    process_block_header(data.state[], data.beaconBlock.message, flags, cache)
+    process_block_header(data.state, data.beaconBlock.message, flags, cache)
 
 proc nfuzz_deposit(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(DepositInput):
-    process_deposit(data.state[], data.deposit, flags)
+    process_deposit(data.state, data.deposit, flags)
 
 proc nfuzz_proposer_slashing(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(ProposerSlashingInput):
-    process_proposer_slashing(data.state[], data.proposerSlashing, flags, cache)
+    process_proposer_slashing(data.state, data.proposerSlashing, flags, cache)
 
 proc nfuzz_voluntary_exit(input: openArray[byte], output: ptr byte,
     output_size: ptr uint, disable_bls: bool): bool {.exportc, raises: [FuzzCrashError, Defect].} =
   decodeAndProcess(VoluntaryExitInput):
-    process_voluntary_exit(data.state[], data.exit, flags)
+    process_voluntary_exit(data.state, data.exit, flags)
 
 # Note: Could also accept raw input pointer and access list_size + seed here.
 # However, list_size needs to be known also outside this proc to allocate output.
diff --git a/research/stackSizes.nim b/research/stackSizes.nim
new file mode 100644
index 0000000000..6a8268ef7a
--- /dev/null
+++ b/research/stackSizes.nim
@@ -0,0 +1,24 @@
+import ../beacon_chain/spec/datatypes
+
+import typetraits, strformat, strutils
+
+proc print(t: auto, n: string, indent: int) =
+  echo fmt"{sizeof(t):>8}  {spaces(indent)}{n}: {typeof(t).name}"
+
+  when t is object|tuple:
+    for n, p in t.fieldPairs:
+      print(p, n, indent + 1)
+
+print(BeaconState(), "state", 0)
+
+echo ""
+
+print(SignedBeaconBlock(), "block", 0)
+
+echo ""
+
+print(Validator(), "validator", 0)
+
+echo ""
+
+print(Attestation(), "attestation", 0)
diff --git a/tests/test_attestation_pool.nim b/tests/test_attestation_pool.nim
index aa891f49aa..71e0e80afe 100644
--- a/tests/test_attestation_pool.nim
+++ b/tests/test_attestation_pool.nim
@@ -29,7 +29,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       var
         blockPool = BlockPool.init(makeTestDB(SLOTS_PER_EPOCH * 3))
         pool = AttestationPool.init(blockPool)
-        state = loadTailState(blockPool)
+        state = newClone(loadTailState(blockPool))
       # Slot 0 is a finalized slot - won't be making attestations for it..
       process_slots(state.data, state.data.data.slot + 1)
 
@@ -38,15 +38,15 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       let
         # Create an attestation for slot 1!
         beacon_committee = get_beacon_committee(
-          state.data.data[], state.data.data.slot, 0.CommitteeIndex, cache)
+          state.data.data, state.data.data.slot, 0.CommitteeIndex, cache)
         attestation = makeAttestation(
-          state.data.data[], state.blck.root, beacon_committee[0], cache)
+          state.data.data, state.blck.root, beacon_committee[0], cache)
 
       pool.add(attestation)
 
       process_slots(state.data, MIN_ATTESTATION_INCLUSION_DELAY.Slot + 1)
 
-      let attestations = pool.getAttestationsForBlock(state.data.data[])
+      let attestations = pool.getAttestationsForBlock(state.data.data)
 
       check:
         attestations.len == 1
@@ -56,17 +56,17 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       let
         # Create an attestation for slot 1!
         bc0 = get_beacon_committee(
-          state.data.data[], state.data.data.slot, 0.CommitteeIndex, cache)
+          state.data.data, state.data.data.slot, 0.CommitteeIndex, cache)
         attestation0 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[0], cache)
+          state.data.data, state.blck.root, bc0[0], cache)
 
       process_slots(state.data, state.data.data.slot + 1)
 
       let
-        bc1 = get_beacon_committee(state.data.data[],
+        bc1 = get_beacon_committee(state.data.data,
           state.data.data.slot, 0.CommitteeIndex, cache)
         attestation1 = makeAttestation(
-          state.data.data[], state.blck.root, bc1[0], cache)
+          state.data.data, state.blck.root, bc1[0], cache)
 
       # test reverse order
       pool.add(attestation1)
@@ -74,7 +74,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
 
       process_slots(state.data, MIN_ATTESTATION_INCLUSION_DELAY.Slot + 1)
 
-      let attestations = pool.getAttestationsForBlock(state.data.data[])
+      let attestations = pool.getAttestationsForBlock(state.data.data)
 
       check:
         attestations.len == 1
@@ -84,18 +84,18 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       let
         # Create an attestation for slot 1!
         bc0 = get_beacon_committee(
-          state.data.data[], state.data.data.slot, 0.CommitteeIndex, cache)
+          state.data.data, state.data.data.slot, 0.CommitteeIndex, cache)
         attestation0 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[0], cache)
+          state.data.data, state.blck.root, bc0[0], cache)
         attestation1 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[1], cache)
+          state.data.data, state.blck.root, bc0[1], cache)
 
       pool.add(attestation0)
       pool.add(attestation1)
 
       process_slots(state.data, MIN_ATTESTATION_INCLUSION_DELAY.Slot + 1)
 
-      let attestations = pool.getAttestationsForBlock(state.data.data[])
+      let attestations = pool.getAttestationsForBlock(state.data.data)
 
       check:
         attestations.len == 1
@@ -106,11 +106,11 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       var
         # Create an attestation for slot 1!
         bc0 = get_beacon_committee(
-          state.data.data[], state.data.data.slot, 0.CommitteeIndex, cache)
+          state.data.data, state.data.data.slot, 0.CommitteeIndex, cache)
         attestation0 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[0], cache)
+          state.data.data, state.blck.root, bc0[0], cache)
         attestation1 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[1], cache)
+          state.data.data, state.blck.root, bc0[1], cache)
 
       attestation0.combine(attestation1, {})
 
@@ -119,7 +119,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
 
       process_slots(state.data, MIN_ATTESTATION_INCLUSION_DELAY.Slot + 1)
 
-      let attestations = pool.getAttestationsForBlock(state.data.data[])
+      let attestations = pool.getAttestationsForBlock(state.data.data)
 
       check:
         attestations.len == 1
@@ -128,12 +128,12 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       var cache = get_empty_per_epoch_cache()
       var
         # Create an attestation for slot 1!
-        bc0 = get_beacon_committee(state.data.data[],
+        bc0 = get_beacon_committee(state.data.data,
           state.data.data.slot, 0.CommitteeIndex, cache)
         attestation0 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[0], cache)
+          state.data.data, state.blck.root, bc0[0], cache)
         attestation1 = makeAttestation(
-          state.data.data[], state.blck.root, bc0[1], cache)
+          state.data.data, state.blck.root, bc0[1], cache)
 
       attestation0.combine(attestation1, {})
 
@@ -142,14 +142,14 @@ when const_preset == "minimal": # Too much stack space used on mainnet
 
       process_slots(state.data, MIN_ATTESTATION_INCLUSION_DELAY.Slot + 1)
 
-      let attestations = pool.getAttestationsForBlock(state.data.data[])
+      let attestations = pool.getAttestationsForBlock(state.data.data)
 
       check:
         attestations.len == 1
 
     timedTest "Fork choice returns latest block with no attestations":
       let
-        b1 = addTestBlock(state.data.data[], blockPool.tail.root)
+        b1 = addTestBlock(state.data.data, blockPool.tail.root)
         b1Root = hash_tree_root(b1.message)
         b1Add = blockPool.add(b1Root, b1)
         head = pool.selectHead()
@@ -158,7 +158,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
         head == b1Add
 
       let
-        b2 = addTestBlock(state.data.data[], b1Root)
+        b2 = addTestBlock(state.data.data, b1Root)
         b2Root = hash_tree_root(b2.message)
         b2Add = blockPool.add(b2Root, b2)
         head2 = pool.selectHead()
@@ -169,7 +169,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
     timedTest "Fork choice returns block with attestation":
       var cache = get_empty_per_epoch_cache()
       let
-        b10 = makeTestBlock(state.data.data[], blockPool.tail.root)
+        b10 = makeTestBlock(state.data.data, blockPool.tail.root)
         b10Root = hash_tree_root(b10.message)
         b10Add = blockPool.add(b10Root, b10)
         head = pool.selectHead()
@@ -178,15 +178,15 @@ when const_preset == "minimal": # Too much stack space used on mainnet
         head == b10Add
 
       let
-        b11 = makeTestBlock(state.data.data[], blockPool.tail.root,
+        b11 = makeTestBlock(state.data.data, blockPool.tail.root,
           graffiti = Eth2Digest(data: [1'u8, 0, 0, 0 ,0 ,0 ,0 ,0 ,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
         )
         b11Root = hash_tree_root(b11.message)
         b11Add = blockPool.add(b11Root, b11)
 
         bc1 = get_beacon_committee(
-          state.data.data[], state.data.data.slot, 1.CommitteeIndex, cache)
-        attestation0 = makeAttestation(state.data.data[], b10Root, bc1[0], cache)
+          state.data.data, state.data.data.slot, 1.CommitteeIndex, cache)
+        attestation0 = makeAttestation(state.data.data, b10Root, bc1[0], cache)
 
       pool.add(attestation0)
 
@@ -197,8 +197,8 @@ when const_preset == "minimal": # Too much stack space used on mainnet
         head2 == b10Add
 
       let
-        attestation1 = makeAttestation(state.data.data[], b11Root, bc1[1], cache)
-        attestation2 = makeAttestation(state.data.data[], b11Root, bc1[2], cache)
+        attestation1 = makeAttestation(state.data.data, b11Root, bc1[1], cache)
+        attestation2 = makeAttestation(state.data.data, b11Root, bc1[2], cache)
       pool.add(attestation1)
 
       let head3 = pool.selectHead()
diff --git a/tests/test_beacon_chain_db.nim b/tests/test_beacon_chain_db.nim
index c9c715944c..7e02a4a4e4 100644
--- a/tests/test_beacon_chain_db.nim
+++ b/tests/test_beacon_chain_db.nim
@@ -8,7 +8,7 @@
 {.used.}
 
 import  options, unittest, sequtils,
-  ../beacon_chain/[beacon_chain_db, extras, interop, ssz, kvstore],
+  ../beacon_chain/[beacon_chain_db, extras, interop, ssz, kvstore, state_transition],
   ../beacon_chain/spec/[beaconstate, datatypes, digest, crypto],
   # test utilies
   ./testutil, ./testblockutil
@@ -17,13 +17,10 @@ suiteReport "Beacon chain DB" & preset():
   timedTest "empty database" & preset():
     var
       db = init(BeaconChainDB, kvStore MemStoreRef.init())
-
+      tmpState = BeaconStateRef()
     check:
-      when const_preset=="minimal":
-        db.getState(Eth2Digest()).isNone and db.getBlock(Eth2Digest()).isNone
-      else:
-        # TODO re-check crash here in mainnet
-        true
+      not db.getState(Eth2Digest(), tmpState[], noRollback)
+      db.getBlock(Eth2Digest()).isNone
 
   timedTest "sanity check blocks" & preset():
     var
@@ -49,13 +46,15 @@ suiteReport "Beacon chain DB" & preset():
 
     let
       state = BeaconStateRef()
+      state2 = BeaconStateRef()
       root = hash_tree_root(state)
 
-    db.putState(state)
+    db.putState(state[])
 
     check:
       db.containsState(root)
-      db.getState(root).get[] == state[]
+      db.getState(root, state2[], noRollback)
+      state2[] == state[]
 
   timedTest "find ancestors" & preset():
     var
@@ -102,8 +101,10 @@ suiteReport "Beacon chain DB" & preset():
         eth1BlockHash, 0, makeInitialDeposits(SLOTS_PER_EPOCH), {skipBlsValidation})
       root = hash_tree_root(state)
 
-    db.putState(state)
+    db.putState(state[])
 
+    var tmpState = BeaconStateRef()
     check:
       db.containsState(root)
-      db.getState(root).get[] == state[]
+      db.getState(root, tmpState[], noRollback)
+      tmpState[] == state[]
diff --git a/tests/test_block_pool.nim b/tests/test_block_pool.nim
index 9423d2748e..bbb8e7cfbb 100644
--- a/tests/test_block_pool.nim
+++ b/tests/test_block_pool.nim
@@ -94,7 +94,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       var
         db = makeTestDB(SLOTS_PER_EPOCH)
         pool = BlockPool.init(db)
-        state = pool.loadTailState().data.data
+        state = newClone(pool.loadTailState().data.data)
         b1 = addTestBlock(state[], pool.tail.root)
         b1Root = hash_tree_root(b1.message)
         b2 = addTestBlock(state[], b1Root)
@@ -242,42 +242,42 @@ when const_preset == "minimal": # Too much stack space used on mainnet
         bs1_3 = b1Add.atSlot(3.Slot)
         bs2_3 = b2Add.atSlot(3.Slot)
 
-      var tmpState = clone(pool.headState)
+      var tmpState = newClone(pool.headState)
 
       # move to specific block
-      pool.updateStateData(tmpState, bs1)
+      pool.updateStateData(tmpState[], bs1)
 
       check:
         tmpState.blck == b1Add
         tmpState.data.data.slot == bs1.slot
 
       # Skip slots
-      pool.updateStateData(tmpState, bs1_3) # skip slots
+      pool.updateStateData(tmpState[], bs1_3) # skip slots
 
       check:
         tmpState.blck == b1Add
         tmpState.data.data.slot == bs1_3.slot
 
       # Move back slots, but not blocks
-      pool.updateStateData(tmpState, bs1_3.parent())
+      pool.updateStateData(tmpState[], bs1_3.parent())
       check:
         tmpState.blck == b1Add
         tmpState.data.data.slot == bs1_3.parent().slot
 
       # Move to different block and slot
-      pool.updateStateData(tmpState, bs2_3)
+      pool.updateStateData(tmpState[], bs2_3)
       check:
         tmpState.blck == b2Add
         tmpState.data.data.slot == bs2_3.slot
 
       # Move back slot and block
-      pool.updateStateData(tmpState, bs1)
+      pool.updateStateData(tmpState[], bs1)
       check:
         tmpState.blck == b1Add
         tmpState.data.data.slot == bs1.slot
 
       # Move back to genesis
-      pool.updateStateData(tmpState, bs1.parent())
+      pool.updateStateData(tmpState[], bs1.parent())
       check:
         tmpState.blck == b1Add.parent
         tmpState.data.data.slot == bs1.parent.slot
@@ -292,7 +292,7 @@ when const_preset == "minimal": # Too much stack space used on mainnet
       block:
         # Create a fork that will not be taken
         var
-          blck = makeTestBlock(pool.headState.data.data[], pool.head.blck.root)
+          blck = makeTestBlock(pool.headState.data.data, pool.head.blck.root)
         discard pool.add(hash_tree_root(blck.message), blck)
 
       for i in 0 ..< (SLOTS_PER_EPOCH * 6):
@@ -304,9 +304,9 @@ when const_preset == "minimal": # Too much stack space used on mainnet
         var
           cache = get_empty_per_epoch_cache()
           blck = makeTestBlock(
-            pool.headState.data.data[], pool.head.blck.root,
+            pool.headState.data.data, pool.head.blck.root,
             attestations = makeFullAttestations(
-              pool.headState.data.data[], pool.head.blck.root,
+              pool.headState.data.data, pool.head.blck.root,
               pool.headState.data.data.slot, cache, {}))
         let added = pool.add(hash_tree_root(blck.message), blck)
         pool.updateHead(added)
diff --git a/tests/testutil.nim b/tests/testutil.nim
index ded38888e7..b4e77adb8b 100644
--- a/tests/testutil.nim
+++ b/tests/testutil.nim
@@ -92,7 +92,7 @@ template timedTest*(name, body) =
   # TODO noto thread-safe as-is
   testTimes.add (f, name)
 
-proc makeTestDB*(tailState: BeaconStateRef, tailBlock: SignedBeaconBlock): BeaconChainDB =
+proc makeTestDB*(tailState: BeaconState, tailBlock: SignedBeaconBlock): BeaconChainDB =
   result = init(BeaconChainDB, kvStore MemStoreRef.init())
   BlockPool.preInit(result, tailState, tailBlock)
 
@@ -103,6 +103,6 @@ proc makeTestDB*(validators: int): BeaconChainDB =
       makeInitialDeposits(validators, flags = {skipBlsValidation}),
         {skipBlsValidation})
     genBlock = get_initial_beacon_block(genState[])
-  makeTestDB(genState, genBlock)
+  makeTestDB(genState[], genBlock)
 
 export inMicroseconds