Update unclaimedAmounts in handleRewardsClaimed

src/entities/merkleDistributor.ts

@@ -68,6 +68,10 @@ export function createOrLoadDistributorClaim(user: Address): DistributorClaim {
   return claim
 }
 
+export function loadDistributorClaim(user: Address): DistributorClaim | null {
+  return DistributorClaim.load(user.toHex())
+}
+
 export function getDistributionType(distData: Bytes): DistributionType {
   // only passed addresses are currently supported
   const distAddress = Address.fromBytes(distData)

src/mappings/merkleDistributor.ts

@@ -1,15 +1,4 @@
-import {
-  Address,
-  BigInt,
-  Bytes,
-  ethereum,
-  ipfs,
-  json,
-  JSONValue,
-  JSONValueKind,
-  log,
-  store,
-} from '@graphprotocol/graph-ts'
+import { Address, BigInt, Bytes, ethereum, ipfs, json, JSONValue, JSONValueKind, log } from '@graphprotocol/graph-ts'
 import { DistributorClaimedAmount, PeriodicDistribution } from '../../generated/schema'
 import {
   OneTimeDistributionAdded,
@@ -27,6 +16,7 @@ import {
   DistributionType,
   getDistributionType,
   loadDistributor,
+  loadDistributorClaim,
 } from '../entities/merkleDistributor'
 import { createTransaction } from '../entities/transaction'
 import { loadUniswapPool } from '../entities/uniswap'
@@ -280,6 +270,12 @@ export function handleRewardsClaimed(event: RewardsClaimed): void {
   const user = event.params.account
   const tokens = event.params.tokens
   const cumulativeAmounts = event.params.cumulativeAmounts
+
+  // The DistributorClaim object is guaranteed to exist at this point because
+  // the user has claimed rewards using the DistributorClaim object.
+  const claim = loadDistributorClaim(user)!
+  const unclaimedAmounts = claim.unclaimedAmounts
+
   for (let i = 0; i < tokens.length; i++) {
     const claimedAmountId = `${tokens[i].toHex()}-${user.toHex()}`
     let claimedAmount = DistributorClaimedAmount.load(claimedAmountId)
@@ -289,8 +285,11 @@ export function handleRewardsClaimed(event: RewardsClaimed): void {
     }
     claimedAmount.cumulativeClaimedAmount = cumulativeAmounts[i]
     claimedAmount.save()
+    unclaimedAmounts[i] = BigInt.zero()
   }
-  store.remove('DistributorClaim', user.toHex())
+
+  claim.unclaimedAmounts = unclaimedAmounts
+  claim.save()
 
   createTransaction(event.transaction.hash.toHex())
   log.info('[MerkleDistributor] RewardsClaimed user={}', [user.toHex()])