From f26d6a4fd3bcf827eb45817dc70e64cb23fc7533 Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Wed, 26 Aug 2020 17:06:40 +0200
Subject: [PATCH] reuse validator key cache better (#1562)

new key cache can be used for old epochs in the same tree
---
 beacon_chain/block_pools/chain_dag.nim | 11 +++++++++++
 beacon_chain/block_pools/clearance.nim |  5 ++++-
 tests/test_block_pool.nim              | 22 +++++++++++++++++++---
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/beacon_chain/block_pools/chain_dag.nim b/beacon_chain/block_pools/chain_dag.nim
index 02d8a5a26..e56159d5b 100644
--- a/beacon_chain/block_pools/chain_dag.nim
+++ b/beacon_chain/block_pools/chain_dag.nim
@@ -136,6 +136,16 @@ proc init*(
 
   epochRef
 
+func updateKeyStores*(epochRef: EpochRef, blck: BlockRef, finalized: BlockRef) =
+  # Because key stores are additive lists, we can use a newer list whereever an
+  # older list is expected - all indices in the new list will be valid for the
+  # old list also
+  var blck = blck
+  while blck != nil and blck.slot >= finalized.slot:
+    for e in blck.epochRefs:
+      e.validator_key_store = epochRef.validator_key_store
+    blck = blck.parent
+
 func link*(parent, child: BlockRef) =
   doAssert (not (parent.root == Eth2Digest() or child.root == Eth2Digest())),
     "blocks missing root!"
@@ -400,6 +410,7 @@ proc getEpochRef*(dag: ChainDAGRef, blck: BlockRef, epoch: Epoch): EpochRef =
 
     # TODO consider constraining the number of epochrefs per state
     ancestor.blck.epochRefs.add newEpochRef
+    newEpochRef.updateKeyStores(blck.parent, dag.finalizedHead.blck)
     newEpochRef
 
 proc getState(
diff --git a/beacon_chain/block_pools/clearance.nim b/beacon_chain/block_pools/clearance.nim
index d6760849c..6731b2cd7 100644
--- a/beacon_chain/block_pools/clearance.nim
+++ b/beacon_chain/block_pools/clearance.nim
@@ -65,7 +65,10 @@ proc addResolvedBlock(
     let prevEpochRef = blockRef.findEpochRef(blockEpoch - 1)
 
     epochRef = EpochRef.init(state.data.data, cache, prevEpochRef)
-    blockRef.epochAncestor(blockEpoch).blck.epochRefs.add epochRef
+    let ancestor = blockRef.epochAncestor(blockEpoch)
+    epochRef.updateKeyStores(ancestor.blck.parent, dag.finalizedHead.blck)
+
+    ancestor.blck.epochRefs.add epochRef
 
   dag.blocks[blockRoot] = blockRef
   trace "Populating block dag", key = blockRoot, val = blockRef
diff --git a/tests/test_block_pool.nim b/tests/test_block_pool.nim
index de26ab444..64cb68136 100644
--- a/tests/test_block_pool.nim
+++ b/tests/test_block_pool.nim
@@ -343,15 +343,31 @@ suiteReport "chain DAG finalization tests" & preset():
     check:
       dag.heads.len() == 1
 
+    let headER = dag.heads[0].findEpochRef(dag.heads[0].slot.epoch)
+    check:
+
       # Epochrefs should share validator key set when the validator set is
       # stable
-      not dag.heads[0].findEpochRef(dag.heads[0].slot.epoch).isNil
+      not headER.isNil
       not dag.heads[0].findEpochRef(dag.heads[0].slot.epoch - 1).isNil
-      dag.heads[0].findEpochRef(dag.heads[0].slot.epoch) !=
+      headER !=
         dag.heads[0].findEpochRef(dag.heads[0].slot.epoch - 1)
-      dag.heads[0].findEpochRef(dag.heads[0].slot.epoch).validator_key_store[1] ==
+      headER.validator_key_store[1] ==
         dag.heads[0].findEpochRef(dag.heads[0].slot.epoch - 1).validator_key_store[1]
 
+    block:
+      var cur = dag.heads[0]
+      while cur != nil:
+        if cur.slot < dag.finalizedHead.blck.slot:
+          # Cache should be cleaned on finalization
+          check: cur.epochRefs.len == 0
+        else:
+          # EpochRef validator keystores should back-propagate to all previous
+          # epochs
+          for e in cur.epochRefs:
+            check (addr headER.validator_keys) == (addr e.validator_keys)
+        cur = cur.parent
+
     block:
       # The late block is a block whose parent was finalized long ago and thus
       # is no longer a viable head candidate