From 5e0e2e578f0da2f2e58abba6f12c2d7e725cae66 Mon Sep 17 00:00:00 2001
From: alecks <0xalecks@gmail.com>
Date: Fri, 5 Apr 2024 14:38:18 -0400
Subject: [PATCH] ui: fix stalk balance bug

---
 projects/ui/src/components/Silo/Overview.tsx | 39 ++++----
 projects/ui/src/state/farmer/silo/updater.ts | 96 ++++++++++++++++----
 2 files changed, 99 insertions(+), 36 deletions(-)

diff --git a/projects/ui/src/components/Silo/Overview.tsx b/projects/ui/src/components/Silo/Overview.tsx
index 40baccf506..9e6f0fd005 100644
--- a/projects/ui/src/components/Silo/Overview.tsx
+++ b/projects/ui/src/components/Silo/Overview.tsx
@@ -7,11 +7,7 @@ import { AppState } from '~/state';
 import useTabs from '~/hooks/display/useTabs';
 import TokenIcon from '~/components/Common/TokenIcon';
 import { SEEDS, STALK } from '~/constants/tokens';
-import {
-  displayPercentage,
-  displayStalk,
-  displayUSD,
-} from '~/util';
+import { displayPercentage, displayStalk, displayUSD } from '~/util';
 import { ChipLabel, StyledTab } from '~/components/Common/Tabs';
 import { ZERO_BN } from '~/constants';
 import Row from '~/components/Common/Row';
@@ -78,7 +74,10 @@ const Overview: FC<{
   const migrationNeeded = useMigrationNeeded();
   const siloBalance = useFarmerSiloBalances();
   //
-  const [tab, handleChange] = useTabs(migrationNeeded ? SLUGS : altSLUGS, 'view');
+  const [tab, handleChange] = useTabs(
+    migrationNeeded ? SLUGS : altSLUGS,
+    'view'
+  );
 
   //
   const ownership =
@@ -86,13 +85,15 @@ const Overview: FC<{
       ? farmerSilo.stalk.active.div(beanstalkSilo.stalk.total)
       : ZERO_BN;
 
-  const deposits = Object.values(siloBalance).map(token => token.deposited.crates).flat(Infinity)
+  const deposits = Object.values(siloBalance)
+    .map((token) => token.deposited.crates)
+    .flat(Infinity);
 
-  let totalStalkGrown = farmerSilo.stalk.grown;
+  const totalStalkGrown = farmerSilo.stalk.grown;
 
-  deposits.forEach((deposit: any) => {
-    totalStalkGrown = totalStalkGrown.plus(deposit.stalk.grown)
-  })
+  // deposits.forEach((deposit: any) => {
+  //   totalStalkGrown = totalStalkGrown.plus(deposit.stalk.grown)
+  // })
 
   const stalkStats = useCallback(
     (s: BigNumber, v: BigNumber[], d: string) => (
@@ -118,9 +119,7 @@ const Overview: FC<{
         <Stat
           title="Total Stalk Grown"
           titleTooltip="The total number of Mown and Mowable Grown Stalk your Deposits have accrued."
-          amount={displayStalk(
-            totalStalkGrown
-          )}
+          amount={displayStalk(totalStalkGrown)}
           color="text.primary"
           gap={0.25}
           sx={{ minWidth: 120, ml: 0 }}
@@ -177,7 +176,9 @@ const Overview: FC<{
           <MigrateTab />
         </Box>
       )}
-      <Box sx={{ display: tab === (migrationNeeded ? 1 : 0) ? 'block' : 'none' }}>
+      <Box
+        sx={{ display: tab === (migrationNeeded ? 1 : 0) ? 'block' : 'none' }}
+      >
         <OverviewPlot
           label="Silo Deposits"
           account={account}
@@ -199,7 +200,9 @@ const Overview: FC<{
           empty={breakdown.states.deposited.value.eq(0)}
         />
       </Box>
-      <Box sx={{ display: tab === (migrationNeeded ? 2 : 1) ? 'block' : 'none' }}>
+      <Box
+        sx={{ display: tab === (migrationNeeded ? 2 : 1) ? 'block' : 'none' }}
+      >
         <OverviewPlot
           label="Stalk Ownership"
           account={account}
@@ -229,7 +232,9 @@ const Overview: FC<{
           empty={farmerSilo.stalk.total.lte(0)}
         />
       </Box>
-      <Box sx={{ display: tab === (migrationNeeded ? 3 : 2) ? 'block' : 'none' }}>
+      <Box
+        sx={{ display: tab === (migrationNeeded ? 3 : 2) ? 'block' : 'none' }}
+      >
         <OverviewPlot
           label="Seeds Ownership"
           account={account}
diff --git a/projects/ui/src/state/farmer/silo/updater.ts b/projects/ui/src/state/farmer/silo/updater.ts
index 296da9d4d1..044c9e7508 100644
--- a/projects/ui/src/state/farmer/silo/updater.ts
+++ b/projects/ui/src/state/farmer/silo/updater.ts
@@ -1,13 +1,17 @@
 import { useCallback, useEffect } from 'react';
 import { useDispatch } from 'react-redux';
+import BigNumber from 'bignumber.js';
 import axios from 'axios';
 import { Token, TokenValue } from '@beanstalk/sdk';
 import { ethers } from 'ethers';
-import { ZERO_BN } from '~/constants';
+import { TokenMap, ZERO_BN } from '~/constants';
 import { useBeanstalkContract } from '~/hooks/ledger/useContract';
 import useChainId from '~/hooks/chain/useChainId';
 import { bigNumberResult, transform } from '~/util';
 import useAccount from '~/hooks/ledger/useAccount';
+import useSdk from '~/hooks/sdk';
+import { LegacyDepositCrate } from '~/state/farmer/silo';
+import useSeason from '~/hooks/beanstalk/useSeason';
 import {
   resetFarmerSilo,
   updateLegacyFarmerSiloBalances,
@@ -19,9 +23,6 @@ import {
   updateFarmerSiloError,
   updateFarmerSiloRan,
 } from './actions';
-import useSdk from '~/hooks/sdk';
-import { LegacyDepositCrate } from '~/state/farmer/silo';
-import useSeason from '~/hooks/beanstalk/useSeason';
 
 type SiloV3StaticData = {
   deposits: {
@@ -37,6 +38,13 @@ type SiloV3StaticData = {
   };
 };
 
+type BaseToGrownStalk = {
+  base: BigNumber;
+  grown: BigNumber;
+  seeds: BigNumber;
+  unclaimed: BigNumber;
+};
+
 export const fetchMigrationData = async (account: string) =>
   axios
     .get(
@@ -145,8 +153,6 @@ export const useFetchFarmerSilo = () => {
           sdk.silo.getSeeds(account),
         ]);
 
-        console.log('Fetched migration data', balances);
-
         // Pre-migration, # of seeds is calc'd from the contract getter
         activeSeedBalance = _activeSeedBalance;
 
@@ -264,9 +270,59 @@ export const useFetchFarmerSilo = () => {
         dispatch(updateFarmerSiloBalanceSdk(balances));
       }
 
-      /// earnedStalk (this is already included in activeStalk)
-      /// earnedSeed  (aka plantable seeds)
-      /// these work because 1 BEAN = 1 BDV.
+      /**
+       * We need to calculate the stalk for un-migrated accounts differently than migrated ones
+       */
+
+      // First aggregate all crates per token
+      const stalkPerTokenForUnMigrated = Object.entries(payload).reduce<
+        TokenMap<BaseToGrownStalk>
+      >((prev, [tokenAddress, tokenBalances]) => {
+        if (!season) return prev;
+        prev[tokenAddress] =
+          tokenBalances.deposited!.crates.reduce<BaseToGrownStalk>(
+            (acc, crate) => {
+              acc.base = acc.base.plus(crate.stalk.base);
+              acc.grown = acc.grown.plus(crate.stalk.grown);
+              acc.seeds = acc.seeds.plus(crate.seeds);
+              acc.unclaimed = ZERO_BN;
+              return acc;
+            },
+            {
+              base: ZERO_BN,
+              grown: ZERO_BN,
+              unclaimed: ZERO_BN,
+              seeds: ZERO_BN,
+            }
+          );
+        return prev;
+      }, {});
+
+      // Then aggregate all tokens
+      const stalkForUnMigrated = Object.entries(
+        stalkPerTokenForUnMigrated
+      ).reduce(
+        (prev, [_, data]) => {
+          prev.base = prev.base.plus(data.base);
+          prev.grown = prev.grown.plus(data.grown);
+
+          return prev;
+        },
+        {
+          base: ZERO_BN,
+          grown: ZERO_BN,
+          earned: transform(
+            sdk.tokens.BEAN.getStalk(earnedBeanBalance),
+            'bnjs'
+          ),
+          total: ZERO_BN,
+        }
+      );
+      stalkForUnMigrated.total = stalkForUnMigrated.base
+        .plus(stalkForUnMigrated.grown)
+        .plus(stalkForUnMigrated.earned);
+      // End of un-migrated stalk calculation
+
       const earnedStalkBalance = sdk.tokens.BEAN.getStalk(earnedBeanBalance);
       const earnedSeedBalance = sdk.tokens.BEAN.getSeeds(earnedBeanBalance);
       const totalStalkBalance = activeStalkBalance.add(grownStalkBalance);
@@ -281,10 +337,18 @@ export const useFetchFarmerSilo = () => {
           earned: transform(earnedBeanBalance, 'bnjs', sdk.tokens.BEAN),
         },
         stalk: {
-          active: transform(activeStalkBalance, 'bnjs', sdk.tokens.STALK),
-          earned: transform(earnedStalkBalance, 'bnjs', sdk.tokens.STALK),
-          grown: transform(grownStalkBalance, 'bnjs', sdk.tokens.STALK),
-          total: transform(totalStalkBalance, 'bnjs', sdk.tokens.STALK),
+          active: migrationNeeded
+            ? stalkForUnMigrated.base
+            : transform(activeStalkBalance, 'bnjs', sdk.tokens.STALK),
+          earned: migrationNeeded
+            ? stalkForUnMigrated.earned
+            : transform(earnedStalkBalance, 'bnjs', sdk.tokens.STALK),
+          grown: migrationNeeded
+            ? stalkForUnMigrated.grown
+            : transform(grownStalkBalance, 'bnjs', sdk.tokens.STALK),
+          total: migrationNeeded
+            ? stalkForUnMigrated.total
+            : transform(totalStalkBalance, 'bnjs', sdk.tokens.STALK),
           grownByToken: grownStalkByToken,
         },
         seeds: {
@@ -297,12 +361,6 @@ export const useFetchFarmerSilo = () => {
         },
       };
 
-      // console.log("Silo Rewards", rewards, {
-      //   totalStalkBalance: totalStalkBalance.toHuman(),
-      //   grownStalkBalance: grownStalkBalance.toHuman(),
-      //   earnedBeanBalance: earnedBeanBalance.toHuman(),
-      // })
-
       dispatch(updateLegacyFarmerSiloRewards(rewards));
 
       // HEADS UP: this has to be called after updateLegacyFarmerSiloRewards