simplify RANDAO recovery in `ShufflingRef` acceleration (#5183)

Current RANDAO recovery logic is quite complex as it optimizes for the
minimum amount of database reads. Loading blocks isn't the bottleneck
though, so rather make the implementation more concise by avoiding the
complex strategy planning step. Note that this also prepares for an even
faster implementation for post-merge blocks in the future that extracts
RANDAO from `ExecutionPayload` directly if available, so even in cases
where efficiency is slightly lower, only historical data is affected.

`time nim c -r tests/test_blockchain_dag` (cached binary):

- new: 145.45s, 133.59s, 144.65s, 127.69s, 136.14s
- old: 149.15s, 150.84s, 135.77s, 137.49s, 133.89s

beacon_chain/consensus_object_pools/blockchain_dag.nim

@@ -1330,7 +1330,7 @@ proc getFinalizedEpochRef*(dag: ChainDAGRef): EpochRef =
     dag.finalizedHead.blck, dag.finalizedHead.slot.epoch, false).expect(
       "getEpochRef for finalized head should always succeed")
 
-func ancestorSlotForShuffling*(
+func ancestorSlotForAttesterShuffling*(
     dag: ChainDAGRef, state: ForkyHashedBeaconState,
     blck: BlockRef, epoch: Epoch): Opt[Slot] =
   ## Return slot of `blck` ancestor to which `state` can be rewinded
@@ -1385,116 +1385,57 @@ func ancestorSlotForShuffling*(
   doAssert dependentSlot >= lowSlot
   ok min(min(stateBid.slot, ancestorBlck.slot), dependentSlot)
 
-proc mixRandao(
+type AttesterRandaoMix = tuple[dependentBid: BlockId, mix: Eth2Digest]
-    dag: ChainDAGRef, mix: var Eth2Digest,
-    bid: BlockId): Opt[void] =
-  ## Mix in/out the RANDAO reveal from the given block.
-  let bdata = ? dag.getForkedBlock(bid)
-  withBlck(bdata):  # See `process_randao` / `process_randao_mixes_reset`
-    mix.data.mxor eth2digest(blck.message.body.randao_reveal.toRaw()).data
-  ok()
 
-proc computeRandaoMix(
+proc computeAttesterRandaoMix(
     dag: ChainDAGRef, state: ForkyHashedBeaconState,
-    blck: BlockRef, epoch: Epoch
+    blck: BlockRef, epoch: Epoch): Opt[AttesterRandaoMix] =
-): Opt[tuple[dependentBid: BlockId, mix: Eth2Digest]] =
   ## Compute the requested RANDAO mix for `blck@epoch` based on `state`.
-  ## `state` must have the correct `get_active_validator_indices` for `epoch`.
+  ## If `state` has unviable `get_active_validator_indices`, return `none`.
-  ## RANDAO reveals of blocks from `state.data.slot` back to `ancestorSlot` are
+
-  ## mixed out from `state.data.randao_mixes`, and RANDAO reveals from blocks
+  # Check `state` has locked-in `get_active_validator_indices` for `epoch`
-  ## up through `epoch.attester_dependent_slot` are mixed in.
   let
     stateSlot = state.data.slot
     dependentSlot = epoch.attester_dependent_slot
-    # Check `state` has locked-in `get_active_validator_indices` for `epoch`
+    ancestorSlot = ? dag.ancestorSlotForAttesterShuffling(state, blck, epoch)
-    ancestorSlot = ? dag.ancestorSlotForShuffling(state, blck, epoch)
   doAssert ancestorSlot <= stateSlot
   doAssert ancestorSlot <= dependentSlot
 
-  # Load initial mix
+  var mix: Eth2Digest
-  var mix {.noinit.}: Eth2Digest
+  proc mixToAncestor(highBid: BlockId): Opt[void] =
+    ## Mix in/out RANDAO reveals back to `ancestorSlot`
+    var bid = highBid
+    while bid.slot > ancestorSlot:
+      let bdata = ? dag.getForkedBlock(bid)
+      withBlck(bdata):  # See `process_randao` / `process_randao_mixes_reset`
+        mix.data.mxor eth2digest(blck.message.body.randao_reveal.toRaw()).data
+      bid = ? dag.parent(bid)
+    ok()
+
+  # Determine block for obtaining RANDAO mix
   let
-    stateEpoch = stateSlot.epoch
+    dependentBid =
-    ancestorEpoch = ancestorSlot.epoch
+      if dependentSlot >= dag.finalizedHead.slot:
-    highRandaoSlot =
+        var b = blck.get_ancestor(dependentSlot)
-      # `randao_mixes[ancestorEpoch]`
-      if stateEpoch == ancestorEpoch:
-        stateSlot
-      else:
-        (ancestorEpoch + 1).start_slot - 1
-    startSlot =
-      if ancestorEpoch == GENESIS_EPOCH:
-        # Can only move backward
-        mix = state.data.get_randao_mix(ancestorEpoch)
-        highRandaoSlot
-      else:
-        # `randao_mixes[ancestorEpoch - 1]`
-        let lowRandaoSlot = ancestorEpoch.start_slot - 1
-        if highRandaoSlot - ancestorSlot < ancestorSlot - lowRandaoSlot:
-          mix = state.data.get_randao_mix(ancestorEpoch)
-          highRandaoSlot
-        else:
-          mix = state.data.get_randao_mix(ancestorEpoch - 1)
-          lowRandaoSlot
-    slotsToMix =
-      if startSlot > ancestorSlot:
-        (ancestorSlot + 1) .. startSlot
-      else:
-        (startSlot + 1) .. ancestorSlot
-    highRoot =
-      if slotsToMix.b == stateSlot:
-        state.latest_block_root
-      else:
-        doAssert slotsToMix.b < stateSlot
-        state.data.get_block_root_at_slot(slotsToMix.b)
-
-  # Move `mix` from `startSlot` to `ancestorSlot`
-  var bid =
-    if slotsToMix.b >= dag.finalizedHead.slot:
-      var b = ? dag.getBlockRef(highRoot)
-      let lowSlot = max(slotsToMix.a, dag.finalizedHead.slot)
-      while b.bid.slot > lowSlot:
-        ? dag.mixRandao(mix, b.bid)
-        b = b.parent
         doAssert b != nil
-      b.bid
+        b.bid
-    else:
-      var highSlot = slotsToMix.b
-      const availableSlots = SLOTS_PER_HISTORICAL_ROOT
-      let lowSlot = max(state.data.slot, availableSlots.Slot) - availableSlots
-      while highSlot > lowSlot and
-          state.data.get_block_root_at_slot(highSlot - 1) == highRoot:
-        dec highSlot
-      if highSlot + SLOTS_PER_HISTORICAL_ROOT > state.data.slot:
-        BlockId(slot: highSlot, root: highRoot)
       else:
-        let bsi = ? dag.getBlockIdAtSlot(highSlot)
+        let bsi = ? dag.getBlockIdAtSlot(dependentSlot)
-        doAssert bsi.bid.root == highRoot
         bsi.bid
-  while bid.slot >= slotsToMix.a:
-    ? dag.mixRandao(mix, bid)
-    bid = ? dag.parent(bid)
 
-  # Move `mix` from `ancestorSlot` to `dependentSlot`
+  # Mix in RANDAO from `blck`
-  var dependentBid {.noinit.}: BlockId
+  if ancestorSlot < dependentBid.slot:
-  bid =
+    ? mixToAncestor(dependentBid)
-    if dependentSlot >= dag.finalizedHead.slot:
+
-      var b = blck.get_ancestor(dependentSlot)
+  # Mix in RANDAO from `state`
-      doAssert b != nil
+  let ancestorEpoch = ancestorSlot.epoch
-      dependentBid = b.bid
+  if ancestorEpoch + EPOCHS_PER_HISTORICAL_VECTOR <= stateSlot.epoch:
-      let lowSlot = max(ancestorSlot, dag.finalizedHead.slot)
+    return Opt.none(AttesterRandaoMix)
-      while b.bid.slot > lowSlot:
+  let mixRoot = state.dependent_root(ancestorEpoch + 1)
-        ? dag.mixRandao(mix, b.bid)
+  if mixRoot.isZero:
-        b = b.parent
+    return Opt.none(AttesterRandaoMix)
-        doAssert b != nil
+  ? mixToAncestor(? dag.getBlockId(mixRoot))
-      b.bid
+  mix.data.mxor state.data.get_randao_mix(ancestorEpoch).data
-    else:
-      let bsi = ? dag.getBlockIdAtSlot(dependentSlot)
-      dependentBid = bsi.bid
-      bsi.bid
-  while bid.slot > ancestorSlot:
-    ? dag.mixRandao(mix, bid)
-    bid = ? dag.parent(bid)
 
   ok (dependentBid: dependentBid, mix: mix)
 
@@ -1502,7 +1443,7 @@ proc computeShufflingRefFromState*(
     dag: ChainDAGRef, state: ForkyHashedBeaconState,
     blck: BlockRef, epoch: Epoch): Opt[ShufflingRef] =
   let (dependentBid, mix) =
-    ? dag.computeRandaoMix(state, blck, epoch)
+    ? dag.computeAttesterRandaoMix(state, blck, epoch)
 
   return ok ShufflingRef(
     epoch: epoch,
@@ -1555,15 +1496,9 @@ proc computeShufflingRefFromDatabase*(
 proc computeShufflingRef*(
     dag: ChainDAGRef, blck: BlockRef, epoch: Epoch): Opt[ShufflingRef] =
   # Try to compute `ShufflingRef` from states available in memory
-  template tryWithState(state: ForkedHashedBeaconState) =
+  let shufflingRef = dag.computeShufflingRefFromMemory(blck, epoch)
-    withState(state):
+  if shufflingRef.isOk:
-      let shufflingRef =
+    return shufflingRef
-        dag.computeShufflingRefFromState(forkyState, blck, epoch)
-      if shufflingRef.isOk:
-        return shufflingRef
-  tryWithState dag.headState
-  tryWithState dag.epochRefState
-  tryWithState dag.clearanceState
 
   # Fall back to database
   dag.computeShufflingRefFromDatabase(blck, epoch)

tests/test_blockchain_dag.nim

@@ -1311,7 +1311,8 @@ suite "Shufflings":
             blckEpoch = blck.bid.slot.epoch
             minEpoch = min(stateEpoch, blckEpoch)
           if compute_activation_exit_epoch(minEpoch) <= epoch or
-              dag.ancestorSlotForShuffling(forkyState, blck, epoch).isNone:
+              dag.ancestorSlotForAttesterShuffling(
+                forkyState, blck, epoch).isNone:
             check shufflingRef.isErr
           else:
             check shufflingRef.isOk