generalize `ShufflingRef` acceleration logic (#5197)

Split up the `ShufflingRef` acceleration logic into generically usable
parts and attester shuffling specific parts. The generic parts could be
used to accelerate other purposes, e.g., REST `/states/xxx/randao` API.

beacon_chain/consensus_object_pools/blockchain_dag.nim

@@ -771,6 +771,26 @@ proc getStateByParent(
   dag.db.getState(
     dag.cfg, summary.parent_root, parentMinSlot..slot, state, rollback)
 
+proc getNearbyState(
+    dag: ChainDAGRef, state: ref ForkedHashedBeaconState, bid: BlockId,
+    lowSlot: Slot): Opt[void] =
+  ## Load state from DB that is close to `bid` and has at least slot `lowSlot`.
+  var
+    e = bid.slot.epoch
+    b = bid
+  while true:
+    let stateSlot = e.start_slot
+    if stateSlot < lowSlot:
+      return err()
+    b = (? dag.atSlot(b, max(stateSlot, 1.Slot) - 1)).bid
+    let bsi = BlockSlotId.init(b, stateSlot)
+    if not dag.getState(bsi, state[]):
+      if e == GENESIS_EPOCH:
+        return err()
+      dec e
+      continue
+    return ok()
+
 proc currentSyncCommitteeForPeriod*(
     dag: ChainDAGRef,
     tmpState: var ForkedHashedBeaconState,
@@ -1368,59 +1388,36 @@ proc ancestorSlot*(
 
   Opt.some stateBid.slot
 
-proc ancestorSlotForAttesterShuffling*(
+proc computeRandaoMix(
-    dag: ChainDAGRef, state: ForkyHashedBeaconState,
+    dag: ChainDAGRef, bdata: ForkedTrustedSignedBeaconBlock): Opt[Eth2Digest] =
-    blck: BlockRef, epoch: Epoch): Opt[Slot] =
+  ## Compute the requested RANDAO mix for `bdata` without `state`, if possible.
-  ## Return slot of `blck` ancestor to which `state` can be rewinded
+  withBlck(bdata):
-  ## so that RANDAO at `epoch.attester_dependent_slot` can be computed.
-  ## Return `err` if `state` is unviable to compute shuffling for `blck@epoch`.
-
-  # A state must be somewhat recent so that `get_active_validator_indices`
-  # for the queried `epoch` cannot be affected by any such skipped processing.
-  const numDelayEpochs = compute_activation_exit_epoch(GENESIS_EPOCH).uint64
-  let
-    lowEpoch = max(epoch, (numDelayEpochs - 1).Epoch) - (numDelayEpochs - 1)
-    ancestorSlot = ? dag.ancestorSlot(state, blck.bid, lowEpoch.start_slot)
-  Opt.some min(ancestorSlot, epoch.attester_dependent_slot)
-
-type AttesterRandaoMix = tuple[dependentBid: BlockId, mix: Eth2Digest]
-
-proc computeAttesterRandaoMix(
-    dag: ChainDAGRef, state: ForkyHashedBeaconState,
-    blck: BlockRef, epoch: Epoch): Opt[AttesterRandaoMix] =
-  ## Compute the requested RANDAO mix for `blck@epoch` based on `state`.
-  ## If `state` has unviable `get_active_validator_indices`, return `none`.
-
-  # Check `state` has locked-in `get_active_validator_indices` for `epoch`
-  let
-    stateSlot = state.data.slot
-    dependentSlot = epoch.attester_dependent_slot
-    ancestorSlot = ? dag.ancestorSlotForAttesterShuffling(state, blck, epoch)
-  doAssert ancestorSlot <= stateSlot
-  doAssert ancestorSlot <= dependentSlot
-
-  # Determine block for obtaining RANDAO mix
-  let
-    dependentBid =
-      if dependentSlot >= dag.finalizedHead.slot:
-        var b = blck.get_ancestor(dependentSlot)
-        doAssert b != nil
-        b.bid
-      else:
-        let bsi = ? dag.getBlockIdAtSlot(dependentSlot)
-        bsi.bid
-    dependentBdata = ? dag.getForkedBlock(dependentBid)
-  var mix {.noinit.}: Eth2Digest
-
-  # If `dependentBid` is post merge, RANDAO information is available
-  withBlck(dependentBdata):
     when consensusFork >= ConsensusFork.Bellatrix:
       if blck.message.is_execution_block:
-        mix = eth2digest(blck.message.body.randao_reveal.toRaw())
+        var mix = eth2digest(blck.message.body.randao_reveal.toRaw())
         mix.data.mxor blck.message.body.execution_payload.prev_randao.data
-        return ok (dependentBid: dependentBid, mix: mix)
+        return ok mix
+  Opt.none(Eth2Digest)
 
-  # RANDAO mix has to be recomputed from `blck` and `state`
+proc computeRandaoMix*(
+    dag: ChainDAGRef, state: ForkyHashedBeaconState, bid: BlockId,
+    lowSlot: Slot): Opt[Eth2Digest] =
+  ## Compute the requested RANDAO mix for `bid` based on `state`.
+  ## Return `none` if `state` and `bid` do not share a common ancestor
+  ## with slot >= `lowSlot`.
+  let ancestorSlot = ? dag.ancestorSlot(state, bid, lowSlot)
+  doAssert ancestorSlot <= state.data.slot
+  doAssert ancestorSlot <= bid.slot
+
+  # If `blck` is post merge, RANDAO information is immediately available
+  let
+    bdata = ? dag.getForkedBlock(bid)
+    fullMix = dag.computeRandaoMix(bdata)
+  if fullMix.isSome:
+    return fullMix
+
+  # RANDAO mix has to be recomputed from `bid` and `state`
+  var mix {.noinit.}: Eth2Digest
   proc mixToAncestor(highBid: BlockId): Opt[void] =
     ## Mix in/out RANDAO reveals back to `ancestorSlot`
     var bid = highBid
@@ -1431,31 +1428,83 @@ proc computeAttesterRandaoMix(
       bid = ? dag.parent(bid)
     ok()
 
-  # Mix in RANDAO from `blck`
+  # Mix in RANDAO from `bid`
-  if ancestorSlot < dependentBid.slot:
+  if ancestorSlot < bid.slot:
-    withBlck(dependentBdata):
+    withBlck(bdata):
       mix = eth2digest(blck.message.body.randao_reveal.toRaw())
-    ? mixToAncestor(? dag.parent(dependentBid))
+    ? mixToAncestor(? dag.parent(bid))
   else:
     mix.reset()
 
   # Mix in RANDAO from `state`
   let ancestorEpoch = ancestorSlot.epoch
-  if ancestorEpoch + EPOCHS_PER_HISTORICAL_VECTOR <= stateSlot.epoch:
+  if ancestorEpoch + EPOCHS_PER_HISTORICAL_VECTOR <= state.data.slot.epoch:
-    return Opt.none(AttesterRandaoMix)
+    return Opt.none(Eth2Digest)
   let mixRoot = state.dependent_root(ancestorEpoch + 1)
   if mixRoot.isZero:
-    return Opt.none(AttesterRandaoMix)
+    return Opt.none(Eth2Digest)
   ? mixToAncestor(? dag.getBlockId(mixRoot))
   mix.data.mxor state.data.get_randao_mix(ancestorEpoch).data
 
-  ok (dependentBid: dependentBid, mix: mix)
+  ok mix
 
-proc computeShufflingRefFromState*(
+proc computeRandaoMixFromMemory*(
+    dag: ChainDAGRef, bid: BlockId, lowSlot: Slot): Opt[Eth2Digest] =
+  ## Compute requested RANDAO mix for `bid` from available states (~5 ms).
+  template tryWithState(state: ForkedHashedBeaconState) =
+    block:
+      withState(state):
+        let mix = dag.computeRandaoMix(forkyState, bid, lowSlot)
+        if mix.isSome:
+          return mix
+  tryWithState dag.headState
+  tryWithState dag.epochRefState
+  tryWithState dag.clearanceState
+
+proc computeRandaoMixFromDatabase*(
+    dag: ChainDAGRef, bid: BlockId, lowSlot: Slot): Opt[Eth2Digest] =
+  ## Compute requested RANDAO mix for `bid` using closest DB state (~500 ms).
+  let state = newClone(dag.headState)
+  ? dag.getNearbyState(state, bid, lowSlot)
+  withState(state[]):
+    dag.computeRandaoMix(forkyState, bid, lowSlot)
+
+proc computeRandaoMix(
+    dag: ChainDAGRef, bid: BlockId, lowSlot: Slot): Opt[Eth2Digest] =
+  # Try to compute from states available in memory
+  let mix = dag.computeRandaoMixFromMemory(bid, lowSlot)
+  if mix.isSome:
+    return mix
+
+  # Fall back to database
+  dag.computeRandaoMixFromDatabase(bid, lowSlot)
+
+proc computeRandaoMix*(dag: ChainDAGRef, bid: BlockId): Opt[Eth2Digest] =
+  ## Compute requested RANDAO mix for `bid`.
+  const maxSlotDistance = SLOTS_PER_HISTORICAL_ROOT
+  let lowSlot = max(bid.slot, maxSlotDistance.Slot) - maxSlotDistance
+  dag.computeRandaoMix(bid, lowSlot)
+
+proc lowSlotForAttesterShuffling*(epoch: Epoch): Slot =
+  ## Return minimum slot that a state must share ancestry with a block history
+  ## so that RANDAO at `epoch.attester_dependent_slot` can be computed.
+
+  # A state must be somewhat recent so that `get_active_validator_indices`
+  # for the queried `epoch` cannot be affected by any such skipped processing.
+  const numDelayEpochs = compute_activation_exit_epoch(GENESIS_EPOCH).uint64
+  let lowEpoch = max(epoch, (numDelayEpochs - 1).Epoch) - (numDelayEpochs - 1)
+  lowEpoch.start_slot
+
+proc computeShufflingRef*(
     dag: ChainDAGRef, state: ForkyHashedBeaconState,
     blck: BlockRef, epoch: Epoch): Opt[ShufflingRef] =
-  let (dependentBid, mix) =
+  ## Compute `ShufflingRef` for `blck@epoch` based on `state`.
-    ? dag.computeAttesterRandaoMix(state, blck, epoch)
+  ## If `state` has unviable `get_active_validator_indices`, return `none`.
+
+  let
+    dependentBid = (? dag.atSlot(blck.bid, epoch.attester_dependent_slot)).bid
+    lowSlot = epoch.lowSlotForAttesterShuffling
+    mix = ? dag.computeRandaoMix(state, dependentBid, lowSlot)
 
   return ok ShufflingRef(
     epoch: epoch,
@@ -1465,12 +1514,11 @@ proc computeShufflingRefFromState*(
 
 proc computeShufflingRefFromMemory*(
     dag: ChainDAGRef, blck: BlockRef, epoch: Epoch): Opt[ShufflingRef] =
-  ## Compute `ShufflingRef` from states available in memory (up to ~5 ms)
+  ## Compute `ShufflingRef` from available states (~5 ms).
   template tryWithState(state: ForkedHashedBeaconState) =
     block:
       withState(state):
-        let shufflingRef =
+        let shufflingRef = dag.computeShufflingRef(forkyState, blck, epoch)
-          dag.computeShufflingRefFromState(forkyState, blck, epoch)
         if shufflingRef.isOk:
           return shufflingRef
   tryWithState dag.headState
@@ -1479,35 +1527,15 @@ proc computeShufflingRefFromMemory*(
 
 proc computeShufflingRefFromDatabase*(
     dag: ChainDAGRef, blck: BlockRef, epoch: Epoch): Opt[ShufflingRef] =
-  ## Load state from DB, for when DAG states are unviable (up to ~500 ms)
+  ## Compute `ShufflingRef` for `blck@epoch` using closest DB state (~500 ms).
-  let
+  let state = newClone(dag.headState)
-    dependentSlot = epoch.attester_dependent_slot
+  ? dag.getNearbyState(state, blck.bid, epoch.lowSlotForAttesterShuffling)
-    state = newClone(dag.headState)
+  withState(state[]):
-  var
+    dag.computeShufflingRef(forkyState, blck, epoch)
-    e = dependentSlot.epoch
-    b = blck
-  while e > GENESIS_EPOCH and compute_activation_exit_epoch(e) > epoch:
-    let boundaryBlockSlot = e.start_slot - 1
-    b = b.get_ancestor(boundaryBlockSlot)  # nil if < finalized head
-    let
-      bid =
-        if b != nil:
-          b.bid
-        else:
-          let bsi = ? dag.getBlockIdAtSlot(boundaryBlockSlot)
-          bsi.bid
-      bsi = BlockSlotId.init(bid, boundaryBlockSlot + 1)
-    if not dag.getState(bsi, state[]):
-      dec e
-      continue
 
-    return withState(state[]):
+proc computeShufflingRef(
-      dag.computeShufflingRefFromState(forkyState, blck, epoch)
-  err()
-
-proc computeShufflingRef*(
     dag: ChainDAGRef, blck: BlockRef, epoch: Epoch): Opt[ShufflingRef] =
-  # Try to compute `ShufflingRef` from states available in memory
+  # Try to compute from states available in memory
   let shufflingRef = dag.computeShufflingRefFromMemory(blck, epoch)
   if shufflingRef.isOk:
     return shufflingRef

tests/test_blockchain_dag.nim

@@ -1568,7 +1568,7 @@ template runShufflingTests(cfg: RuntimeConfig, numRandomTests: int) =
     ## Check that computed shuffling matches the one from `EpochRef`.
     block:
       let computedShufflingRef = computedShufflingRefParam
-      if computedShufflingRef.isOk:
+      if computedShufflingRef.isSome:
         check computedShufflingRef.get[] == epochRef.get.shufflingRef[]
 
   test "Accelerated shuffling computation":
@@ -1583,6 +1583,14 @@ template runShufflingTests(cfg: RuntimeConfig, numRandomTests: int) =
       let epochRef = dag.getEpochRef(blck, epoch, true)
       check epochRef.isOk
 
+      let dependentBsi = dag.atSlot(blck.bid, epoch.attester_dependent_slot)
+      check dependentBsi.isSome
+      let
+        memoryMix = dag.computeRandaoMixFromMemory(
+          dependentBsi.get.bid, epoch.lowSlotForAttesterShuffling)
+        databaseMix = dag.computeRandaoMixFromDatabase(
+          dependentBsi.get.bid, epoch.lowSlotForAttesterShuffling)
+
       # If shuffling is computable from DAG, check its correctness
       epochRef.checkShuffling dag.computeShufflingRefFromMemory(blck, epoch)
 
@@ -1593,18 +1601,32 @@ template runShufflingTests(cfg: RuntimeConfig, numRandomTests: int) =
       for state in states:
         withState(state[]):
           let
-            shufflingRef =
-              dag.computeShufflingRefFromState(forkyState, blck, epoch)
             stateEpoch = forkyState.data.get_current_epoch
             blckEpoch = blck.bid.slot.epoch
             minEpoch = min(stateEpoch, blckEpoch)
+            lowSlot = epoch.lowSlotForAttesterShuffling
+            shufflingRef = dag.computeShufflingRef(forkyState, blck, epoch)
+            mix = dag.computeRandaoMix(forkyState,
+              dependentBsi.get.bid, epoch.lowSlotForAttesterShuffling)
           if compute_activation_exit_epoch(minEpoch) <= epoch or
-              dag.ancestorSlotForAttesterShuffling(
+              dag.ancestorSlot(
-                forkyState, blck, epoch).isNone:
+                forkyState, dependentBsi.get.bid,
-            check shufflingRef.isErr
+                epoch.lowSlotForAttesterShuffling).isNone:
+            check:
+              shufflingRef.isNone
+              mix.isNone
           else:
-            check shufflingRef.isOk
+            check shufflingRef.isSome
             epochRef.checkShuffling shufflingRef
+            check:
+              mix.isSome
+              memoryMix.isNone or mix == memoryMix
+              databaseMix.isNone or mix == databaseMix
+            epochRef.checkShuffling Opt.some ShufflingRef(
+              epoch: epoch,
+              attester_dependent_root: dependentBsi.get.bid.root,
+              shuffled_active_validator_indices: forkyState.data
+                .get_shuffled_active_validator_indices(epoch, mix.get))
 
   test "Accelerated shuffling computation (with epochRefState jump)":
     # Test cases where `epochRefState` is set to a very old block