From 6e40ca5e1502c1b3741aa16401874000f39473e1 Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Tue, 10 Nov 2020 14:48:59 +0100
Subject: [PATCH] Use separate cache for getting epochref

...else tmpState gets overwritten when producing blocks and requesting
epochRef for finalized state.
---
 .../block_pools/block_pools_types.nim         | 20 ++++++++++---------
 beacon_chain/block_pools/chain_dag.nim        | 13 +++++++-----
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/beacon_chain/block_pools/block_pools_types.nim b/beacon_chain/block_pools/block_pools_types.nim
index 6ea736d76..f7e556b13 100644
--- a/beacon_chain/block_pools/block_pools_types.nim
+++ b/beacon_chain/block_pools/block_pools_types.nim
@@ -110,10 +110,6 @@ type
     heads*: seq[BlockRef] ##\
     ## Candidate heads of candidate chains
 
-    head*: BlockRef ##\
-    ## The latest block we know about, that's been chosen as a head by the fork
-    ## choice rule
-
     finalizedHead*: BlockSlot ##\
     ## The latest block that was finalized according to the block in head
     ## Ancestors of this block are guaranteed to have 1 child only.
@@ -122,14 +118,18 @@ type
     # Rewinder - Mutable state processing
 
     headState*: StateData ##\
-    ## State given by the head block; only update in `updateHead`, not anywhere
-    ## else via `withState`
+    ## State given by the head block - must only be updated in `updateHead` -
+    ## always matches dag.head
 
-    tmpState*: StateData ## Scratchpad - may be any state
+    epochRefState*: StateData ##\
+      ## State used to produce epochRef instances - must only be used in
+      ## `getEpochRef`
 
     clearanceState*: StateData ##\
-      ## Cached state used during block clearance - should only be used in the
-      ## clearance module to avoid the risk of modifying it in a callback
+      ## Cached state used during block clearance - must only be used in
+      ## clearance module
+
+    tmpState*: StateData ## Scratchpad - may be any state
 
     updateFlags*: UpdateFlags
 
@@ -202,6 +202,8 @@ type
 
 template validator_keys*(e: EpochRef): untyped = e.validator_key_store[1][]
 
+template head*(v: ChainDagRef): BlockRef = v.headState.blck
+
 func shortLog*(v: BlockSlot): string =
   if v.blck.slot == v.slot:
     &"{v.blck.root.data.toOpenArray(0, 3).toHex()}:{v.blck.slot}"
diff --git a/beacon_chain/block_pools/chain_dag.nim b/beacon_chain/block_pools/chain_dag.nim
index 199cbac4b..ce151df48 100644
--- a/beacon_chain/block_pools/chain_dag.nim
+++ b/beacon_chain/block_pools/chain_dag.nim
@@ -389,13 +389,13 @@ proc init*(T: type ChainDAGRef,
   let res = ChainDAGRef(
     blocks: blocks,
     tail: tailRef,
-    head: headRef,
     genesis: genesisRef,
     db: db,
     heads: @[headRef],
     headState: tmpState[],
-    tmpState: tmpState[],
+    epochRefState: tmpState[],
     clearanceState: tmpState[],
+    tmpState: tmpState[],
 
     # The only allowed flag right now is verifyFinalization, as the others all
     # allow skipping some validation.
@@ -446,7 +446,7 @@ proc getEpochRef*(dag: ChainDAGRef, blck: BlockRef, epoch: Epoch): EpochRef =
   let
     ancestor = blck.epochAncestor(epoch)
 
-  dag.withState(dag.tmpState, ancestor):
+  dag.withState(dag.epochRefState, ancestor):
     let
       prevEpochRef = if dag.tail.slot.epoch >= epoch: nil
                      else: blck.findEpochRef(epoch - 1)
@@ -702,6 +702,11 @@ proc updateStateData*(
       found = true
       break
 
+    if canAdvance(dag.epochRefState, cur):
+      assign(state, dag.epochRefState)
+      found = true
+      break
+
     if cur.slot == cur.blck.slot:
       # This is not an empty slot, so the block will need to be applied to
       # eventually reach bs
@@ -812,8 +817,6 @@ proc updateHead*(
   updateStateData(
     dag, dag.headState, newHead.atSlot(newHead.slot), false, cache)
 
-  dag.head = newHead
-
   if not lastHead.isAncestorOf(newHead):
     notice "Updated head block with chain reorg",
       lastHead = shortLog(lastHead),