From f836634ddd8d84f6c1606f687b86cab524e0d34f Mon Sep 17 00:00:00 2001
From: "Eugene Y. Q. Shen" <eyqs@plumenetwork.xyz>
Date: Wed, 25 Sep 2024 12:54:42 -0700
Subject: [PATCH] [NES-200] accrue yield in the FakeComponentToken (#34)

---
 nest/script/DeployNestContracts.s.sol |  4 +--
 nest/src/FakeComponentToken.sol       | 46 ++++++++++++++++++++-------
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/nest/script/DeployNestContracts.s.sol b/nest/script/DeployNestContracts.s.sol
index 29530b9..f403e4b 100644
--- a/nest/script/DeployNestContracts.s.sol
+++ b/nest/script/DeployNestContracts.s.sol
@@ -15,12 +15,12 @@ contract DeployNestContracts is Script {
 
     address private constant ARC_ADMIN_ADDRESS = 0x1c9d94FAD4ccCd522804a955103899e0D6A4405a;
     address private constant NEST_ADMIN_ADDRESS = 0xb015762405De8fD24d29A6e0799c12e0Ea81c1Ff;
-    address private constant USDC_ADDRESS = 0x849c25e6cCB03cdc23ba91d92440dA7bC8486be2;
+    address private constant P_ADDRESS = 0xEa0c23A2411729073Ed52fF94b38FceffE82FDE3;
 
     function run() external {
         vm.startBroadcast(ARC_ADMIN_ADDRESS);
 
-        IComponentToken currencyToken = IComponentToken(USDC_ADDRESS);
+        IComponentToken currencyToken = IComponentToken(P_ADDRESS);
 
         FakeComponentToken fakeComponentToken = new FakeComponentToken();
         FakeComponentTokenProxy fakeComponentTokenProxy = new FakeComponentTokenProxy(
diff --git a/nest/src/FakeComponentToken.sol b/nest/src/FakeComponentToken.sol
index 27d08ac..f283a4f 100644
--- a/nest/src/FakeComponentToken.sol
+++ b/nest/src/FakeComponentToken.sol
@@ -33,6 +33,14 @@ contract FakeComponentToken is
         uint8 decimals;
         /// @dev Version of the FakeComponentToken
         uint256 version;
+        /// @dev Total amount of yield that has ever been accrued by all users
+        uint256 totalYieldAccrued;
+        /// @dev Total amount of yield that has ever been withdrawn by all users
+        uint256 totalYieldWithdrawn;
+        /// @dev Total amount of yield that has ever been accrued by each user
+        mapping(address user => uint256 currencyTokenAmount) yieldAccrued;
+        /// @dev Total amount of yield that has ever been withdrawn by each user
+        mapping(address user => uint256 currencyTokenAmount) yieldWithdrawn;
     }
 
     // keccak256(abi.encode(uint256(keccak256("plume.storage.FakeComponentToken")) - 1)) & ~bytes32(uint256(0xff))
@@ -202,8 +210,22 @@ contract FakeComponentToken is
      * @param user Address of the user for which to claim yield
      */
     function claimYield(address user) external returns (uint256 amount) {
+        FakeComponentTokenStorage storage $ = _getFakeComponentTokenStorage();
         amount = unclaimedYield(user);
-        _getFakeComponentTokenStorage().currencyToken.transfer(user, amount);
+        $.currencyToken.transfer(user, amount);
+        $.yieldWithdrawn[user] += amount;
+        $.totalYieldWithdrawn += amount;
+    }
+
+    /**
+     * @notice Accrue yield for the given user
+     * @dev Anyone can call this function to accrue yield for any user
+     * @param user Address of the user for which to accrue yield
+     */
+    function accrueYield(address user, uint256 amount) external {
+        FakeComponentTokenStorage storage $ = _getFakeComponentTokenStorage();
+        $.yieldAccrued[user] += amount;
+        $.totalYieldAccrued += amount;
     }
 
     // Admin Setter Functions
@@ -243,18 +265,19 @@ contract FakeComponentToken is
     }
 
     /// @notice Total yield distributed to all FakeComponentTokens for all users
-    function totalYield() public view returns (uint256 amount) {
-        return 6;
+    function totalYield() external view returns (uint256 amount) {
+        return _getFakeComponentTokenStorage().totalYieldAccrued;
     }
 
     /// @notice Claimed yield across all FakeComponentTokens for all users
-    function claimedYield() public view returns (uint256 amount) {
-        return 4;
+    function claimedYield() external view returns (uint256 amount) {
+        return _getFakeComponentTokenStorage().totalYieldWithdrawn;
     }
 
     /// @notice Unclaimed yield across all FakeComponentTokens for all users
     function unclaimedYield() external view returns (uint256 amount) {
-        return totalYield() - claimedYield();
+        FakeComponentTokenStorage storage $ = _getFakeComponentTokenStorage();
+        return $.totalYieldAccrued - $.totalYieldWithdrawn;
     }
 
     /**
@@ -262,8 +285,8 @@ contract FakeComponentToken is
      * @param user Address of the user for which to get the total yield
      * @return amount Total yield distributed to the user
      */
-    function totalYield(address user) public view returns (uint256 amount) {
-        return 3;
+    function totalYield(address user) external view returns (uint256 amount) {
+        return _getFakeComponentTokenStorage().yieldAccrued[user];
     }
 
     /**
@@ -271,8 +294,8 @@ contract FakeComponentToken is
      * @param user Address of the user for which to get the claimed yield
      * @return amount Amount of yield that the user has claimed
      */
-    function claimedYield(address user) public view returns (uint256 amount) {
-        return 2;
+    function claimedYield(address user) external view returns (uint256 amount) {
+        return _getFakeComponentTokenStorage().yieldWithdrawn[user];
     }
 
     /**
@@ -281,7 +304,8 @@ contract FakeComponentToken is
      * @return amount Amount of yield that the user has not yet claimed
      */
     function unclaimedYield(address user) public view returns (uint256 amount) {
-        return totalYield(user) - claimedYield(user);
+        FakeComponentTokenStorage storage $ = _getFakeComponentTokenStorage();
+        return $.yieldAccrued[user] - $.yieldWithdrawn[user];
     }
 
 }