simplify fork choice code (#1521)

* standardize init
* avoid loading state on init
* avoid some inefficient exception-based code
* remove some TODO

beacon_chain/attestation_pool.nim

@@ -29,11 +29,12 @@ proc init*(T: type AttestationPool, chainDag: ChainDAGRef, quarantine: Quarantin
   #      should probably be removed as a dependency of AttestationPool (or some other
   #      smart refactoring)
 
-  let tmpState = newClone(chainDag.headState)
-  chainDag.withState(tmpState[], chainDag.finalizedHead):
-    var cache = StateCache()
-    var forkChoice = initForkChoice(
-      tmpState[], getEpochInfo(chainDag.finalizedHead.blck, state, cache)).get()
+  let
+    finalizedEpochRef = chainDag.getEpochRef(
+      chainDag.finalizedHead.blck, chainDag.finalizedHead.slot.epoch())
+
+  var forkChoice = ForkChoice.init(
+    finalizedEpochRef, chainDag.finalizedHead.blck)
 
   # Feed fork choice with unfinalized history - during startup, block pool only
   # keeps track of a single history so we just need to follow it

beacon_chain/fork_choice/fork_choice.nim

@@ -28,8 +28,6 @@ export sets, results, fork_choice_types
 # - Prysmatic writeup: https://hackmd.io/bABJiht3Q9SyV3Ga4FT9lQ#High-level-concept
 # - Gasper Whitepaper: https://arxiv.org/abs/2003.03052
 
-const DefaultPruneThreshold = 256
-
 # Forward declarations
 # ----------------------------------------------------------------------
 
@@ -49,54 +47,41 @@ func compute_deltas(
 logScope:
   topics = "fork_choice"
 
-proc initForkChoiceBackend*(justified_epoch: Epoch,
-                            finalized_epoch: Epoch,
-                            finalized_root: Eth2Digest,
-                            ): FcResult[ForkChoiceBackend] =
-  var proto_array = ProtoArray(
-    prune_threshold: DefaultPruneThreshold,
-    justified_epoch: finalized_epoch,
-    finalized_epoch: finalized_epoch
+proc init*(T: type ForkChoiceBackend,
+           justified_epoch: Epoch,
+           finalized_root: Eth2Digest,
+           finalized_epoch: Epoch): T =
+  T(
+    proto_array: ProtoArray.init(
+      justified_epoch,
+      finalized_root,
+      finalized_epoch
+    )
   )
 
-  ? proto_array.on_block(
-    finalized_root,
-    hasParentInForkChoice = false,
-    Eth2Digest(),
-    finalized_epoch,
-    finalized_epoch
-  )
-
-  ok(ForkChoiceBackend(
-    proto_array: proto_array,
-  ))
-
-proc initForkChoice*(finalizedState: StateData,
-                     epochRef: EpochRef): FcResult[ForkChoice] =
-  ## Initialize a fork choice context
+proc init*(T: type ForkChoice,
+           epochRef: EpochRef,
+           blck: BlockRef): T =
+  ## Initialize a fork choice context for a genesis state - in the genesis
+  ## state, the justified and finalized checkpoints are the same, so only one
+  ## is used here
   debug "Initializing fork choice",
-    state_epoch = finalizedState.data.data.get_current_epoch(),
-    blck = shortLog(finalizedState.blck)
-
-  let finalized_epoch = finalizedState.data.data.get_current_epoch()
+    epoch = epochRef.epoch, blck = shortLog(blck)
 
   let
-    justified = BalanceCheckpoint(
-      blck: finalizedState.blck, epochRef: epochRef)
-    finalized = Checkpoint(
-      root: finalizedState.blck.root, epoch: finalized_epoch)
+    justified = BalanceCheckpoint(blck: blck, epochRef: epochRef)
+    finalized = Checkpoint(root: blck.root, epoch: epochRef.epoch)
+    best_justified = Checkpoint(
+      root: justified.blck.root, epoch: justified.epochRef.epoch)
 
-  let backend = ? initForkChoiceBackend(
-    finalized_epoch, finalized_epoch, finalizedState.blck.root)
-
-  ok(ForkChoice(
-    backend: backend,
+  ForkChoice(
+    backend: ForkChoiceBackend.init(
+      epochRef.epoch, blck.root, epochRef.epoch),
     checkpoints: Checkpoints(
       justified: justified,
-      best_justified:
-        Checkpoint(root: justified.blck.root, epoch: justified.epochRef.epoch),
-      finalized: finalized)
-  ))
+      finalized: finalized,
+      best_justified: best_justified)
+  )
 
 func extend[T](s: var seq[T], minLen: int) =
   ## Extend a sequence so that it can contains at least `minLen` elements.
@@ -288,8 +273,7 @@ proc process_block*(self: var ForkChoiceBackend,
                     justified_epoch: Epoch,
                     finalized_epoch: Epoch): FcResult[void] =
   self.proto_array.on_block(
-    block_root, hasParentInForkChoice = true, parent_root,
-    justified_epoch, finalized_epoch)
+    block_root, parent_root, justified_epoch, finalized_epoch)
 
 proc process_block*(self: var ForkChoice,
                     dag: ChainDAGRef,

beacon_chain/fork_choice/proto_array.nim

@@ -23,6 +23,8 @@ logScope:
 
 export results
 
+const DefaultPruneThreshold* = 256
+
 # https://github.com/ethereum/eth2.0-specs/blob/v0.11.1/specs/phase0/fork-choice.md
 # This is a port of https://github.com/sigp/lighthouse/pull/804
 # which is a port of "Proto-Array": https://github.com/protolambda/lmd-ghost
@@ -46,24 +48,14 @@ func tiebreak(a, b: Eth2Digest): bool =
     # else we have equality so far
   return true
 
-template getOrFailcase*[K, V](table: Table[K, V], key: K, failcase: untyped): V =
-  ## Get a value from a Nim Table, turning KeyError into
-  ## the "failcase"
-  block:
-    # TODO: try/except expression with Nim v1.2.0:
-    #       https://github.com/status-im/nim-beacon-chain/pull/865#discussion_r404856551
-    var value: V
-    try:
-      value = table[key]
-    except KeyError:
-      failcase
-    value
-
 template unsafeGet*[K, V](table: Table[K, V], key: K): V =
   ## Get a value from a Nim Table, turning KeyError into
   ## an AssertionError defect
-  getOrFailcase(table, key):
-    doAssert false, "The " & astToStr(table) & " table shouldn't miss a key"
+  # Pointer is used to work around the lack of a `var` withValue
+  try:
+    table[key]
+  except KeyError as exc:
+    raiseAssert(exc.msg)
 
 # Forward declarations
 # ----------------------------------------------------------------------
@@ -76,6 +68,29 @@ func node_leads_to_viable_head(self: ProtoArray, node: ProtoNode): FcResult[bool
 # ProtoArray routines
 # ----------------------------------------------------------------------
 
+func init*(T: type ProtoArray,
+           justified_epoch: Epoch,
+           finalized_root: Eth2Digest,
+           finalized_epoch: Epoch,
+           prune_threshold = DefaultPruneThreshold): T =
+  let node = ProtoNode(
+    root: finalized_root,
+    parent: none(int),
+    justified_epoch: justified_epoch,
+    finalized_epoch: finalized_epoch,
+    weight: 0,
+    best_child: none(int),
+    best_descendant: none(int)
+  )
+
+  T(
+    prune_threshold: DefaultPruneThreshold,
+    justified_epoch: justified_epoch,
+    finalized_epoch: finalized_epoch,
+    nodes: @[node],
+    indices: {node.root: 0}.toTable()
+  )
+
 func apply_score_changes*(
        self: var ProtoArray,
        deltas: var openarray[Delta],
@@ -158,7 +173,6 @@ func apply_score_changes*(
 func on_block*(
        self: var ProtoArray,
        root: Eth2Digest,
-       hasParentInForkChoice: bool,
        parent: Eth2Digest,
        justified_epoch: Epoch,
        finalized_epoch: Epoch
@@ -175,24 +189,21 @@ func on_block*(
   if root in self.indices:
     return ok()
 
-  var parent_index: Option[int]
-  if not hasParentInForkChoice:
-      # Genesis (but Genesis might not be default(Eth2Digest))
-    parent_index = none(int)
-  elif parent notin self.indices:
+  var parent_index: Index
+  self.indices.withValue(parent, index) do:
+    parent_index = index[]
+  do:
     return err ForkChoiceError(
       kind: fcUnknownParent,
       child_root: root,
       parent_root: parent
     )
-  else:
-    parent_index = some(self.indices.unsafeGet(parent))
 
   let node_index = self.nodes.len
 
   let node = ProtoNode(
     root: root,
-    parent: parent_index,
+    parent: some(parent_index),
     justified_epoch: justified_epoch,
     finalized_epoch: finalized_epoch,
     weight: 0,
@@ -201,10 +212,9 @@ func on_block*(
   )
 
   self.indices[node.root] = node_index
-  self.nodes.add node # TODO: if this is costly, we can setLen + construct the node in-place
+  self.nodes.add node
 
-  if parent_index.isSome(): # parent_index is always valid except for Genesis
-    ? self.maybe_update_best_child_and_descendant(parent_index.unsafeGet(), node_index)
+  ? self.maybe_update_best_child_and_descendant(parent_index, node_index)
 
   return ok()
 
@@ -220,7 +230,10 @@ func find_head*(
   ## is not followed by `apply_score_changes` as `on_new_block` does not
   ## update the whole tree.
 
-  let justified_index = self.indices.getOrFailcase(justified_root):
+  var justified_index: Index
+  self.indices.withValue(justified_root, value) do:
+    justified_index = value[]
+  do:
     return err ForkChoiceError(
       kind: fcJustifiedNodeUnknown,
       block_root: justified_root
@@ -235,11 +248,7 @@ func find_head*(
   template justified_node: untyped = self.nodes[justified_index]
     # Alias, IndexError are defects
 
-  let best_descendant_index = block:
-    if justified_node.best_descendant.isSome():
-      justified_node.best_descendant.unsafeGet()
-    else:
-      justified_index
+  let best_descendant_index = justified_node.best_descendant.get(justified_index)
 
   if best_descendant_index notin {0..self.nodes.len-1}:
     return err ForkChoiceError(
@@ -281,7 +290,11 @@ func maybe_prune*(
   ## - The finalized epoch is less than the current one
   ## - The finalized epoch matches the current one but the finalized root is different
   ## - Internal error due to invalid indices in `self`
-  let finalized_index = self.indices.getOrFailcase(finalized_root):
+
+  var finalized_index: int
+  self.indices.withValue(finalized_root, value) do:
+    finalized_index = value[]
+  do:
     return err ForkChoiceError(
       kind: fcFinalizedNodeUnknown,
       block_root: finalized_root
@@ -382,8 +395,8 @@ func maybe_update_best_child_and_descendant(
     )
 
   # Aliases
-  template child: untyped {.dirty.} = self.nodes[child_index]
-  template parent: untyped {.dirty.} = self.nodes[parent_index]
+  template child: untyped = self.nodes[child_index]
+  template parent: untyped = self.nodes[parent_index]
 
   let child_leads_to_viable_head = ? self.node_leads_to_viable_head(child)
 

tests/fork_choice/scenarios/ffg_01.nim

@@ -12,11 +12,11 @@ proc setup_finality_01(): tuple[fork_choice: ForkChoiceBackend, ops: seq[Operati
   let GenesisRoot = fakeHash(0)
 
   # Initialize the fork choice context
-  result.fork_choice = initForkChoiceBackend(
+  result.fork_choice = ForkChoiceBackend.init(
     justified_epoch = Epoch(1),
     finalized_epoch = Epoch(1),
     finalized_root = GenesisRoot
-  ).get()
+  )
 
   # ----------------------------------
 

tests/fork_choice/scenarios/ffg_02.nim

@@ -12,11 +12,11 @@ proc setup_finality_02(): tuple[fork_choice: ForkChoiceBackend, ops: seq[Operati
   let GenesisRoot = fakeHash(0)
 
   # Initialize the fork choice context
-  result.fork_choice = initForkChoiceBackend(
+  result.fork_choice = ForkChoiceBackend.init(
     justified_epoch = Epoch(1),
     finalized_epoch = Epoch(1),
     finalized_root = GenesisRoot
-  ).get()
+  )
 
   # ----------------------------------
 

tests/fork_choice/scenarios/no_votes.nim

@@ -12,11 +12,11 @@ proc setup_no_votes(): tuple[fork_choice: ForkChoiceBackend, ops: seq[Operation]
   let GenesisRoot = fakeHash(0)
 
   # Initialize the fork choice context
-  result.fork_choice = initForkChoiceBackend(
+  result.fork_choice = ForkChoiceBackend.init(
     justified_epoch = Epoch(1),
     finalized_epoch = Epoch(1),
     finalized_root = GenesisRoot
-  ).get()
+  )
 
   # ----------------------------------
 

tests/fork_choice/scenarios/votes.nim

@@ -12,11 +12,11 @@ proc setup_votes(): tuple[fork_choice: ForkChoiceBackend, ops: seq[Operation]] =
   let GenesisRoot = fakeHash(0)
 
   # Initialize the fork choice context
-  result.fork_choice = initForkChoiceBackend(
+  result.fork_choice = ForkChoiceBackend.init(
     justified_epoch = Epoch(1),
     finalized_epoch = Epoch(1),
     finalized_root = GenesisRoot
-  ).get()
+  )
 
   # ----------------------------------