fix(staking): account for unstakes pro-rata in stake position rewards

src/views/yield-dtf/staking/atoms.ts

@@ -8,8 +8,14 @@ import {
   rTokenAtom,
   rTokenConfigurationAtom,
   rTokenStateAtom,
+  stRsrBalanceAtom,
   walletAtom,
 } from '@/state/atoms'
+import {
+  AccountStakeRecord,
+  calculateStakeLots,
+  calculateStakeRewards,
+} from './stake-accounting'
 
 export const unstakeDelayAtom = atom((get) => {
   const params = get(rTokenConfigurationAtom)
@@ -55,7 +61,7 @@ export const pendingRSRSummaryAtom = atom<{
   )
 })
 
-const accountStakeHistoryAtom = atomWithLoadable(async (get) => {
+const accountStakeLotsAtom = atomWithLoadable(async (get) => {
   const gqlClient = get(gqlClientAtom)
   const wallet = get(walletAtom)
   const rToken = get(rTokenAtom)
@@ -67,7 +73,12 @@ const accountStakeHistoryAtom = atomWithLoadable(async (get) => {
   const request: any = await gqlClient.request(
     gql`
       query getAccountStakeHistory($id: String!) {
-        accountStakeRecords(orderBy: blockNumber, where: { account: $id }) {
+        accountStakeRecords(
+          first: 1000
+          orderBy: blockNumber
+          orderDirection: asc
+          where: { account: $id }
+        ) {
           exchangeRate
           amount
           rsrAmount
@@ -82,83 +93,22 @@ const accountStakeHistoryAtom = atomWithLoadable(async (get) => {
     return null
   }
 
-  let stakes: [number, number, number][] = []
-  let totalRewardBalance = 0
-
-  for (const record of request.accountStakeRecords as {
-    exchangeRate: string
-    amount: string
-    rsrAmount: bigint
-    isStake: string | boolean
-  }[]) {
-    const recordAmount = Number(record.amount)
-    const recordExchangeRate = Number(record.exchangeRate)
-
-    if (record.isStake === 'true' || record.isStake === true) {
-      stakes.push([recordAmount, recordExchangeRate, Number(record.rsrAmount)])
-    } else {
-      let stakesRewarded = 0
-      let unstake = recordAmount
-
-      // Calculate current stake rewards
-      for (let i = 0; i < stakes.length; i++) {
-        const [stakeAmount, stakeExchangeRate, stakeRsrAmount] = stakes[i]
-        // Calculate rewards from this stake and keep going
-        const snapshotRsrAmount =
-          Math.min(unstake, stakeAmount) * stakeExchangeRate
-        const currentRsrAmount =
-          Math.min(unstake, stakeAmount) * recordExchangeRate
-
-        // Count rewards
-        totalRewardBalance += currentRsrAmount - snapshotRsrAmount
-
-        if (stakeAmount > unstake) {
-          stakes[i] = [
-            stakeAmount - unstake,
-            stakeExchangeRate,
-            stakeRsrAmount - snapshotRsrAmount,
-          ]
-          break
-        } else if (stakeAmount === unstake) {
-          stakesRewarded++
-          break
-        } else {
-          // Continue counting rewards
-          unstake = unstake - stakeAmount
-          stakesRewarded++
-        }
-      }
-      // Remove accrued stakes
-      stakes = stakes.slice(stakesRewarded)
-    }
-  }
-
-  return {
-    stakes,
-    totalRewardBalance,
-  }
+  return calculateStakeLots(request.accountStakeRecords as AccountStakeRecord[])
 })
 
 const exchangeRateAtom = atom((get) => get(rTokenStateAtom).exchangeRate)
 
 // TODO: Check re-renders on exchangeRateUpdate improve memo
 export const accountCurrentPositionAtom = atom((get) => {
-  const stakeHistory = get(accountStakeHistoryAtom)
+  const lots = get(accountStakeLotsAtom)
   const exchangeRate = get(exchangeRateAtom)
+  const stRsrBalance = get(stRsrBalanceAtom)
 
-  let stBalance = 0
-  let rsrBalance = 0
-
-  if (!stakeHistory) {
+  if (!lots) {
     return 0
   }
 
-  for (const [stakeAmount, _, stakeRsrAmount] of stakeHistory.stakes) {
-    stBalance += stakeAmount
-    rsrBalance += stakeRsrAmount
-  }
-
-  return stBalance * exchangeRate - rsrBalance
+  return calculateStakeRewards(lots, exchangeRate, +stRsrBalance.balance)
 })
 
 export const rateAtom = atom((get) => {

src/views/yield-dtf/staking/components/stake-position.tsx

@@ -16,12 +16,7 @@ const StakePosition = ({ className }: StakePositionProps) => {
   const rate = useAtomValue(rateAtom)
   const balance = useAtomValue(stRsrBalanceAtom)
   const rsrPrice = useAtomValue(rsrPriceAtom)
-  let rewards = useAtomValue(accountCurrentPositionAtom)
-
-  // Prevent the case when the user withdraws and rewards get stuck awaiting for subgraph
-  if (!balance.value && rewards) {
-    rewards = 0
-  }
+  const rewards = useAtomValue(accountCurrentPositionAtom)
 
   return (
     <div className={cn(className)}>

src/views/yield-dtf/staking/stake-accounting.ts

@@ -0,0 +1,87 @@
+export interface AccountStakeRecord {
+  exchangeRate: string
+  amount: string
+  rsrAmount: string
+  isStake: string | boolean
+}
+
+// Open stake, in stRSR, together with the RSR that was paid for it
+export interface StakeLot {
+  amount: number
+  rsrAmount: number
+}
+
+const isStakeRecord = (record: AccountStakeRecord) =>
+  record.isStake === true || record.isStake === 'true'
+
+/**
+ * stRSR is fungible, so an unstake burns a proportional slice of every open
+ * stake rather than a specific one: unstaking x% of the position removes x%
+ * of both its stRSR and its RSR cost basis, leaving the remaining rewards at
+ * (100 - x)% of what they were.
+ */
+export const calculateStakeLots = (
+  records: AccountStakeRecord[]
+): StakeLot[] => {
+  let lots: StakeLot[] = []
+
+  for (const record of records) {
+    const amount = Number(record.amount)
+
+    if (isStakeRecord(record)) {
+      lots.push({ amount, rsrAmount: Number(record.rsrAmount) })
+      continue
+    }
+
+    const staked = lots.reduce((total, lot) => total + lot.amount, 0)
+
+    if (!staked) {
+      continue
+    }
+
+    const remainingShare = Math.max(1 - amount / staked, 0)
+
+    if (!remainingShare) {
+      lots = []
+      continue
+    }
+
+    lots = lots.map(({ amount, rsrAmount }) => ({
+      amount: amount * remainingShare,
+      rsrAmount: rsrAmount * remainingShare,
+    }))
+  }
+
+  return lots
+}
+
+/**
+ * RSR accrued by the open position: what it is worth now minus what it cost.
+ *
+ * Stake records go missing from the subgraph — accounts that unstaked
+ * everything can still have open stakes according to their records — so the
+ * reconstruction is trusted only up to the stRSR the account actually holds.
+ * Rewards are scaled by the share of the reconstructed position still held,
+ * and stRSR the records cannot account for accrues nothing.
+ */
+export const calculateStakeRewards = (
+  lots: StakeLot[],
+  exchangeRate: number,
+  stRsrBalance: number
+) => {
+  const staked = lots.reduce((total, lot) => total + lot.amount, 0)
+
+  if (!staked) {
+    return 0
+  }
+
+  const heldShare = Math.min(stRsrBalance / staked, 1)
+
+  return (
+    lots.reduce(
+      (rewards, { amount, rsrAmount }) =>
+        rewards + amount * exchangeRate - rsrAmount,
+      0
+    ) * heldShare
+  )
+}