llamafolio · 0xpeluche · Mar 8, 2024 · Mar 8, 2024
diff --git a/src/adapters/balancer/arbitrum/index.ts b/src/adapters/balancer/arbitrum/index.ts
@@ -1,4 +1,4 @@
-import { getBalancerBalances } from '@adapters/balancer/common/balance'
+import { getAllBalancerBalances } from '@adapters/balancer/common/balance'
 import type { BaseContext, Contract, GetBalancesHandler } from '@lib/adapter'
 import { resolveBalances } from '@lib/balance'
 
@@ -22,7 +22,7 @@ export const getContracts = async (ctx: BaseContext) => {
 
 export const getBalances: GetBalancesHandler<typeof getContracts> = async (ctx, contracts) => {
   const balances = await resolveBalances<typeof getContracts>(ctx, contracts, {
-    pools: (...args) => getBalancerBalances(...args, vault),
+    pools: (...args) => getAllBalancerBalances(...args, vault),
   })
 
   return {

diff --git a/src/adapters/balancer/avalanche/index.ts b/src/adapters/balancer/avalanche/index.ts
@@ -1,4 +1,4 @@
-import { getBalancerBalances } from '@adapters/balancer/common/balance'
+import { getAllBalancerBalances } from '@adapters/balancer/common/balance'
 import type { BaseContext, Contract, GetBalancesHandler } from '@lib/adapter'
 import { resolveBalances } from '@lib/balance'
 
@@ -22,7 +22,7 @@ export const getContracts = async (ctx: BaseContext) => {
 
 export const getBalances: GetBalancesHandler<typeof getContracts> = async (ctx, contracts) => {
   const balances = await resolveBalances<typeof getContracts>(ctx, contracts, {
-    pools: (...args) => getBalancerBalances(...args, vault),
+    pools: (...args) => getAllBalancerBalances(...args, vault),
   })
 
   return {

diff --git a/src/adapters/balancer/common/balance.ts b/src/adapters/balancer/common/balance.ts
@@ -1,6 +1,34 @@
 import type { Balance, BalancesContext, Contract } from '@lib/adapter'
+import { groupBy, mapSuccess } from '@lib/array'
 import { getUnderlyingsBalancesFromBalancer, type IBalancerBalance } from '@lib/balancer/underlying'
 import { getBalancesOf } from '@lib/erc20'
+import { multicall, type Call } from '@lib/multicall'
+import { isNotNullish } from '@lib/type'
+
+const abi = {
+  claimable_tokens: {
+    stateMutability: 'nonpayable',
+    type: 'function',
+    name: 'claimable_tokens',
+    inputs: [{ name: 'addr', type: 'address' }],
+    outputs: [{ name: '', type: 'uint256' }],
+  },
+  claimable_reward: {
+    stateMutability: 'view',
+    type: 'function',
+    name: 'claimable_reward',
+    inputs: [
+      { name: '_user', type: 'address' },
+      { name: '_reward_token', type: 'address' },
+    ],
+    outputs: [{ name: '', type: 'uint256' }],
+  },
+} as const
+
+export async function getAllBalancerBalances(ctx: BalancesContext, pools: Contract[], vault: Contract) {
+  const { lp, farm } = groupBy(await getBalancerBalances(ctx, pools, vault), 'category')
+  return [...lp, ...(await getBalancerRewards(ctx, farm))]
+}
 
 export async function getBalancerBalances(ctx: BalancesContext, pools: Contract[], vault: Contract) {
   const [lpBalances, farmBalances]: Balance[][] = await Promise.all([
@@ -25,3 +53,59 @@ export async function getBalancerBalances(ctx: BalancesContext, pools: Contract[
     getCategory: (balance: Balance) => balance.category,
   })
 }
+
+async function getBalancerRewards(ctx: BalancesContext, pools: Contract[]): Promise<Balance[]> {
+  const rewardCalls: Call<typeof abi.claimable_reward>[] = pools.flatMap(
+    (pool, poolIndex) =>
+      pool.rewards?.map(
+        (rewardToken, rewardTokenIndex) =>
+          ({
+            target: pool.gauge,
+            params: [ctx.address, (rewardToken as Contract).address],
+            meta: { poolIndex, rewardTokenIndex },
+          }) as any,
+      ),
+  )
+
+  const [pendingBals, extraRewardsBalances] = await Promise.all([
+    multicall({
+      ctx,
+      calls: pools.map((pool) => ({ target: pool.gauge, params: [ctx.address] }) as const),
+      abi: abi.claimable_tokens,
+    }),
+    multicall({
+      ctx,
+      calls: rewardCalls,
+      abi: abi.claimable_reward,
+    }),
+  ])
+
+  const rewardsByPool = extraRewardsBalances.reduce((acc: any, result) => {
+    const { meta } = result.input as any
+    const { poolIndex, rewardTokenIndex } = meta
+
+    if (!acc[poolIndex]) {
+      acc[poolIndex] = { ...pools[poolIndex], rewards: [...(pools[poolIndex].rewards as Contract[])] }
+    }
+
+    acc[poolIndex].rewards[rewardTokenIndex] = {
+      ...acc[poolIndex].rewards[rewardTokenIndex],
+      amount: result.output,
+    }
+
+    return acc
+  }, {})
+
+  return mapSuccess(pendingBals, (res, index) => {
+    const pool = Object.values(rewardsByPool)[index] as Contract
+    if (!pool || !pool.rewards) return null
+
+    const rewards = pool.rewards as Contract[]
+
+    if (rewards.length > 0 && rewards[0].amount === 0n) {
+      rewards[0] = { ...rewards[0], amount: res.output }
+    }
+
+    return { ...pool, rewards }
+  }).filter(isNotNullish)
+}
diff --git a/src/adapters/balancer/common/pool.ts b/src/adapters/balancer/common/pool.ts
@@ -40,8 +40,37 @@ const abi = {
     inputs: [],
     outputs: [{ name: '', type: 'address' }],
   },
+  bal_token: {
+    stateMutability: 'view',
+    type: 'function',
+    name: 'bal_token',
+    inputs: [],
+    outputs: [{ name: '', type: 'address' }],
+  },
+  reward_count: {
+    stateMutability: 'view',
+    type: 'function',
+    name: 'reward_count',
+    inputs: [],
+    outputs: [{ name: '', type: 'uint256' }],
+  },
+  reward_tokens: {
+    stateMutability: 'view',
+    type: 'function',
+    name: 'reward_tokens',
+    inputs: [{ name: 'arg0', type: 'uint256' }],
+    outputs: [{ name: '', type: 'address' }],
+  },
 } as const
 
+const BAL: { [key: string]: `0x${string}` } = {
+  arbitrum: '0x040d1edc9569d4bab2d15287dc5a4f10f56a56b8',
+  avalanche: '0xA2A035Dd93b0e963864FA14A240401d6CeAc5558',
+  ethereum: '0xba100000625a3754423978a60c9317c58a424e3d',
+  gnosis: '0x7ef541e2a22058048904fe5744f9c7e4c57af717',
+  polygon: '0x9a71012b13ca4d3d0cdc72a177df3ef03b0e76a3',
+}
+
 export async function getBalancerPools(ctx: BaseContext, url: string, gaugeController?: Contract): Promise<Contract[]> {
   const contracts: Contract[] = []
 
@@ -77,9 +106,11 @@ export async function getBalancerPools(ctx: BaseContext, url: string, gaugeContr
   }
 
   if (ctx.chain === 'ethereum' && gaugeController) {
-    return getBalancerEthGauges(ctx, gaugeController, contracts)
+    const balancerEthGauges = await getBalancerEthGauges(ctx, gaugeController, contracts)
+    return getBalancerRewards(ctx, balancerEthGauges)
   } else {
-    return getBalancerChildGauges(ctx, contracts)
+    const balancerChildGauges = await getBalancerChildGauges(ctx, contracts)
+    return getBalancerRewards(ctx, balancerChildGauges)
   }
 }
 
@@ -160,3 +191,46 @@ async function getBalancerChildGauges(ctx: BaseContext, pools: Contract[]): Prom
 
   return mergedPools
 }
+
+async function getBalancerRewards(ctx: BaseContext, pools: Contract[]): Promise<Contract[]> {
+  const poolsWithoutGauges = pools.filter((pool) => pool.gauge === undefined)
+  const poolsWithGauges = pools.filter((pool) => pool.gauge !== undefined)
+
+  const rewardsLengths = await multicall({
+    ctx,
+    calls: poolsWithGauges.map((pool) => ({ target: pool.gauge }) as const),
+    abi: abi.reward_count,
+  })
+
+  const rewardsTokensCalls = rewardsLengths
+    .map((res, poolIndex) => {
+      return {
+        pool: poolsWithGauges[poolIndex],
+        calls: rangeBI(0n, res.output as bigint).map(
+          (i) =>
+            ({
+              target: res.input.target,
+              params: [i],
+              poolIndex,
+            }) as const,
+        ),
+      }
+    })
+    .flat()
+
+  const rewardsTokens = await multicall({
+    ctx,
+    calls: rewardsTokensCalls.map((c) => c.calls).flat(),
+    abi: abi.reward_tokens,
+  })
+
+  const poolWithRewards: Contract[] = poolsWithGauges.map((pool, index) => {
+    const tokensForThisPool = rewardsTokens
+      .filter((token) => token.input.poolIndex === index)
+      .map((token) => token.output)
+
+    return { ...pool, rewards: [BAL[ctx.chain], ...tokensForThisPool] as `0x${string}`[] }
+  })
+
+  return [...poolsWithoutGauges, ...poolWithRewards]
+}
diff --git a/src/adapters/balancer/ethereum/index.ts b/src/adapters/balancer/ethereum/index.ts
@@ -1,4 +1,4 @@
-import { getBalancerBalances } from '@adapters/balancer/common/balance'
+import { getAllBalancerBalances } from '@adapters/balancer/common/balance'
 import { getLockerBalances } from '@adapters/balancer/ethereum/locker'
 import type { BaseContext, Contract, GetBalancesHandler } from '@lib/adapter'
 import { resolveBalances } from '@lib/balance'
@@ -66,7 +66,7 @@ export const getContracts = async (ctx: BaseContext) => {
 
 export const getBalances: GetBalancesHandler<typeof getContracts> = async (ctx, contracts) => {
   const balances = await resolveBalances<typeof getContracts>(ctx, contracts, {
-    pools: (...args) => getBalancerBalances(...args, vault),
+    pools: (...args) => getAllBalancerBalances(...args, vault),
     oldPools: (ctx, oldPools) => getPoolsBalances(ctx, oldPools, { getPoolAddress: (pool) => pool.address }),
     votingEscrow: (...args) => getLockerBalances(...args, vault),
   })

diff --git a/src/adapters/balancer/gnosis/index.ts b/src/adapters/balancer/gnosis/index.ts
@@ -1,4 +1,4 @@
-import { getBalancerBalances } from '@adapters/balancer/common/balance'
+import { getAllBalancerBalances } from '@adapters/balancer/common/balance'
 import type { BaseContext, Contract, GetBalancesHandler } from '@lib/adapter'
 import { resolveBalances } from '@lib/balance'
 
@@ -22,7 +22,7 @@ export const getContracts = async (ctx: BaseContext) => {
 
 export const getBalances: GetBalancesHandler<typeof getContracts> = async (ctx, contracts) => {
   const balances = await resolveBalances<typeof getContracts>(ctx, contracts, {
-    pools: (...args) => getBalancerBalances(...args, vault),
+    pools: (...args) => getAllBalancerBalances(...args, vault),
   })
 
   return {

diff --git a/src/adapters/balancer/polygon/index.ts b/src/adapters/balancer/polygon/index.ts
@@ -1,4 +1,4 @@
-import { getBalancerBalances } from '@adapters/balancer/common/balance'
+import { getAllBalancerBalances } from '@adapters/balancer/common/balance'
 import type { BaseContext, Contract, GetBalancesHandler } from '@lib/adapter'
 import { resolveBalances } from '@lib/balance'
 
@@ -22,7 +22,7 @@ export const getContracts = async (ctx: BaseContext) => {
 
 export const getBalances: GetBalancesHandler<typeof getContracts> = async (ctx, contracts) => {
   const balances = await resolveBalances<typeof getContracts>(ctx, contracts, {
-    pools: (...args) => getBalancerBalances(...args, vault),
+    pools: (...args) => getAllBalancerBalances(...args, vault),
   })
 
   return {