fix: align timestamps for sample values (#16849)

Please read [contributing guidelines](CONTRIBUTING.md) and remove this
line.

For audit-related pull requests, please use the [audit PR
template](?expand=1&template=audit.md).

Author: LHerskind

l1-contracts/src/core/Rollup.sol

@@ -387,6 +387,14 @@ contract Rollup is IStaking, IValidatorSelection, IRollup, RollupCore {
     return ValidatorOperationsExtLib.getSampleSeedAt(getEpochAt(_ts));
   }
 
+  function getSamplingSizeAt(Timestamp _ts) external view override(IValidatorSelection) returns (uint256) {
+    return ValidatorOperationsExtLib.getSamplingSizeAt(getEpochAt(_ts));
+  }
+
+  function getLagInEpochs() external view override(IValidatorSelection) returns (uint256) {
+    return ValidatorOperationsExtLib.getLagInEpochs();
+  }
+
   /**
    * @notice  Get the sample seed for the current epoch
    *

l1-contracts/src/core/RollupCore.sol

@@ -272,7 +272,7 @@ contract RollupCore is EIP712("Aztec Rollup", "1"), Ownable, IStakingCore, IVali
     StakingLib.initialize(
       _stakingAsset, _gse, exitDelay, address(slasher), _config.stakingQueueConfig, _config.localEjectionThreshold
     );
-    ValidatorOperationsExtLib.initializeValidatorSelection(_config.targetCommitteeSize);
+    ValidatorOperationsExtLib.initializeValidatorSelection(_config.targetCommitteeSize, _config.lagInEpochs);
 
     // If no booster is specifically provided, deploy one.
     if (address(_config.rewardConfig.booster) == address(0)) {

l1-contracts/src/core/interfaces/IRollup.sol

@@ -57,6 +57,7 @@ struct RollupConfigInput {
   uint256 aztecSlotDuration;
   uint256 aztecEpochDuration;
   uint256 targetCommitteeSize;
+  uint256 lagInEpochs;
   uint256 aztecProofSubmissionEpochs;
   uint256 slashingQuorum;
   uint256 slashingRoundSize;

l1-contracts/src/core/interfaces/IValidatorSelection.sol

@@ -11,7 +11,8 @@ struct ValidatorSelectionStorage {
   mapping(Epoch => bytes32 committeeCommitment) committeeCommitments;
   // Checkpointed map of epoch -> randao value
   Checkpoints.Trace224 randaos;
-  uint256 targetCommitteeSize;
+  uint32 targetCommitteeSize;
+  uint32 lagInEpochs;
 }
 
 interface IValidatorSelectionCore {
@@ -36,6 +37,8 @@ interface IValidatorSelection is IValidatorSelectionCore, IEmperor {
   function getTimestampForSlot(Slot _slotNumber) external view returns (Timestamp);
 
   function getSampleSeedAt(Timestamp _ts) external view returns (uint256);
+  function getSamplingSizeAt(Timestamp _ts) external view returns (uint256);
+  function getLagInEpochs() external view returns (uint256);
   function getCurrentSampleSeed() external view returns (uint256);
 
   function getEpochAt(Timestamp _ts) external view returns (Epoch);

l1-contracts/src/core/libraries/rollup/ValidatorOperationsExtLib.sol

@@ -66,8 +66,8 @@ library ValidatorOperationsExtLib {
     StakingLib.finalizeWithdraw(_attester);
   }
 
-  function initializeValidatorSelection(uint256 _targetCommitteeSize) external {
-    ValidatorSelectionLib.initialize(_targetCommitteeSize);
+  function initializeValidatorSelection(uint256 _targetCommitteeSize, uint256 _lagInEpochs) external {
+    ValidatorSelectionLib.initialize(_targetCommitteeSize, _lagInEpochs);
   }
 
   function setupEpoch() external {
@@ -125,6 +125,14 @@ library ValidatorOperationsExtLib {
     return ValidatorSelectionLib.getSampleSeed(_epoch);
   }
 
+  function getSamplingSizeAt(Epoch _epoch) external view returns (uint256) {
+    return ValidatorSelectionLib.getSamplingSize(_epoch);
+  }
+
+  function getLagInEpochs() external view returns (uint256) {
+    return ValidatorSelectionLib.getLagInEpochs();
+  }
+
   function getTargetCommitteeSize() external view returns (uint256) {
     return ValidatorSelectionLib.getStorage().targetCommitteeSize;
   }

l1-contracts/src/core/libraries/rollup/ValidatorSelectionLib.sol

@@ -130,12 +130,12 @@ library ValidatorSelectionLib {
    *      The first two epochs use maximum seed values for startup.
    * @param _targetCommitteeSize The desired number of validators in each epoch's committee
    */
-  function initialize(uint256 _targetCommitteeSize) internal {
+  function initialize(uint256 _targetCommitteeSize, uint256 _lagInEpochs) internal {
     ValidatorSelectionStorage storage store = getStorage();
-    store.targetCommitteeSize = _targetCommitteeSize;
+    store.targetCommitteeSize = _targetCommitteeSize.toUint32();
+    store.lagInEpochs = _lagInEpochs.toUint32();
 
-    // Set the initial randao
-    store.randaos.push(0, uint224(block.prevrandao));
+    checkpointRandao(Epoch.wrap(0));
   }
 
   /**
@@ -453,33 +453,22 @@ library ValidatorSelectionLib {
 
   /**
    * @notice Checkpoints randao value for future usage
-   * @dev Checks if already stored before computing and storing the randao value.
-   *      Offset the epoch by 2 to maintain the two-epoch advance requirement.
-   *      This ensures randomness are set well in advance to prevent manipulation.
-   * @param _epoch The current epoch (randao will be set for _epoch + 2)
-   *       Passed to reduce recomputation
+   * @dev Checks if already stored before storing the randao value.
+   * @param _epoch The current epoch
    */
   function checkpointRandao(Epoch _epoch) internal {
     ValidatorSelectionStorage storage store = getStorage();
 
-    // Compute the offset
-    uint32 epoch = Epoch.unwrap(_epoch).toUint32() + 2;
-
     // Check if the latest checkpoint is for the next epoch
     // It should be impossible that zero epoch snapshots exist, as in the genesis state we push the first values
     // into the store
-    (, uint32 mostRecentEpoch,) = store.randaos.latestCheckpoint();
-
-    // If the randao for the next epoch is already set, we can skip the computation
-    if (mostRecentEpoch == epoch) {
-      return;
-    }
+    (, uint32 mostRecentTs,) = store.randaos.latestCheckpoint();
+    uint32 ts = Timestamp.unwrap(_epoch.toTimestamp()).toUint32();
 
     // If the most recently stored epoch is less than the epoch we are querying, then we need to store randao for
-    // later use
-    if (mostRecentEpoch < epoch) {
-      // Truncate the randao to be used for future sampling.
-      store.randaos.push(epoch, uint224(block.prevrandao));
+    // later use. We truncate to save storage costs.
+    if (mostRecentTs < ts) {
+      store.randaos.push(ts, uint224(block.prevrandao));
     }
   }
 
@@ -539,22 +528,16 @@ library ValidatorSelectionLib {
    * @notice Converts an epoch number to the timestamp used for validator set sampling
    * @dev Calculates the sampling timestamp by:
    *      1. Taking the epoch start timestamp
-   *      2. Subtracting one full epoch duration to ensure stability
-   *      3. Subtracting 1 second to get end-of-block state
+   *      2. Subtracting `lagInEpochs` full epoch duration to ensure stability
    *
    *      This ensures validator set sampling uses stable historical data that won't be
    *      affected by last-minute changes or L1 reorgs during synchronization.
    * @param _epoch The epoch to calculate sampling time for
    * @return The Unix timestamp (uint32) to use for validator set sampling
    */
   function epochToSampleTime(Epoch _epoch) internal view returns (uint32) {
-    // We do -1, as the snapshots practically happen at the end of the block, e.g.,
-    // a tx manipulating the set in at $t$ would be visible already at lookup $t$ if after that
-    // transactions. But reading at $t-1$ would be the state at the end of $t-1$ which is the state
-    // as we "start" time $t$. We then shift that back by an entire L2 epoch to guarantee
-    // we are not hit by last-minute changes or L1 reorgs when syncing validators from our clients.
-
-    return Timestamp.unwrap(_epoch.toTimestamp()).toUint32() - uint32(TimeLib.getEpochDurationInSeconds()) - 1;
+    uint32 sub = getStorage().lagInEpochs * TimeLib.getEpochDurationInSeconds().toUint32();
+    return Timestamp.unwrap(_epoch.toTimestamp()).toUint32() - sub;
   }
 
   /**
@@ -566,7 +549,17 @@ library ValidatorSelectionLib {
    */
   function getSampleSeed(Epoch _epoch) internal view returns (uint256) {
     ValidatorSelectionStorage storage store = getStorage();
-    return uint256(keccak256(abi.encode(_epoch, store.randaos.upperLookup(Epoch.unwrap(_epoch).toUint32()))));
+    uint32 ts = epochToSampleTime(_epoch);
+    return uint256(keccak256(abi.encode(_epoch, store.randaos.upperLookup(ts))));
+  }
+
+  function getSamplingSize(Epoch _epoch) internal view returns (uint256) {
+    uint32 ts = epochToSampleTime(_epoch);
+    return StakingLib.getAttesterCountAtTime(Timestamp.wrap(ts));
+  }
+
+  function getLagInEpochs() internal view returns (uint256) {
+    return getStorage().lagInEpochs;
   }
 
   /**

l1-contracts/test/harnesses/TestConstants.sol

@@ -18,6 +18,7 @@ library TestConstants {
   uint256 internal constant AZTEC_SLOT_DURATION = 36;
   uint256 internal constant AZTEC_EPOCH_DURATION = 32;
   uint256 internal constant AZTEC_TARGET_COMMITTEE_SIZE = 48;
+  uint256 internal constant AZTEC_LAG_IN_EPOCHS = 2;
   uint256 internal constant AZTEC_PROOF_SUBMISSION_EPOCHS = 1;
   uint256 internal constant AZTEC_SLASHING_QUORUM = 6;
   uint256 internal constant AZTEC_SLASHING_ROUND_SIZE = 10;
@@ -95,6 +96,7 @@ library TestConstants {
       aztecEpochDuration: AZTEC_EPOCH_DURATION,
       aztecProofSubmissionEpochs: AZTEC_PROOF_SUBMISSION_EPOCHS,
       targetCommitteeSize: AZTEC_TARGET_COMMITTEE_SIZE,
+      lagInEpochs: AZTEC_LAG_IN_EPOCHS,
       slashingQuorum: AZTEC_SLASHING_QUORUM,
       slashingRoundSize: AZTEC_SLASHING_ROUND_SIZE,
       slashingLifetimeInRounds: AZTEC_SLASHING_LIFETIME_IN_ROUNDS,

l1-contracts/test/staking/TimeCheater.sol

@@ -42,6 +42,10 @@ contract TimeCheater {
     return Epoch.wrap(currentSlot / epochDuration);
   }
 
+  function getCurrentSlot() public view returns (uint256) {
+    return currentSlot;
+  }
+
   function cheat__setTimeStorage(TimeStorage memory _timeStorage) public {
     vm.store(
       target,

l1-contracts/test/validator-selection/seedAndSizeSnapshots.t.sol

@@ -0,0 +1,78 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright 2025 Aztec Labs.
+pragma solidity >=0.8.27;
+
+import {ValidatorSelectionTestBase, CheatDepositArgs} from "./ValidatorSelectionBase.sol";
+import {Epoch, Timestamp, TimeLib} from "@aztec/core/libraries/TimeLib.sol";
+import {Checkpoints} from "@oz/utils/structs/Checkpoints.sol";
+import {IValidatorSelection} from "@aztec/core/interfaces/IValidatorSelection.sol";
+import {TestConstants} from "../harnesses/TestConstants.sol";
+import {BN254Lib, G1Point, G2Point} from "@aztec/shared/libraries/BN254Lib.sol";
+import {GSE} from "@aztec/governance/GSE.sol";
+
+contract SeedAndSizeSnapshotsTest is ValidatorSelectionTestBase {
+  using TimeLib for Timestamp;
+  using TimeLib for Epoch;
+
+  uint256 internal $currentsize = 4;
+  mapping(uint256 slot => uint256 size) internal $sizes;
+  mapping(uint256 slot => uint256 randao) internal $randaos;
+
+  function test_seedAndSizeSnapshots() public setup(4, 4) {
+    // We set up the initial
+    $sizes[timeCheater.getCurrentSlot()] = $currentsize;
+    $randaos[timeCheater.getCurrentSlot()] = block.prevrandao;
+
+    uint256 endEpoch = Epoch.unwrap(timeCheater.getCurrentEpoch()) + 10;
+
+    GSE gse = rollup.getGSE();
+
+    vm.prank(address(testERC20.owner()));
+    testERC20.mint(address(rollup), type(uint128).max);
+    vm.prank(address(rollup));
+    testERC20.approve(address(gse), type(uint128).max);
+
+    while (Epoch.unwrap(timeCheater.getCurrentEpoch()) < endEpoch) {
+      timeCheater.cheat__progressSlot();
+
+      uint256 nextRandao = uint256(keccak256(abi.encode(block.prevrandao)));
+      vm.prevrandao(nextRandao);
+
+      rollup.checkpointRandao();
+
+      vm.prank(address(rollup));
+      gse.deposit(
+        address(uint160(nextRandao)),
+        address(uint160(nextRandao)),
+        BN254Lib.g1Zero(),
+        BN254Lib.g2Zero(),
+        BN254Lib.g1Zero(),
+        true
+      );
+
+      $currentsize += 1;
+      $sizes[timeCheater.getCurrentSlot()] = $currentsize;
+      $randaos[timeCheater.getCurrentSlot()] = nextRandao;
+
+      // We will add one node to the GSE for rollup (impersonate rollup to avoid the queue).
+      // We want to see that the lag between current values are the same between randaos and size values
+
+      uint256 epochIndex = Epoch.unwrap(timeCheater.getCurrentEpoch());
+      if (epochIndex >= 2) {
+        (uint256 seed, uint256 size) = getValues();
+        uint256 slot = (epochIndex - 2) * timeCheater.epochDuration();
+
+        assertEq(size, $sizes[slot], "invalid size");
+        assertEq(seed, uint256(keccak256(abi.encode(epochIndex, uint224($randaos[slot])))), "invalid seed");
+      }
+    }
+  }
+
+  function getValues() internal view returns (uint256 sampleSeed, uint256 size) {
+    // We are always using for an epoch, so we are going to do that here as well
+    Timestamp ts = Timestamp.wrap(block.timestamp);
+
+    sampleSeed = rollup.getSampleSeedAt(ts);
+    size = rollup.getSamplingSizeAt(ts);
+  }
+}

spartan/environments/scenario.local.env

@@ -18,6 +18,7 @@ AZTEC_EJECTION_THRESHOLD=50000000000000000000
 AZTEC_SLASH_AMOUNT_SMALL=5000000000000000000
 AZTEC_SLASH_AMOUNT_MEDIUM=10000000000000000000
 AZTEC_SLASH_AMOUNT_LARGE=15000000000000000000
+AZTEC_LAG_IN_EPOCHS=0
 
 # The following need to be set manually
 # AZTEC_DOCKER_IMAGE=aztecprotocol/aztec:whatever