feat: staking rewards accounting with on-chain validation

foundry/src/StakingV1.sol

@@ -177,6 +177,7 @@ contract StakingV1 is
         StakingInfo storage info = stakingInfo[msg.sender];
         info.stakingBalance += amount;
         info.runeAddress = runeAddress;
+
         totalStaked += amount;
 
         emit Stake(msg.sender, amount, runeAddress);
@@ -201,6 +202,7 @@ contract StakingV1 is
         // Set staking / unstaking amounts
         info.stakingBalance -= amount;
         info.unstakingBalance += amount;
+
         totalStaked -= amount;
         totalCoolingDown += amount;
 

scripts/rewards-distribution/calculateRewards/calculateRewards.ts

@@ -1,57 +1,263 @@
-import { Address } from "viem";
+import { Address, Block } from "viem";
 import { RFoxLog, StakeLog, UnstakeLog } from "../events";
 import { getStakingAmount, isLogType } from "../helpers";
+import { REWARD_RATE, WAD } from "../constants";
+import assert from "assert";
 
-export const calculateRewards = (
-  fromBlock: bigint,
-  toBlock: bigint,
-  logsByBlockNumber: Record<string, RFoxLog[]>,
-  epochBlockReward: bigint,
+type StakingInfo = {
+  stakingBalance: bigint;
+  earnedRewards: bigint;
+  rewardPerTokenStored: bigint;
+  runeAddress: String;
+};
+
+const getEmptyStakingInfo = () => {
+  return {
+    stakingBalance: 0n,
+    earnedRewards: 0n,
+    rewardPerTokenStored: 0n,
+    runeAddress: "",
+  };
+};
+
+const rewardPerToken = (
+  rewardPerTokenStored: bigint,
+  totalStaked: bigint,
+  currentTimestamp: bigint,
+  lastUpdateTimestamp: bigint,
+) => {
+  if (totalStaked == 0n) {
+    return rewardPerTokenStored;
+  }
+  return (
+    rewardPerTokenStored +
+    ((currentTimestamp - lastUpdateTimestamp) * REWARD_RATE * WAD) / totalStaked
+  );
+};
+
+const earned = (
+  stakingInfo: StakingInfo,
+  rewardPerTokenStored: bigint,
+  totalStaked: bigint,
+  currentTimestamp: bigint,
+  lastUpdateTimestamp: bigint,
+) => {
+  return (
+    (stakingInfo.stakingBalance *
+      (rewardPerToken(
+        rewardPerTokenStored,
+        totalStaked,
+        currentTimestamp,
+        lastUpdateTimestamp,
+      ) -
+        stakingInfo.rewardPerTokenStored)) /
+      WAD +
+    stakingInfo.earnedRewards
+  );
+};
+
+const updateReward = (
+  stakingInfo: StakingInfo,
+  rewardPerTokenStored: bigint,
   totalStaked: bigint,
+  currentTimestamp: bigint,
+  lastUpdateTimestamp: bigint,
 ) => {
-  const balanceByAccountBaseUnit: Record<Address, bigint> = {};
+  rewardPerTokenStored = rewardPerToken(
+    rewardPerTokenStored,
+    totalStaked,
+    currentTimestamp,
+    lastUpdateTimestamp,
+  );
+  lastUpdateTimestamp = currentTimestamp;
+  stakingInfo.earnedRewards = earned(
+    stakingInfo,
+    rewardPerTokenStored,
+    totalStaked,
+    currentTimestamp,
+    lastUpdateTimestamp,
+  );
+  stakingInfo.rewardPerTokenStored = rewardPerTokenStored;
+  return { stakingInfo, rewardPerTokenStored, lastUpdateTimestamp };
+};
 
-  // this must be initialized to empty
-  const epochRewardByAccount: Record<Address, bigint> = {};
+const stake = (
+  amount: bigint,
+  runeAddress: String,
+  stakingInfo: StakingInfo,
+  rewardPerTokenStored: bigint,
+  totalStaked: bigint,
+  currentTimestamp: bigint,
+  lastUpdateTimestamp: bigint,
+) => {
+  ({ stakingInfo, rewardPerTokenStored, lastUpdateTimestamp } = updateReward(
+    stakingInfo,
+    rewardPerTokenStored,
+    totalStaked,
+    currentTimestamp,
+    lastUpdateTimestamp,
+  ));
 
-  for (let blockNumber = fromBlock; blockNumber <= toBlock; blockNumber++) {
-    const incomingLogs: RFoxLog[] =
-      logsByBlockNumber[blockNumber.toString()] ?? [];
+  stakingInfo.stakingBalance += amount;
+  stakingInfo.runeAddress = runeAddress;
+  totalStaked += amount;
 
-    const stakingLogs = incomingLogs.filter(
-      (log): log is StakeLog | UnstakeLog =>
-        isLogType("Stake", log) || isLogType("Unstake", log),
-    );
+  return {
+    stakingInfo,
+    rewardPerTokenStored,
+    lastUpdateTimestamp,
+    totalStaked,
+  };
+};
 
-    // process logs if there are any
-    for (const log of stakingLogs) {
-      const account = log.args.account;
-      if (!balanceByAccountBaseUnit[account]) {
-        balanceByAccountBaseUnit[account] = 0n;
-      }
-      const stakingAmountBaseUnit = getStakingAmount(log);
-      balanceByAccountBaseUnit[account] += stakingAmountBaseUnit;
-      totalStaked += stakingAmountBaseUnit;
+const unstake = (
+  amount: bigint,
+  stakingInfo: StakingInfo,
+  rewardPerTokenStored: bigint,
+  totalStaked: bigint,
+  currentTimestamp: bigint,
+  lastUpdateTimestamp: bigint,
+) => {
+  ({ stakingInfo, rewardPerTokenStored, lastUpdateTimestamp } = updateReward(
+    stakingInfo,
+    rewardPerTokenStored,
+    totalStaked,
+    currentTimestamp,
+    lastUpdateTimestamp,
+  ));
 
-      // clear empty balances
-      if (balanceByAccountBaseUnit[account] === 0n) {
-        delete balanceByAccountBaseUnit[account];
+  stakingInfo.stakingBalance -= amount;
+  totalStaked -= amount;
+
+  return {
+    stakingInfo,
+    rewardPerTokenStored,
+    lastUpdateTimestamp,
+    totalStaked,
+  };
+};
+
+export const calculateRewards = (
+  contractCreationBlock: Block,
+  // The reward is computed as an all-time value, so we need to subtract the rewards at the end of
+  // the previous epoch. This prevents us missing rewards for the first block in the epoch.
+  previousEpochEndBlock: Block,
+  epochEndBlock: Block,
+  logs: { log: RFoxLog; timestamp: bigint }[],
+) => {
+  let totalStaked = 0n;
+  let rewardPerTokenStored = 0n;
+  let lastUpdateTimestamp = contractCreationBlock.timestamp;
+  const stakingInfoByAccount: Record<Address, StakingInfo> = {};
+
+  const stakingLogs = logs.filter(
+    (
+      logWithTimestamp,
+    ): logWithTimestamp is { log: StakeLog | UnstakeLog; timestamp: bigint } =>
+      isLogType("Stake", logWithTimestamp.log) ||
+      isLogType("Unstake", logWithTimestamp.log),
+  );
+
+  const epochStartRewardsByAccount: Record<Address, bigint> = {};
+  const epochEndRewardsByAccount: Record<Address, bigint> = {};
+
+  const previousEpochEndBlockNumber = previousEpochEndBlock.number;
+  const epochEndBlockNumber = epochEndBlock.number;
+
+  assert(
+    previousEpochEndBlockNumber !== null,
+    "Epoch start block number is null",
+  );
+  assert(epochEndBlockNumber !== null, "Epoch end block number is null");
+
+  let hasCalcedStartRewards = false;
+
+  // process logs
+  for (const { log, timestamp: currentTimestamp } of stakingLogs) {
+    // Paranoia in case we get logs past the end of the epoch
+    assert(log.blockNumber <= epochEndBlockNumber);
+
+    // When the block number passes the start of the epoch, assign the reward values for the start of the epoch
+    if (
+      !hasCalcedStartRewards &&
+      log.blockNumber > previousEpochEndBlockNumber
+    ) {
+      for (const [account, stakingInfo] of Object.entries(
+        stakingInfoByAccount,
+      )) {
+        epochStartRewardsByAccount[account as Address] = earned(
+          stakingInfo,
+          rewardPerTokenStored,
+          totalStaked,
+          previousEpochEndBlock.timestamp,
+          lastUpdateTimestamp,
+        );
+
+        hasCalcedStartRewards = true;
       }
     }
 
-    for (const account of Object.keys(balanceByAccountBaseUnit) as Address[]) {
-      // calculate rewards for the current block
-      const reward =
-        totalStaked > 0n
-          ? (epochBlockReward * balanceByAccountBaseUnit[account]) / totalStaked
-          : 0n;
+    let stakingInfo =
+      stakingInfoByAccount[log.args.account] ?? getEmptyStakingInfo();
 
-      if (epochRewardByAccount[account] == undefined) {
-        epochRewardByAccount[account] = 0n;
-      }
-      epochRewardByAccount[account] += reward;
+    switch (true) {
+      case isLogType<StakeLog>("Stake", log):
+        ({
+          stakingInfo,
+          rewardPerTokenStored,
+          lastUpdateTimestamp,
+          totalStaked,
+        } = stake(
+          log.args.amount,
+          log.args.runeAddress,
+          stakingInfo,
+          rewardPerTokenStored,
+          totalStaked,
+          currentTimestamp,
+          lastUpdateTimestamp,
+        ));
+        break;
+      case isLogType<UnstakeLog>("Unstake", log):
+        ({
+          stakingInfo,
+          rewardPerTokenStored,
+          lastUpdateTimestamp,
+          totalStaked,
+        } = unstake(
+          log.args.amount,
+          stakingInfo,
+          rewardPerTokenStored,
+          totalStaked,
+          currentTimestamp,
+          lastUpdateTimestamp,
+        ));
+        break;
+      default:
+        break;
     }
+
+    stakingInfoByAccount[log.args.account] = stakingInfo;
+  }
+
+  // Grab the reward values for the end of the epoch
+  for (const [account, stakingInfo] of Object.entries(stakingInfoByAccount)) {
+    epochEndRewardsByAccount[account as Address] = earned(
+      stakingInfo,
+      rewardPerTokenStored,
+      totalStaked,
+      epochEndBlock.timestamp,
+      lastUpdateTimestamp,
+    );
+  }
+
+  const earnedRewardsByAccount: Record<Address, bigint> = {};
+
+  for (const [account, epochEndReward] of Object.entries(
+    epochEndRewardsByAccount,
+  )) {
+    earnedRewardsByAccount[account as Address] =
+      epochEndReward - (epochStartRewardsByAccount[account as Address] ?? 0n);
   }
 
-  return epochRewardByAccount;
+  return earnedRewardsByAccount;
 };

scripts/rewards-distribution/constants.ts

@@ -2,9 +2,12 @@ import { Address } from "viem";
 
 export const RUNE_DECIMALS = 8;
 
+export const WAD = 1n * 10n ** 18n;
+export const REWARD_RATE = 1_000_000_000n;
+
 // The number of blocks to query at a time, when fetching logs
 export const GET_LOGS_BLOCK_STEP_SIZE = 20000n;
 
 // RFOX on Arbitrum ERC1967Proxy contract address
 export const ARBITRUM_RFOX_PROXY_CONTRACT_ADDRESS: Address =
-  "0x0c66f315542fdec1d312c415b14eef614b0910ef";
+  "0xd612B64A134f3D4830542B7463CE8ca8a29D7268";