feat: rewards distribution helper (#59)

* feat: added integer rewards distribution helper with tests

* feat: ci job for typescript tests

* feat: print out rune allocation per account as a table

* fix: order logs by block number and timestamp, get latest rune adddress by account

* chore: cleanup

* chore: rename typescript ci job

.github/workflows/typescript.yml

@@ -0,0 +1,41 @@
+name: Rewards distribution
+
+on:
+  workflow_dispatch:
+  push:
+    branches: [main, develop]
+  pull_request:
+    branches: [main, develop]
+
+jobs:
+  check:
+    strategy:
+      fail-fast: true
+
+    name: Typescript scripts
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Setup Node.js
+        id: setup-node
+        uses: actions/setup-node@v4
+        with:
+          node-version-file: .nvmrc
+          cache: yarn
+
+      - name: Install Dependencies
+        id: install
+        run: |
+          yarn
+          cd scripts/rewards-distribution
+          yarn
+
+      - name: Run Typescript tests
+        id: test
+        run: |
+          cd scripts/rewards-distribution
+          yarn test

scripts/rewards-distribution/calculateRewards/calculateRewards.ts

@@ -1,14 +1,14 @@
 import { Address, Block } from "viem";
 import { RFoxLog, StakeLog, UnstakeLog } from "../events";
-import { getStakingAmount, isLogType } from "../helpers";
+import { isLogType } from "../helpers";
 import { REWARD_RATE, WAD } from "../constants";
 import assert from "assert";
 
-type StakingInfo = {
+export type StakingInfo = {
   stakingBalance: bigint;
   earnedRewards: bigint;
   rewardPerTokenStored: bigint;
-  runeAddress: String;
+  runeAddress: string;
 };
 
 const getEmptyStakingInfo = () => {
@@ -83,7 +83,7 @@ const updateReward = (
 
 const stake = (
   amount: bigint,
-  runeAddress: String,
+  runeAddress: string,
   stakingInfo: StakingInfo,
   rewardPerTokenStored: bigint,
   totalStaked: bigint,
@@ -143,14 +143,14 @@ export const calculateRewards = (
   // the previous epoch. This prevents us missing rewards for the first block in the epoch.
   previousEpochEndBlock: Block,
   epochEndBlock: Block,
-  logs: { log: RFoxLog; timestamp: bigint }[],
+  orderedLogs: { log: RFoxLog; timestamp: bigint }[],
 ) => {
   let totalStaked = 0n;
   let rewardPerTokenStored = 0n;
   let lastUpdateTimestamp = contractCreationBlock.timestamp;
   const stakingInfoByAccount: Record<Address, StakingInfo> = {};
 
-  const stakingLogs = logs.filter(
+  const stakingLogs = orderedLogs.filter(
     (
       logWithTimestamp,
     ): logWithTimestamp is { log: StakeLog | UnstakeLog; timestamp: bigint } =>

scripts/rewards-distribution/distributeAmount/distributeAmount.test.ts

@@ -0,0 +1,152 @@
+import { distributeAmount } from "./distributeAmount";
+
+describe("distributeAmount", () => {
+  test("basic distribution", () => {
+    const totalRuneAmountToDistroBaseUnit = 100n;
+    const earnedRewardsByAccount = {
+      "0x1": 50n,
+      "0x2": 50n,
+    };
+    const expectedOutput = {
+      "0x1": 50n,
+      "0x2": 50n,
+    };
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+
+  test("uneven distribution", () => {
+    const totalRuneAmountToDistroBaseUnit = 100n;
+    const earnedRewardsByAccount = {
+      "0x1": 60n,
+      "0x2": 40n,
+    };
+    const expectedOutput = {
+      "0x1": 60n,
+      "0x2": 40n,
+    };
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+
+  test("remainder distribution", () => {
+    const totalRuneAmountToDistroBaseUnit = 101n;
+    const earnedRewardsByAccount = {
+      "0x1": 50n,
+      "0x2": 50n,
+    };
+    const output = distributeAmount(
+      totalRuneAmountToDistroBaseUnit,
+      earnedRewardsByAccount,
+    );
+    const totalDistributed = Object.values(output).reduce(
+      (sum, value) => sum + value,
+      0n,
+    );
+    expect(totalDistributed).toBe(101n);
+    expect(output["0x1"] + output["0x2"]).toBe(101n);
+  });
+
+  test("zero total distribution", () => {
+    const totalRuneAmountToDistroBaseUnit = 0n;
+    const earnedRewardsByAccount = {
+      "0x1": 50n,
+      "0x2": 50n,
+    };
+    const expectedOutput = {
+      "0x1": 0n,
+      "0x2": 0n,
+    };
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+
+  test("single account", () => {
+    const totalRuneAmountToDistroBaseUnit = 100n;
+    const earnedRewardsByAccount = {
+      "0x1": 100n,
+    };
+    const expectedOutput = {
+      "0x1": 100n,
+    };
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+
+  test("no earned rewards", () => {
+    const totalRuneAmountToDistroBaseUnit = 100n;
+    const earnedRewardsByAccount = {};
+    const expectedOutput = {};
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+
+  test("large number", () => {
+    const totalRuneAmountToDistroBaseUnit =
+      1000000000000000000000000000000000000000000000000000000n;
+    const earnedRewardsByAccount = {
+      "0x1": 999999999999999999999999999999999999999999999999999999n,
+      "0x2": 1n,
+    };
+    const expectedOutput = {
+      "0x1": 999999999999999999999999999999999999999999999999999999n,
+      "0x2": 1n,
+    };
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+
+  test("negative rewards test (invalid input)", () => {
+    const totalRuneAmountToDistroBaseUnit = 100n;
+    const earnedRewardsByAccount = {
+      "0x1": -50n,
+      "0x2": 50n,
+    };
+    expect(() =>
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toThrow();
+  });
+
+  test("remainder distribution test with many accounts", () => {
+    const totalRuneAmountToDistroBaseUnit = 1003n;
+    const earnedRewardsByAccount = {
+      "0x1": 10n,
+      "0x2": 20n,
+      "0x3": 30n,
+      "0x4": 40n,
+    };
+    const output = distributeAmount(
+      totalRuneAmountToDistroBaseUnit,
+      earnedRewardsByAccount,
+    );
+    const totalDistributed = Object.values(output).reduce(
+      (sum, value) => sum + value,
+      0n,
+    );
+    expect(totalDistributed).toBe(1003n);
+    expect(output["0x1"] + output["0x2"] + output["0x3"] + output["0x4"]).toBe(
+      1003n,
+    );
+  });
+
+  test("edge case test with zero rewards for some accounts", () => {
+    const totalRuneAmountToDistroBaseUnit = 100n;
+    const earnedRewardsByAccount = {
+      "0x1": 50n,
+      "0x2": 0n,
+    };
+    const expectedOutput = {
+      "0x1": 100n,
+      "0x2": 0n,
+    };
+    expect(
+      distributeAmount(totalRuneAmountToDistroBaseUnit, earnedRewardsByAccount),
+    ).toEqual(expectedOutput);
+  });
+});

scripts/rewards-distribution/distributeAmount/distributeAmount.ts

@@ -0,0 +1,75 @@
+import BigNumber from "bignumber.js";
+import { Address } from "viem";
+
+// Distributes a total amount of RUNE to a set of accounts based on their earned rewards.
+// If there is a remainder of RUNE base units, the remainder is distributed to the accounts
+// with the largest proportion of the total rewards.
+export const distributeAmount = (
+  totalRuneAmountToDistroBaseUnit: bigint,
+  earnedRewardsByAccount: Record<Address, bigint>,
+) => {
+  // Set the precision to a high value to avoid rounding errors
+  BigNumber.config({ DECIMAL_PLACES: 100 });
+
+  const totalEarnedRewards = Object.values(earnedRewardsByAccount).reduce(
+    (sum, earnedRewards) => sum + earnedRewards,
+    0n,
+  );
+
+  const earnedRewardsByAccountArray = Object.entries(
+    earnedRewardsByAccount,
+  ).map(([account, earnedRewards]) => {
+    const proportionOfRewards = BigNumber(earnedRewards.toString()).div(
+      totalEarnedRewards.toString(),
+    );
+    return { account: account as Address, earnedRewards, proportionOfRewards };
+  });
+
+  const runeAllocationBaseUnitByAccount: Record<Address, bigint> = {};
+
+  // Calculate each user's share ignoring remainder (we'll add it later)
+  for (const { account, proportionOfRewards } of earnedRewardsByAccountArray) {
+    // Calculate the integer allocation ignoring remainder.
+    // Rounds towards nearest neighbor. If equidistant, rounds towards -Infinity.
+    runeAllocationBaseUnitByAccount[account as Address] = BigInt(
+      proportionOfRewards
+        .times(totalRuneAmountToDistroBaseUnit.toString())
+        .toFixed(0, BigNumber.ROUND_HALF_FLOOR),
+    );
+  }
+
+  // Calculate the allocated amount so far (to determine remainder)
+  const totalAllocationRuneBaseUnitBeforeRemainder = Object.values(
+    runeAllocationBaseUnitByAccount,
+  ).reduce((sum, runeAllocationBaseUnit) => sum + runeAllocationBaseUnit, 0n);
+
+  // Determine the remainder
+  let remainderRuneAmountBaseUnit =
+    totalRuneAmountToDistroBaseUnit -
+    totalAllocationRuneBaseUnitBeforeRemainder;
+
+  // If there's no remaining amount, return the distribution by account
+  if (remainderRuneAmountBaseUnit === 0n) {
+    return runeAllocationBaseUnitByAccount;
+  }
+
+  // Sort the accounts by their proportion of the total rewards in descending order
+  // so the accounts with the largest reward get preference for the remainder
+  earnedRewardsByAccountArray.sort((a, b) => {
+    return b.proportionOfRewards.minus(a.proportionOfRewards).toNumber();
+  });
+
+  // Distribute the remainder one base unit at a time until everything is distributed
+  let i = 0;
+  while (remainderRuneAmountBaseUnit !== 0n) {
+    if (!earnedRewardsByAccountArray[i]) break;
+
+    const account = earnedRewardsByAccountArray[i].account;
+    const baseUnitToAdd: bigint = remainderRuneAmountBaseUnit > 0n ? 1n : -1n;
+    runeAllocationBaseUnitByAccount[account] += baseUnitToAdd;
+    remainderRuneAmountBaseUnit -= baseUnitToAdd;
+    i = (i + 1) % earnedRewardsByAccountArray.length;
+  }
+
+  return runeAllocationBaseUnitByAccount;
+};

scripts/rewards-distribution/getLatestRuneAddressByAccount.ts

@@ -0,0 +1,24 @@
+import { Address } from "viem";
+import { RFoxLog, StakeLog, SetRuneAddressLog } from "./events";
+import { isLogType } from "./helpers";
+
+export const getLatestRuneAddressByAccount = (
+  orderedLogs: { log: RFoxLog; timestamp: bigint }[],
+) => {
+  const runeAddressByAccount: Record<Address, string> = {};
+
+  for (const { log } of orderedLogs) {
+    if (isLogType("Stake", log)) {
+      const stakeLog = log as StakeLog;
+      runeAddressByAccount[stakeLog.args.account] = stakeLog.args.runeAddress;
+    }
+
+    if (isLogType("SetRuneAddress", log)) {
+      const setRuneAddressLog = log as SetRuneAddressLog;
+      runeAddressByAccount[setRuneAddressLog.args.account] =
+        setRuneAddressLog.args.newRuneAddress;
+    }
+  }
+
+  return runeAddressByAccount;
+};

scripts/rewards-distribution/helpers.ts

@@ -3,10 +3,10 @@ import {
   ARBITRUM_RFOX_PROXY_CONTRACT_ADDRESS,
   GET_LOGS_BLOCK_STEP_SIZE,
 } from "./constants";
-import { AbiEvent, Block, Log, PublicClient } from "viem";
+import { AbiEvent, Log, PublicClient } from "viem";
 import cliProgress from "cli-progress";
 import colors from "ansi-colors";
-import { RFoxEvent, RFoxLog, StakeLog, UnstakeLog } from "./events";
+import { RFoxLog, StakeLog, UnstakeLog } from "./events";
 import { stakingV1Abi } from "./generated/abi-types";
 
 // we cache promises to prevent async race conditions hydrating the cache