fix state cache loading

* load the cache of the current state epoch instead of the target state
epoch, when applying states and slots
* load state cache for each slot/block (for longer slot jumps)
* load state cache after full updateStateData
* look up two state cache epochs, instead of the same epoch twice :)

beacon_chain/consensus_object_pools/blockchain_dag.nim

@@ -266,16 +266,15 @@ func loadStateCache(
   # functions
 
   template load(e: Epoch) =
-    if epoch notin cache.shuffled_active_validator_indices:
+    if e notin cache.shuffled_active_validator_indices:
-      let epochRef = dag.findEpochRef(blck, epoch)
+      let epochRef = dag.findEpochRef(blck, e)
       if epochRef != nil:
         cache.shuffled_active_validator_indices[epochRef.epoch] =
           epochRef.shuffled_active_validator_indices
 
-        if epochRef.epoch == epoch:
         for i, idx in epochRef.beacon_proposers:
           cache.beacon_proposer_indices[
-              epoch.compute_start_slot_at_epoch + i] = idx
+            epochRef.epoch.compute_start_slot_at_epoch + i] = idx
 
   load(epoch)
 
@@ -662,6 +661,8 @@ proc advanceSlots(
   # target
   doAssert getStateField(state, slot) <= slot
   while getStateField(state, slot) < slot:
+    loadStateCache(dag, cache, state.blck, getStateField(state, slot).epoch)
+
     doAssert process_slots(
         state.data, getStateField(state, slot) + 1, cache, rewards,
         dag.updateFlags),
@@ -683,7 +684,7 @@ proc applyBlock(
     doAssert (addr(statePtr.data) == addr v)
     statePtr[] = dag.headState
 
-  loadStateCache(dag, cache, blck.refs, blck.data.message.slot.epoch)
+  loadStateCache(dag, cache, state.blck, getStateField(state, slot).epoch)
 
   let ok = state_transition(
     dag.runtimePreset, state.data, blck.data,
@@ -832,11 +833,12 @@ proc updateStateData*(
       dag.applyBlock(state, dag.get(ancestors[i]), {}, cache, rewards)
     doAssert ok, "Blocks in database should never fail to apply.."
 
-  loadStateCache(dag, cache, bs.blck, bs.slot.epoch)
-
   # ...and make sure to process empty slots as requested
   dag.advanceSlots(state, bs.slot, save, cache, rewards)
 
+  # ...and make sure to load the state cache, if it exists
+  loadStateCache(dag, cache, state.blck, getStateField(state, slot).epoch)
+
   let
     assignDur = assignTick - startTick
     replayDur = Moment.now() - assignTick

tests/test_block_pool.nim

@@ -357,9 +357,6 @@ suite "chain DAG finalization tests" & preset():
 
     assign(tmpState[], dag.headState.data)
 
-    # StateCache is designed to only hold a single linear history at a time
-    cache = StateCache()
-
     for i in 0 ..< (SLOTS_PER_EPOCH * 6):
       if i == 1:
         # There are 2 heads now because of the fork at slot 1
@@ -393,6 +390,22 @@ suite "chain DAG finalization tests" & preset():
       for er in dag.epochRefs:
         check: er[1] == nil or er[1].epoch >= dag.finalizedHead.slot.epoch
 
+    block:
+      let tmpStateData = assignClone(dag.headState)
+
+      # Check that cached data is available after updateStateData - since we
+      # just processed the head the relevant epochrefs should not have been
+      # evicted yet
+      cache = StateCache()
+      updateStateData(
+        dag, tmpStateData[], dag.head.atSlot(dag.head.slot), false, cache)
+
+      check:
+        dag.head.slot.epoch in cache.shuffled_active_validator_indices
+        (dag.head.slot.epoch - 1) in cache.shuffled_active_validator_indices
+
+        dag.head.slot in cache.beacon_proposer_indices
+
     block:
       # The late block is a block whose parent was finalized long ago and thus
       # is no longer a viable head candidate