From 6813203202780b9bfd28ba708d9330d7b98d3fb3 Mon Sep 17 00:00:00 2001
From: aalavandhan1984 <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 14:01:35 -0500
Subject: [PATCH 01/10] removed stakeFor function

---
 contracts/TokenGeyser.sol | 33 ++++-----------------------------
 1 file changed, 4 insertions(+), 29 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 1f7008b..0de0e33 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -148,32 +148,7 @@ contract TokenGeyser is IStaking, Ownable {
      * @param data Not used.
      */
     function stake(uint256 amount, bytes calldata data) external {
-        _stakeFor(msg.sender, msg.sender, amount);
-    }
-
-    /**
-     * @dev Transfers amount of deposit tokens from the caller on behalf of user.
-     * @param user User address who gains credit for this stake operation.
-     * @param amount Number of deposit tokens to stake.
-     * @param data Not used.
-     */
-    function stakeFor(
-        address user,
-        uint256 amount,
-        bytes calldata data
-    ) external onlyOwner {
-        _stakeFor(msg.sender, user, amount);
-    }
-
-    /**
-     * @dev Private implementation of staking methods.
-     * @param staker User address who deposits tokens to stake.
-     * @param beneficiary User address who gains credit for this stake operation.
-     * @param amount Number of deposit tokens to stake.
-     */
-    function _stakeFor(address staker, address beneficiary, uint256 amount) private {
         require(amount > 0, "TokenGeyser: stake amount is zero");
-        require(beneficiary != address(0), "TokenGeyser: beneficiary is zero address");
         require(
             totalStakingShares == 0 || totalStaked() > 0,
             "TokenGeyser: Invalid state. Staking shares exist, but no staking tokens do"
@@ -187,12 +162,12 @@ contract TokenGeyser is IStaking, Ownable {
         updateAccounting();
 
         // 1. User Accounting
-        UserTotals storage totals = _userTotals[beneficiary];
+        UserTotals storage totals = _userTotals[msg.sender];
         totals.stakingShares = totals.stakingShares.add(mintedStakingShares);
         totals.lastAccountingTimestampSec = block.timestamp;
 
         Stake memory newStake = Stake(mintedStakingShares, block.timestamp);
-        _userStakes[beneficiary].push(newStake);
+        _userStakes[msg.sender].push(newStake);
 
         // 2. Global Accounting
         totalStakingShares = totalStakingShares.add(mintedStakingShares);
@@ -201,11 +176,11 @@ contract TokenGeyser is IStaking, Ownable {
 
         // interactions
         require(
-            _stakingPool.token().transferFrom(staker, address(_stakingPool), amount),
+            _stakingPool.token().transferFrom(msg.sender, address(_stakingPool), amount),
             "TokenGeyser: transfer into staking pool failed"
         );
 
-        emit Staked(beneficiary, amount, totalStakedFor(beneficiary), "");
+        emit Staked(msg.sender, amount, totalStakedFor(msg.sender), "");
     }
 
     /**

From 8c73361cc5bd63a180d8377f0b19f216ae008345 Mon Sep 17 00:00:00 2001
From: aalavandhan1984 <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 14:28:09 -0500
Subject: [PATCH 02/10] updated interface

---
 contracts/IStaking.sol             |  15 ---
 contracts/ITokenGeyser.sol         |  18 +++
 contracts/TokenGeyser.sol          | 184 +++++++++++++----------------
 contracts/_external/ampleforth.sol |   2 +-
 test/helper.ts                     |   3 +-
 test/staking.ts                    |  48 +++-----
 test/token_pool.ts                 |   6 +-
 test/token_unlock.ts               |  11 +-
 test/unstake.ts                    | 102 ++++++++--------
 9 files changed, 177 insertions(+), 212 deletions(-)
 delete mode 100644 contracts/IStaking.sol
 create mode 100644 contracts/ITokenGeyser.sol

diff --git a/contracts/IStaking.sol b/contracts/IStaking.sol
deleted file mode 100644
index 26c5f74..0000000
--- a/contracts/IStaking.sol
+++ /dev/null
@@ -1,15 +0,0 @@
-pragma solidity ^0.8.24;
-
-/**
- * @title Staking interface, as defined by EIP-900.
- * @dev https://github.com/ethereum/EIPs/blob/master/EIPS/eip-900.md
- */
-interface IStaking {
-    function stake(uint256 amount, bytes calldata data) external;
-    function stakeFor(address user, uint256 amount, bytes calldata data) external;
-    function unstake(uint256 amount, bytes calldata data) external;
-    function totalStakedFor(address addr) external view returns (uint256);
-    function totalStaked() external view returns (uint256);
-    function token() external view returns (address);
-    function supportsHistory() external view returns (bool);
-}
diff --git a/contracts/ITokenGeyser.sol b/contracts/ITokenGeyser.sol
new file mode 100644
index 0000000..8bfb847
--- /dev/null
+++ b/contracts/ITokenGeyser.sol
@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.24;
+
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+/**
+ * @title Geyser staking interface
+ */
+interface ITokenGeyser {
+    function stake(uint256 amount) external;
+    function unstake(uint256 amount) external returns (uint256);
+    function totalStakedFor(address addr) external view returns (uint256);
+    function totalStaked() external view returns (uint256);
+    function totalLocked() external view returns (uint256);
+    function totalUnlocked() external view returns (uint256);
+    function stakingToken() external view returns (IERC20);
+    function distributionToken() external view returns (IERC20);
+}
diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 0de0e33..36b944a 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -6,7 +6,7 @@ import { SafeMathCompatibility } from "./_utils/SafeMathCompatibility.sol";
 import { TokenPool } from "./TokenPool.sol";
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import { IStaking } from "./IStaking.sol";
+import { ITokenGeyser } from "./ITokenGeyser.sol";
 
 /**
  * @title Token Geyser
@@ -26,16 +26,22 @@ import { IStaking } from "./IStaking.sol";
  *      More background and motivation available at:
  *      https://github.com/ampleforth/RFCs/blob/master/RFCs/rfc-1.md
  */
-contract TokenGeyser is IStaking, Ownable {
+contract TokenGeyser is ITokenGeyser, Ownable {
     using SafeMathCompatibility for uint256;
 
-    event Staked(address indexed user, uint256 amount, uint256 total, bytes data);
-    event Unstaked(address indexed user, uint256 amount, uint256 total, bytes data);
+    //-------------------------------------------------------------------------
+    // Events
+
+    event Staked(address indexed user, uint256 amount, uint256 total);
+    event Unstaked(address indexed user, uint256 amount, uint256 total);
     event TokensClaimed(address indexed user, uint256 amount);
     event TokensLocked(uint256 amount, uint256 durationSec, uint256 total);
     // amount: Unlocked tokens, total: Total locked tokens
     event TokensUnlocked(uint256 amount, uint256 total);
 
+    //-------------------------------------------------------------------------
+    // Storage
+
     TokenPool private _stakingPool;
     TokenPool private _unlockedPool;
     TokenPool private _lockedPool;
@@ -93,9 +99,12 @@ contract TokenGeyser is IStaking, Ownable {
 
     UnlockSchedule[] public unlockSchedules;
 
+    //-------------------------------------------------------------------------
+    // Constructor
+
     /**
-     * @param stakingToken The token users deposit as stake.
-     * @param distributionToken The token users receive as they unstake.
+     * @param stakingToken_ The token users deposit as stake.
+     * @param distributionToken_ The token users receive as they unstake.
      * @param maxUnlockSchedules Max number of unlock stages, to guard against hitting gas limit.
      * @param startBonus_ Starting time bonus, BONUS_DECIMALS fixed point.
      *                    e.g. 25% means user gets 25% of max distribution tokens.
@@ -103,8 +112,8 @@ contract TokenGeyser is IStaking, Ownable {
      * @param initialSharesPerToken Number of shares to mint per staking token on first stake.
      */
     constructor(
-        IERC20 stakingToken,
-        IERC20 distributionToken,
+        IERC20 stakingToken_,
+        IERC20 distributionToken_,
         uint256 maxUnlockSchedules,
         uint256 startBonus_,
         uint256 bonusPeriodSec_,
@@ -118,26 +127,29 @@ contract TokenGeyser is IStaking, Ownable {
         require(initialSharesPerToken > 0, "TokenGeyser: initialSharesPerToken is zero");
 
         // TODO: use a factory here.
-        _stakingPool = new TokenPool(stakingToken);
-        _unlockedPool = new TokenPool(distributionToken);
-        _lockedPool = new TokenPool(distributionToken);
+        _stakingPool = new TokenPool(stakingToken_);
+        _unlockedPool = new TokenPool(distributionToken_);
+        _lockedPool = new TokenPool(distributionToken_);
         startBonus = startBonus_;
         bonusPeriodSec = bonusPeriodSec_;
         _maxUnlockSchedules = maxUnlockSchedules;
         _initialSharesPerToken = initialSharesPerToken;
     }
 
+    //-------------------------------------------------------------------------
+    // External and public methods
+
     /**
      * @return The token users deposit as stake.
      */
-    function getStakingToken() public view returns (IERC20) {
+    function stakingToken() public view override returns (IERC20) {
         return _stakingPool.token();
     }
 
     /**
      * @return The token users receive as they unstake.
      */
-    function getDistributionToken() public view returns (IERC20) {
+    function distributionToken() public view override returns (IERC20) {
         assert(_unlockedPool.token() == _lockedPool.token());
         return _unlockedPool.token();
     }
@@ -145,13 +157,12 @@ contract TokenGeyser is IStaking, Ownable {
     /**
      * @dev Transfers amount of deposit tokens from the user.
      * @param amount Number of deposit tokens to stake.
-     * @param data Not used.
      */
-    function stake(uint256 amount, bytes calldata data) external {
+    function stake(uint256 amount) external {
         require(amount > 0, "TokenGeyser: stake amount is zero");
         require(
             totalStakingShares == 0 || totalStaked() > 0,
-            "TokenGeyser: Invalid state. Staking shares exist, but no staking tokens do"
+            "TokenGeyser: Staking shares exist, but no staking tokens do"
         );
 
         uint256 mintedStakingShares = (totalStakingShares > 0)
@@ -180,34 +191,15 @@ contract TokenGeyser is IStaking, Ownable {
             "TokenGeyser: transfer into staking pool failed"
         );
 
-        emit Staked(msg.sender, amount, totalStakedFor(msg.sender), "");
+        emit Staked(msg.sender, amount, totalStakedFor(msg.sender));
     }
 
     /**
      * @dev Unstakes a certain amount of previously deposited tokens. User also receives their
      * alotted number of distribution tokens.
      * @param amount Number of deposit tokens to unstake / withdraw.
-     * @param data Not used.
-     */
-    function unstake(uint256 amount, bytes calldata data) external {
-        _unstake(amount);
-    }
-
-    /**
-     * @param amount Number of deposit tokens to unstake / withdraw.
-     * @return The total number of distribution tokens that would be rewarded.
      */
-    function unstakeQuery(uint256 amount) public returns (uint256) {
-        return _unstake(amount);
-    }
-
-    /**
-     * @dev Unstakes a certain amount of previously deposited tokens. User also receives their
-     * alotted number of distribution tokens.
-     * @param amount Number of deposit tokens to unstake / withdraw.
-     * @return The total number of distribution tokens rewarded.
-     */
-    function _unstake(uint256 amount) private returns (uint256) {
+    function unstake(uint256 amount) external returns (uint256) {
         updateAccounting();
 
         // checks
@@ -287,12 +279,12 @@ contract TokenGeyser is IStaking, Ownable {
             "TokenGeyser: transfer out of unlocked pool failed"
         );
 
-        emit Unstaked(msg.sender, amount, totalStakedFor(msg.sender), "");
+        emit Unstaked(msg.sender, amount, totalStakedFor(msg.sender));
         emit TokensClaimed(msg.sender, rewardAmount);
 
         require(
             totalStakingShares == 0 || totalStaked() > 0,
-            "TokenGeyser: Error unstaking. Staking shares exist, but no staking tokens do"
+            "TokenGeyser: Staking shares exist, but no staking tokens do"
         );
         return rewardAmount;
     }
@@ -315,7 +307,7 @@ contract TokenGeyser is IStaking, Ownable {
         uint256 currentRewardTokens,
         uint256 stakingShareSeconds,
         uint256 stakeTimeSec
-    ) private view returns (uint256) {
+    ) public view returns (uint256) {
         uint256 newRewardTokens = totalUnlocked().mul(stakingShareSeconds).div(
             _totalStakingShareSeconds
         );
@@ -352,15 +344,6 @@ contract TokenGeyser is IStaking, Ownable {
         return _stakingPool.balance();
     }
 
-    /**
-     * @dev Note that this application has a staking token as well as a distribution token, which
-     * may be different. This function is required by EIP-900.
-     * @return The deposit token used for staking.
-     */
-    function token() external view returns (address) {
-        return address(getStakingToken());
-    }
-
     /**
      * @dev A globally callable function to update the accounting state of the system.
      *      Global state and state for the caller are updated.
@@ -415,14 +398,14 @@ contract TokenGeyser is IStaking, Ownable {
     /**
      * @return Total number of locked distribution tokens.
      */
-    function totalLocked() public view returns (uint256) {
+    function totalLocked() public view override returns (uint256) {
         return _lockedPool.balance();
     }
 
     /**
      * @return Total number of unlocked distribution tokens.
      */
-    function totalUnlocked() public view returns (uint256) {
+    function totalUnlocked() public view override returns (uint256) {
         return _unlockedPool.balance();
     }
 
@@ -433,6 +416,40 @@ contract TokenGeyser is IStaking, Ownable {
         return unlockSchedules.length;
     }
 
+    /**
+     * @dev Moves distribution tokens from the locked pool to the unlocked pool, according to the
+     *      previously defined unlock schedules. Publicly callable.
+     * @return Number of newly unlocked distribution tokens.
+     */
+    function unlockTokens() public returns (uint256) {
+        uint256 unlockedTokens = 0;
+        uint256 lockedTokens = totalLocked();
+
+        if (totalLockedShares == 0) {
+            unlockedTokens = lockedTokens;
+        } else {
+            uint256 unlockedShares = 0;
+            for (uint256 s = 0; s < unlockSchedules.length; s++) {
+                unlockedShares = unlockedShares.add(_unlockScheduleShares(s));
+            }
+            unlockedTokens = unlockedShares.mul(lockedTokens).div(totalLockedShares);
+            totalLockedShares = totalLockedShares.sub(unlockedShares);
+        }
+
+        if (unlockedTokens > 0) {
+            require(
+                _lockedPool.transfer(address(_unlockedPool), unlockedTokens),
+                "TokenGeyser: transfer out of locked pool failed"
+            );
+            emit TokensUnlocked(unlockedTokens, totalLocked());
+        }
+
+        return unlockedTokens;
+    }
+
+    //-------------------------------------------------------------------------
+    // Admin only methods
+
     /**
      * @dev This funcion allows the contract owner to add more locked distribution tokens, along
      *      with the associated "unlock schedule". These locked tokens immediately begin unlocking
@@ -471,36 +488,23 @@ contract TokenGeyser is IStaking, Ownable {
     }
 
     /**
-     * @dev Moves distribution tokens from the locked pool to the unlocked pool, according to the
-     *      previously defined unlock schedules. Publicly callable.
-     * @return Number of newly unlocked distribution tokens.
+     * @dev Lets the owner rescue funds air-dropped to the staking pool.
+     * @param tokenToRescue Address of the token to be rescued.
+     * @param to Address to which the rescued funds are to be sent.
+     * @param amount Amount of tokens to be rescued.
+     * @return Transfer success.
      */
-    function unlockTokens() public returns (uint256) {
-        uint256 unlockedTokens = 0;
-        uint256 lockedTokens = totalLocked();
-
-        if (totalLockedShares == 0) {
-            unlockedTokens = lockedTokens;
-        } else {
-            uint256 unlockedShares = 0;
-            for (uint256 s = 0; s < unlockSchedules.length; s++) {
-                unlockedShares = unlockedShares.add(unlockScheduleShares(s));
-            }
-            unlockedTokens = unlockedShares.mul(lockedTokens).div(totalLockedShares);
-            totalLockedShares = totalLockedShares.sub(unlockedShares);
-        }
-
-        if (unlockedTokens > 0) {
-            require(
-                _lockedPool.transfer(address(_unlockedPool), unlockedTokens),
-                "TokenGeyser: transfer out of locked pool failed"
-            );
-            emit TokensUnlocked(unlockedTokens, totalLocked());
-        }
-
-        return unlockedTokens;
+    function rescueFundsFromStakingPool(
+        address tokenToRescue,
+        address to,
+        uint256 amount
+    ) public onlyOwner returns (bool) {
+        return _stakingPool.rescueFunds(tokenToRescue, to, amount);
     }
 
+    //-------------------------------------------------------------------------
+    // Private methods
+
     /**
      * @dev Returns the number of unlockable shares from a given schedule. The returned value
      *      depends on the time since the last unlock. This function updates schedule accounting,
@@ -508,7 +512,7 @@ contract TokenGeyser is IStaking, Ownable {
      * @param s Index of the unlock schedule.
      * @return The number of unlocked shares.
      */
-    function unlockScheduleShares(uint256 s) private returns (uint256) {
+    function _unlockScheduleShares(uint256 s) private returns (uint256) {
         UnlockSchedule storage schedule = unlockSchedules[s];
 
         if (schedule.unlockedShares >= schedule.initialLockedShares) {
@@ -532,26 +536,4 @@ contract TokenGeyser is IStaking, Ownable {
         schedule.unlockedShares = schedule.unlockedShares.add(sharesToUnlock);
         return sharesToUnlock;
     }
-
-    /**
-     * @dev Lets the owner rescue funds air-dropped to the staking pool.
-     * @param tokenToRescue Address of the token to be rescued.
-     * @param to Address to which the rescued funds are to be sent.
-     * @param amount Amount of tokens to be rescued.
-     * @return Transfer success.
-     */
-    function rescueFundsFromStakingPool(
-        address tokenToRescue,
-        address to,
-        uint256 amount
-    ) public onlyOwner returns (bool) {
-        return _stakingPool.rescueFunds(tokenToRescue, to, amount);
-    }
-
-    /**
-     * @return False. This application does not support staking history.
-     */
-    function supportsHistory() external pure returns (bool) {
-        return false;
-    }
 }
diff --git a/contracts/_external/ampleforth.sol b/contracts/_external/ampleforth.sol
index 0f37b7c..98e239b 100644
--- a/contracts/_external/ampleforth.sol
+++ b/contracts/_external/ampleforth.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
-// solhint-disable-next-line compiler-version
+pragma solidity ^0.8.0;
 
 // solhint-disable-next-line no-unused-import
 import { UFragments } from "ampleforth-contracts/contracts/UFragments.sol";
diff --git a/test/helper.ts b/test/helper.ts
index 0af3e4a..61633a6 100644
--- a/test/helper.ts
+++ b/test/helper.ts
@@ -1,5 +1,4 @@
-import { ethers } from "hardhat";
-import { promisify } from "util";
+import hre, { ethers } from "hardhat";
 import { expect } from "chai";
 
 const AMPL_DECIMALS = 9;
diff --git a/test/staking.ts b/test/staking.ts
index 0ba4a05..cd362a2 100644
--- a/test/staking.ts
+++ b/test/staking.ts
@@ -1,8 +1,8 @@
 import { ethers } from "hardhat";
 import { expect } from "chai";
-import { loadFixture } from "@nomicfoundation/hardhat-network-helpers";
-import { TimeHelpers, $AMPL, invokeRebase } from "../test/helper";
-import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
+import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
+import { $AMPL, invokeRebase } from "../test/helper";
+import { SignerWithAddress } from "ethers";
 
 let ampl: any, dist: any, owner: SignerWithAddress, anotherAccount: SignerWithAddress;
 const InitialSharesPerToken = BigInt(10 ** 6);
@@ -58,21 +58,9 @@ describe("staking", function () {
     });
   });
 
-  describe("getStakingToken", function () {
+  describe("stakingToken", function () {
     it("should return the staking token", async function () {
-      expect(await dist.getStakingToken()).to.equal(ampl.target);
-    });
-  });
-
-  describe("token", function () {
-    it("should return the staking token", async function () {
-      expect(await dist.token()).to.equal(ampl.target);
-    });
-  });
-
-  describe("supportsHistory", function () {
-    it("should return supportsHistory", async function () {
-      expect(await dist.supportsHistory()).to.be.false;
+      expect(await dist.stakingToken()).to.equal(ampl.target);
     });
   });
 
@@ -80,7 +68,7 @@ describe("staking", function () {
     describe("when the amount is 0", function () {
       it("should fail", async function () {
         await ampl.approve(dist.target, $AMPL(1000));
-        await expect(dist.stake($AMPL(0), "0x")).to.be.revertedWith(
+        await expect(dist.stake($AMPL(0))).to.be.revertedWith(
           "TokenGeyser: stake amount is zero",
         );
       });
@@ -88,7 +76,7 @@ describe("staking", function () {
 
     describe("when token transfer has not been approved", function () {
       it("should fail", async function () {
-        await expect(dist.stake($AMPL(100), "0x")).to.be.reverted;
+        await expect(dist.stake($AMPL(100))).to.be.reverted;
       });
     });
 
@@ -98,7 +86,7 @@ describe("staking", function () {
         await ampl.approve(dist.target, $AMPL(100));
       });
       it("should update the total staked", async function () {
-        await dist.stake($AMPL(100), "0x");
+        await dist.stake($AMPL(100));
         expect(await dist.totalStaked()).to.equal($AMPL(100));
         expect(await dist.totalStakedFor(await owner.getAddress())).to.equal($AMPL(100));
         expect(await dist.totalStakingShares()).to.equal(
@@ -106,10 +94,10 @@ describe("staking", function () {
         );
       });
       it("should log Staked", async function () {
-        const tx = await dist.stake($AMPL(100), "0x");
+        const tx = await dist.stake($AMPL(100));
         await expect(tx)
           .to.emit(dist, "Staked")
-          .withArgs(await owner.getAddress(), $AMPL(100), $AMPL(100), "0x");
+          .withArgs(await owner.getAddress(), $AMPL(100), $AMPL(100));
       });
     });
 
@@ -118,9 +106,9 @@ describe("staking", function () {
         expect(await dist.totalStaked()).to.equal($AMPL(0));
         await ampl.transfer(await anotherAccount.getAddress(), $AMPL(50));
         await ampl.connect(anotherAccount).approve(dist.target, $AMPL(50));
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
         await ampl.approve(dist.target, $AMPL(150));
-        await dist.stake($AMPL(150), "0x");
+        await dist.stake($AMPL(150));
       });
       it("should update the total staked", async function () {
         expect(await dist.totalStaked()).to.equal($AMPL(200));
@@ -139,11 +127,11 @@ describe("staking", function () {
         expect(await dist.totalStaked()).to.equal($AMPL(0));
         await ampl.transfer(await anotherAccount.getAddress(), $AMPL(50));
         await ampl.connect(anotherAccount).approve(dist.target, $AMPL(50));
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
         await ampl.approve(dist.target, $AMPL(150));
         await invokeRebase(ampl, 100);
         expect(await dist.totalStaked()).to.equal($AMPL(100));
-        await dist.stake($AMPL(150), "0x");
+        await dist.stake($AMPL(150));
       });
       it("should updated the total staked shares", async function () {
         expect(await dist.totalStaked()).to.equal($AMPL(250));
@@ -160,11 +148,11 @@ describe("staking", function () {
     describe("when totalStaked>0, when rebase increases supply", function () {
       beforeEach(async function () {
         await ampl.approve(dist.target, $AMPL(51));
-        await dist.stake($AMPL(50), "0x");
+        await dist.stake($AMPL(50));
       });
       it("should fail if there are too few mintedStakingShares", async function () {
         await invokeRebase(ampl, 100n * InitialSharesPerToken);
-        await expect(dist.stake(1, "0x")).to.be.revertedWith(
+        await expect(dist.stake(1)).to.be.revertedWith(
           "TokenGeyser: Stake amount is too small",
         );
       });
@@ -175,11 +163,11 @@ describe("staking", function () {
         expect(await dist.totalStaked()).to.equal($AMPL(0));
         await ampl.transfer(await anotherAccount.getAddress(), $AMPL(50));
         await ampl.connect(anotherAccount).approve(dist.target, $AMPL(50));
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
         await ampl.approve(dist.target, $AMPL(150));
         await invokeRebase(ampl, -50);
         expect(await dist.totalStaked()).to.equal($AMPL(25));
-        await dist.stake($AMPL(150), "0x");
+        await dist.stake($AMPL(150));
       });
       it("should updated the total staked shares", async function () {
         expect(await dist.totalStaked()).to.equal($AMPL(175));
diff --git a/test/token_pool.ts b/test/token_pool.ts
index 67d2f4f..4843ea5 100644
--- a/test/token_pool.ts
+++ b/test/token_pool.ts
@@ -1,7 +1,7 @@
-import { ethers, waffle } from "hardhat";
+import { ethers } from "hardhat";
 import { expect } from "chai";
-import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
-import { loadFixture } from "@nomicfoundation/hardhat-network-helpers";
+import { SignerWithAddress } from "ethers";
+import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
 
 let owner: SignerWithAddress, anotherAccount: SignerWithAddress;
 
diff --git a/test/token_unlock.ts b/test/token_unlock.ts
index cbd5e23..459f67c 100644
--- a/test/token_unlock.ts
+++ b/test/token_unlock.ts
@@ -1,15 +1,14 @@
 import { ethers } from "hardhat";
 import { expect } from "chai";
-import { loadFixture } from "@nomicfoundation/hardhat-network-helpers";
+import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
 import {
   TimeHelpers,
   $AMPL,
   invokeRebase,
   checkAmplAprox,
   checkSharesAprox,
-  setTimeForNextTransaction,
 } from "../test/helper";
-import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
+import { SignerWithAddress } from "ethers";
 
 let ampl: any, dist: any, owner: SignerWithAddress, anotherAccount: SignerWithAddress;
 const InitialSharesPerToken = BigInt(10 ** 6);
@@ -50,9 +49,9 @@ describe("LockedPool", function () {
     ({ ampl, dist, owner, anotherAccount } = await loadFixture(setupContracts));
   });
 
-  describe("getDistributionToken", function () {
+  describe("distributionToken", function () {
     it("should return the staking token", async function () {
-      expect(await dist.getDistributionToken.staticCall()).to.equal(ampl.target);
+      expect(await dist.distributionToken.staticCall()).to.equal(ampl.target);
     });
   });
 
@@ -462,7 +461,7 @@ describe("LockedPool", function () {
       _r = await dist.updateAccounting.staticCall({ from: owner });
       _t = await TimeHelpers.currentTime();
       await ampl.approve(dist.target, $AMPL(300));
-      await dist.stake($AMPL(100), "0x");
+      await dist.stake($AMPL(100));
       await dist.lockTokens($AMPL(100), ONE_YEAR);
       await TimeHelpers.increaseTime(ONE_YEAR / 2);
       await dist.lockTokens($AMPL(100), ONE_YEAR);
diff --git a/test/unstake.ts b/test/unstake.ts
index b0b3c2b..cdbd06e 100644
--- a/test/unstake.ts
+++ b/test/unstake.ts
@@ -1,8 +1,8 @@
 import { ethers } from "hardhat";
 import { expect } from "chai";
-import { loadFixture, time } from "@nomicfoundation/hardhat-network-helpers";
+import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
 import { $AMPL, invokeRebase, checkAmplAprox, TimeHelpers } from "../test/helper";
-import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
+import { SignerWithAddress } from "ethers";
 
 let ampl: any, dist: any, owner: SignerWithAddress, anotherAccount: SignerWithAddress;
 const InitialSharesPerToken = 10 ** 6;
@@ -54,33 +54,31 @@ describe("unstaking", function () {
   describe("unstake", function () {
     describe("when amount is 0", function () {
       it("should fail", async function () {
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
-        await expect(
-          dist.connect(anotherAccount).unstake($AMPL(0), "0x"),
-        ).to.be.revertedWith("TokenGeyser: unstake amount is zero");
+        await dist.connect(anotherAccount).stake($AMPL(50));
+        await expect(dist.connect(anotherAccount).unstake($AMPL(0))).to.be.revertedWith(
+          "TokenGeyser: unstake amount is zero",
+        );
       });
     });
 
     describe("when rebase increases supply", function () {
       beforeEach(async function () {
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
-        await time.increase(1);
+        await dist.connect(anotherAccount).stake($AMPL(50));
+        await TimeHelpers.increaseTime(1);
       });
       it("should fail if user tries to unstake more than his balance", async function () {
         await invokeRebase(ampl, +50);
-        await expect(
-          dist.connect(anotherAccount).unstake($AMPL(85), "0x"),
-        ).to.be.revertedWith(
+        await expect(dist.connect(anotherAccount).unstake($AMPL(85))).to.be.revertedWith(
           "TokenGeyser: unstake amount is greater than total user stakes",
         );
       });
       it("should NOT fail if user tries to unstake his balance", async function () {
         await invokeRebase(ampl, +50);
-        await dist.connect(anotherAccount).unstake($AMPL(75), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(75));
       });
       it("should fail if there are too few stakingSharesToBurn", async function () {
         await invokeRebase(ampl, 100 * InitialSharesPerToken);
-        await expect(dist.connect(anotherAccount).unstake(1, "0x")).to.be.revertedWith(
+        await expect(dist.connect(anotherAccount).unstake(1)).to.be.revertedWith(
           "TokenGeyser: Unable to unstake amount this small",
         );
       });
@@ -88,20 +86,18 @@ describe("unstaking", function () {
 
     describe("when rebase decreases supply", function () {
       beforeEach(async function () {
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
-        await time.increase(1);
+        await dist.connect(anotherAccount).stake($AMPL(50));
+        await TimeHelpers.increaseTime(1);
       });
       it("should fail if user tries to unstake more than his balance", async function () {
         await invokeRebase(ampl, -50);
-        await expect(
-          dist.connect(anotherAccount).unstake($AMPL(50), "0x"),
-        ).to.be.revertedWith(
+        await expect(dist.connect(anotherAccount).unstake($AMPL(50))).to.be.revertedWith(
           "TokenGeyser: unstake amount is greater than total user stakes",
         );
       });
       it("should NOT fail if user tries to unstake his balance", async function () {
         await invokeRebase(ampl, -50);
-        await dist.connect(anotherAccount).unstake($AMPL(25), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(25));
       });
     });
 
@@ -112,13 +108,13 @@ describe("unstaking", function () {
       // user's final balance is 90 ampl, (20 remains staked), eligible rewards (40 ampl)
       beforeEach(async function () {
         await dist.lockTokens($AMPL(100), ONE_YEAR);
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
         await TimeHelpers.increaseTime(ONE_YEAR);
         await dist.connect(anotherAccount).updateAccounting();
         checkAmplAprox(await totalRewardsFor(anotherAccount), 100);
       });
       it("should update the total staked and rewards", async function () {
-        await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(30));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(20));
         expect(
           await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
@@ -127,21 +123,20 @@ describe("unstaking", function () {
       });
       it("should transfer back staked tokens + rewards", async function () {
         const _b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
-        await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(30));
         const b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
         checkAmplAprox(b - _b, 90);
       });
       it("should log Unstaked", async function () {
-        const r = await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        const r = await dist.connect(anotherAccount).unstake($AMPL(30));
         await expectEvent(r, "Unstaked", [
           await anotherAccount.getAddress(),
           $AMPL(30),
           $AMPL(20),
-          "0x",
         ]);
       });
       it("should log TokensClaimed", async function () {
-        const r = await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        const r = await dist.connect(anotherAccount).unstake($AMPL(30));
         await expectEvent(r, "TokensClaimed", [
           await anotherAccount.getAddress(),
           $AMPL(60),
@@ -160,13 +155,13 @@ describe("unstaking", function () {
       beforeEach(async function () {
         await dist.lockTokens($AMPL(1000), ONE_HOUR);
 
-        await dist.connect(anotherAccount).stake($AMPL(500), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(500));
         await TimeHelpers.increaseTime(12 * ONE_HOUR);
         await dist.connect(anotherAccount).updateAccounting();
         checkAmplAprox(await totalRewardsFor(anotherAccount), 1000);
       });
       it("should update the total staked and rewards", async function () {
-        await dist.connect(anotherAccount).unstake($AMPL(250), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(250));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(250));
         expect(
           await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
@@ -175,21 +170,20 @@ describe("unstaking", function () {
       });
       it("should transfer back staked tokens + rewards", async function () {
         const _b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
-        await dist.connect(anotherAccount).unstake($AMPL(250), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(250));
         const b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
         checkAmplAprox(b - _b, 625);
       });
       it("should log Unstaked", async function () {
-        const r = await dist.connect(anotherAccount).unstake($AMPL(250), "0x");
+        const r = await dist.connect(anotherAccount).unstake($AMPL(250));
         await expectEvent(r, "Unstaked", [
           await anotherAccount.getAddress(),
           $AMPL(250),
           $AMPL(250),
-          "0x",
         ]);
       });
       it("should log TokensClaimed", async function () {
-        const r = await dist.connect(anotherAccount).unstake($AMPL(250), "0x");
+        const r = await dist.connect(anotherAccount).unstake($AMPL(250));
         await expectEvent(r, "TokensClaimed", [
           await anotherAccount.getAddress(),
           $AMPL(375), // .5 * .75 * 1000
@@ -206,10 +200,10 @@ describe("unstaking", function () {
         await dist.lockTokens($AMPL(100), ONE_YEAR);
 
         await TimeHelpers.increaseTime(ONE_YEAR / 100);
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
 
         await TimeHelpers.increaseTime(ONE_YEAR / 4);
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
         await TimeHelpers.increaseTime(ONE_YEAR / 4);
         await dist.connect(anotherAccount).updateAccounting();
       });
@@ -217,7 +211,7 @@ describe("unstaking", function () {
         checkAmplAprox(await totalRewardsFor(anotherAccount), 51);
       });
       it("should update the total staked and rewards", async function () {
-        await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(30));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(70));
         expect(
           await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
@@ -226,7 +220,7 @@ describe("unstaking", function () {
       });
       it("should transfer back staked tokens + rewards", async function () {
         const _b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
-        await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(30));
         const b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
         checkAmplAprox(b - _b, 40.2);
       });
@@ -240,27 +234,27 @@ describe("unstaking", function () {
       beforeEach(async function () {
         await dist.lockTokens($AMPL(100), ONE_YEAR);
 
-        await dist.connect(anotherAccount).stake($AMPL(10), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(10));
         await TimeHelpers.increaseTime(ONE_YEAR);
-        await dist.connect(anotherAccount).stake($AMPL(10), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(10));
         await TimeHelpers.increaseTime(ONE_YEAR);
         await dist.connect(anotherAccount).updateAccounting();
         checkAmplAprox(await totalRewardsFor(anotherAccount), 100);
       });
 
       it("should use updated user accounting", async function () {
-        const r1 = await dist.connect(anotherAccount).unstake($AMPL(5), "0x");
+        const r1 = await dist.connect(anotherAccount).unstake($AMPL(5));
         await expectEvent(r1, "TokensClaimed", [
           await anotherAccount.getAddress(),
           16666666842n,
         ]);
         const claim1 = 16666666842n;
-        const r2 = await dist.connect(anotherAccount).unstake($AMPL(5), "0x");
+        const r2 = await dist.connect(anotherAccount).unstake($AMPL(5));
         await expectEvent(r2, "TokensClaimed", [
           await anotherAccount.getAddress(),
           16666667054n,
         ]);
-        const r3 = await dist.connect(anotherAccount).unstake($AMPL(5), "0x");
+        const r3 = await dist.connect(anotherAccount).unstake($AMPL(5));
         await expectEvent(r3, "TokensClaimed", [
           await anotherAccount.getAddress(),
           33333333052n,
@@ -280,10 +274,10 @@ describe("unstaking", function () {
         await dist.lockTokens($AMPL(100), ONE_YEAR);
 
         await TimeHelpers.increaseTime(ONE_YEAR / 100);
-        await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(50));
 
         await TimeHelpers.increaseTime(ONE_YEAR / 4);
-        await dist.stake($AMPL(50), "0x");
+        await dist.stake($AMPL(50));
         await TimeHelpers.increaseTime(ONE_YEAR / 2);
         await dist.connect(anotherAccount).updateAccounting();
         await dist.updateAccounting();
@@ -297,7 +291,7 @@ describe("unstaking", function () {
         checkAmplAprox(await totalRewardsFor(owner), 30.4);
       });
       it("should update the total staked and rewards", async function () {
-        await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(30));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(70));
         expect(
           await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
@@ -310,7 +304,7 @@ describe("unstaking", function () {
       });
       it("should transfer back staked tokens + rewards", async function () {
         const _b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
-        await dist.connect(anotherAccount).unstake($AMPL(30), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(30));
         const b = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
         checkAmplAprox(b - _b, 57.36);
       });
@@ -328,13 +322,13 @@ describe("unstaking", function () {
       const rewardsOwner = 32500.0 / 11.0;
       beforeEach(async function () {
         await dist.lockTokens($AMPL(10000), ONE_YEAR);
-        await dist.connect(anotherAccount).stake($AMPL(5000), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(5000));
 
         await TimeHelpers.increaseTime(ONE_YEAR / 4);
-        await dist.stake($AMPL(5000), "0x");
+        await dist.stake($AMPL(5000));
         await TimeHelpers.increaseTime(ONE_YEAR / 4);
-        await dist.connect(anotherAccount).stake($AMPL(5000), "0x");
-        await dist.stake($AMPL(3000), "0x");
+        await dist.connect(anotherAccount).stake($AMPL(5000));
+        await dist.stake($AMPL(3000));
         await TimeHelpers.increaseTime(ONE_YEAR / 4);
         await dist.connect(anotherAccount).updateAccounting();
         await dist.updateAccounting();
@@ -343,7 +337,7 @@ describe("unstaking", function () {
         checkAmplAprox(await totalRewardsFor(owner), rewardsOwner);
       });
       it("should update the total staked and rewards", async function () {
-        await dist.connect(anotherAccount).unstake($AMPL(10000), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(10000));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(8000));
         expect(await dist.totalStakedFor.staticCall(ethers.ZeroAddress)).to.eq($AMPL(0));
         expect(await dist.totalStakedFor.staticCall(await owner.getAddress())).to.eq(
@@ -351,7 +345,7 @@ describe("unstaking", function () {
         );
         checkAmplAprox(await totalRewardsFor(anotherAccount), 0);
         checkAmplAprox(await totalRewardsFor(owner), rewardsOwner);
-        await dist.unstake($AMPL(8000), "0x");
+        await dist.unstake($AMPL(8000));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(0));
         expect(
           await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
@@ -364,11 +358,11 @@ describe("unstaking", function () {
       });
       it("should transfer back staked tokens + rewards", async function () {
         const b1 = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
-        await dist.connect(anotherAccount).unstake($AMPL(10000), "0x");
+        await dist.connect(anotherAccount).unstake($AMPL(10000));
         const b2 = await ampl.balanceOf.staticCall(await anotherAccount.getAddress());
         checkAmplAprox(b2 - b1, 10000 + rewardsAnotherAccount);
         const b3 = await ampl.balanceOf.staticCall(await owner.getAddress());
-        await dist.unstake($AMPL(8000), "0x");
+        await dist.unstake($AMPL(8000));
         const b4 = await ampl.balanceOf.staticCall(await owner.getAddress());
         checkAmplAprox(b4 - b3, 8000 + rewardsOwner);
       });
@@ -381,14 +375,14 @@ describe("unstaking", function () {
     // unstakes 30 ampls, gets 60% of the reward (60 ampl)
     beforeEach(async function () {
       await dist.lockTokens($AMPL(100), ONE_YEAR);
-      await dist.connect(anotherAccount).stake($AMPL(50), "0x");
+      await dist.connect(anotherAccount).stake($AMPL(50));
       await TimeHelpers.increaseTime(ONE_YEAR);
       await dist.connect(anotherAccount).updateAccounting();
     });
     it("should return the reward amount", async function () {
       checkAmplAprox(await totalRewardsFor(anotherAccount), 100);
       checkAmplAprox(
-        await dist.connect(anotherAccount).unstakeQuery.staticCall($AMPL(30)),
+        await dist.connect(anotherAccount).unstake.staticCall($AMPL(30)),
         60,
       );
     });

From 3b91c11cfee5fa1f32d9f55418062f313a77c47e Mon Sep 17 00:00:00 2001
From: aalavandhann <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 20:23:10 -0500
Subject: [PATCH 03/10] made storage variables public

---
 contracts/TokenGeyser.sol | 100 +++++++++++++++++++-------------------
 1 file changed, 50 insertions(+), 50 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 36b944a..9abd265 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -42,9 +42,9 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     //-------------------------------------------------------------------------
     // Storage
 
-    TokenPool private _stakingPool;
-    TokenPool private _unlockedPool;
-    TokenPool private _lockedPool;
+    TokenPool public stakingPool;
+    TokenPool public unlockedPool;
+    TokenPool public lockedPool;
 
     //
     // Time-bonus params
@@ -58,10 +58,10 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     //
     uint256 public totalLockedShares = 0;
     uint256 public totalStakingShares = 0;
-    uint256 private _totalStakingShareSeconds = 0;
-    uint256 private _lastAccountingTimestampSec = block.timestamp;
-    uint256 private _maxUnlockSchedules = 0;
-    uint256 private _initialSharesPerToken = 0;
+    uint256 public totalStakingShareSeconds = 0;
+    uint256 public lastAccountingTimestampSec = block.timestamp;
+    uint256 public maxUnlockSchedules = 0;
+    uint256 public initialSharesPerToken = 0;
 
     //
     // User accounting state
@@ -81,10 +81,10 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     }
 
     // Aggregated staking values per user
-    mapping(address => UserTotals) private _userTotals;
+    mapping(address => UserTotals) public userTotals;
 
     // The collection of stakes for each user. Ordered by timestamp, earliest to latest.
-    mapping(address => Stake[]) private _userStakes;
+    mapping(address => Stake[]) public userStakes;
 
     //
     // Locked/Unlocked Accounting state
@@ -105,35 +105,35 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     /**
      * @param stakingToken_ The token users deposit as stake.
      * @param distributionToken_ The token users receive as they unstake.
-     * @param maxUnlockSchedules Max number of unlock stages, to guard against hitting gas limit.
+     * @param maxUnlockSchedules_ Max number of unlock stages, to guard against hitting gas limit.
      * @param startBonus_ Starting time bonus, BONUS_DECIMALS fixed point.
      *                    e.g. 25% means user gets 25% of max distribution tokens.
      * @param bonusPeriodSec_ Length of time for bonus to increase linearly to max.
-     * @param initialSharesPerToken Number of shares to mint per staking token on first stake.
+     * @param initialSharesPerToken_ Number of shares to mint per staking token on first stake.
      */
     constructor(
         IERC20 stakingToken_,
         IERC20 distributionToken_,
-        uint256 maxUnlockSchedules,
+        uint256 maxUnlockSchedules_,
         uint256 startBonus_,
         uint256 bonusPeriodSec_,
-        uint256 initialSharesPerToken
+        uint256 initialSharesPerToken_
     ) Ownable(msg.sender) {
         // The start bonus must be some fraction of the max. (i.e. <= 100%)
         require(startBonus_ <= 10 ** BONUS_DECIMALS, "TokenGeyser: start bonus too high");
         // If no period is desired, instead set startBonus = 100%
         // and bonusPeriod to a small value like 1sec.
         require(bonusPeriodSec_ != 0, "TokenGeyser: bonus period is zero");
-        require(initialSharesPerToken > 0, "TokenGeyser: initialSharesPerToken is zero");
+        require(initialSharesPerToken_ > 0, "TokenGeyser: initialSharesPerToken is zero");
 
         // TODO: use a factory here.
-        _stakingPool = new TokenPool(stakingToken_);
-        _unlockedPool = new TokenPool(distributionToken_);
-        _lockedPool = new TokenPool(distributionToken_);
+        stakingPool = new TokenPool(stakingToken_);
+        unlockedPool = new TokenPool(distributionToken_);
+        lockedPool = new TokenPool(distributionToken_);
         startBonus = startBonus_;
         bonusPeriodSec = bonusPeriodSec_;
-        _maxUnlockSchedules = maxUnlockSchedules;
-        _initialSharesPerToken = initialSharesPerToken;
+        maxUnlockSchedules = maxUnlockSchedules_;
+        initialSharesPerToken = initialSharesPerToken_;
     }
 
     //-------------------------------------------------------------------------
@@ -143,15 +143,15 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      * @return The token users deposit as stake.
      */
     function stakingToken() public view override returns (IERC20) {
-        return _stakingPool.token();
+        return stakingPool.token();
     }
 
     /**
      * @return The token users receive as they unstake.
      */
     function distributionToken() public view override returns (IERC20) {
-        assert(_unlockedPool.token() == _lockedPool.token());
-        return _unlockedPool.token();
+        assert(unlockedPool.token() == lockedPool.token());
+        return unlockedPool.token();
     }
 
     /**
@@ -167,27 +167,27 @@ contract TokenGeyser is ITokenGeyser, Ownable {
 
         uint256 mintedStakingShares = (totalStakingShares > 0)
             ? totalStakingShares.mul(amount).div(totalStaked())
-            : amount.mul(_initialSharesPerToken);
+            : amount.mul(initialSharesPerToken);
         require(mintedStakingShares > 0, "TokenGeyser: Stake amount is too small");
 
         updateAccounting();
 
         // 1. User Accounting
-        UserTotals storage totals = _userTotals[msg.sender];
+        UserTotals storage totals = userTotals[msg.sender];
         totals.stakingShares = totals.stakingShares.add(mintedStakingShares);
         totals.lastAccountingTimestampSec = block.timestamp;
 
         Stake memory newStake = Stake(mintedStakingShares, block.timestamp);
-        _userStakes[msg.sender].push(newStake);
+        userStakes[msg.sender].push(newStake);
 
         // 2. Global Accounting
         totalStakingShares = totalStakingShares.add(mintedStakingShares);
         // Already set in updateAccounting()
-        // _lastAccountingTimestampSec = block.timestamp;
+        // lastAccountingTimestampSec = block.timestamp;
 
         // interactions
         require(
-            _stakingPool.token().transferFrom(msg.sender, address(_stakingPool), amount),
+            stakingPool.token().transferFrom(msg.sender, address(stakingPool), amount),
             "TokenGeyser: transfer into staking pool failed"
         );
 
@@ -215,8 +215,8 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         );
 
         // 1. User Accounting
-        UserTotals storage totals = _userTotals[msg.sender];
-        Stake[] storage accountStakes = _userStakes[msg.sender];
+        UserTotals storage totals = userTotals[msg.sender];
+        Stake[] storage accountStakes = userStakes[msg.sender];
 
         // Redeem from most recent stake and go backwards in time.
         uint256 stakingShareSecondsToBurn = 0;
@@ -262,20 +262,20 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         // totals.lastAccountingTimestampSec = block.timestamp;
 
         // 2. Global Accounting
-        _totalStakingShareSeconds = _totalStakingShareSeconds.sub(
+        totalStakingShareSeconds = totalStakingShareSeconds.sub(
             stakingShareSecondsToBurn
         );
         totalStakingShares = totalStakingShares.sub(stakingSharesToBurn);
         // Already set in updateAccounting
-        // _lastAccountingTimestampSec = block.timestamp;
+        // lastAccountingTimestampSec = block.timestamp;
 
         // interactions
         require(
-            _stakingPool.transfer(msg.sender, amount),
+            stakingPool.transfer(msg.sender, amount),
             "TokenGeyser: transfer out of staking pool failed"
         );
         require(
-            _unlockedPool.transfer(msg.sender, rewardAmount),
+            unlockedPool.transfer(msg.sender, rewardAmount),
             "TokenGeyser: transfer out of unlocked pool failed"
         );
 
@@ -309,7 +309,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         uint256 stakeTimeSec
     ) public view returns (uint256) {
         uint256 newRewardTokens = totalUnlocked().mul(stakingShareSeconds).div(
-            _totalStakingShareSeconds
+            totalStakingShareSeconds
         );
 
         if (stakeTimeSec >= bonusPeriodSec) {
@@ -331,7 +331,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     function totalStakedFor(address addr) public view returns (uint256) {
         return
             totalStakingShares > 0
-                ? totalStaked().mul(_userTotals[addr].stakingShares).div(
+                ? totalStaked().mul(userTotals[addr].stakingShares).div(
                     totalStakingShares
                 )
                 : 0;
@@ -341,7 +341,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      * @return The total number of deposit tokens staked globally, by all users.
      */
     function totalStaked() public view returns (uint256) {
-        return _stakingPool.balance();
+        return stakingPool.balance();
     }
 
     /**
@@ -363,13 +363,13 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         // Global accounting
         uint256 newStakingShareSeconds = block
             .timestamp
-            .sub(_lastAccountingTimestampSec)
+            .sub(lastAccountingTimestampSec)
             .mul(totalStakingShares);
-        _totalStakingShareSeconds = _totalStakingShareSeconds.add(newStakingShareSeconds);
-        _lastAccountingTimestampSec = block.timestamp;
+        totalStakingShareSeconds = totalStakingShareSeconds.add(newStakingShareSeconds);
+        lastAccountingTimestampSec = block.timestamp;
 
         // User Accounting
-        UserTotals storage totals = _userTotals[msg.sender];
+        UserTotals storage totals = userTotals[msg.sender];
         uint256 newUserStakingShareSeconds = block
             .timestamp
             .sub(totals.lastAccountingTimestampSec)
@@ -379,9 +379,9 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         );
         totals.lastAccountingTimestampSec = block.timestamp;
 
-        uint256 totalUserRewards = (_totalStakingShareSeconds > 0)
+        uint256 totalUserRewards = (totalStakingShareSeconds > 0)
             ? totalUnlocked().mul(totals.stakingShareSeconds).div(
-                _totalStakingShareSeconds
+                totalStakingShareSeconds
             )
             : 0;
 
@@ -389,7 +389,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
             totalLocked(),
             totalUnlocked(),
             totals.stakingShareSeconds,
-            _totalStakingShareSeconds,
+            totalStakingShareSeconds,
             totalUserRewards,
             block.timestamp
         );
@@ -399,14 +399,14 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      * @return Total number of locked distribution tokens.
      */
     function totalLocked() public view override returns (uint256) {
-        return _lockedPool.balance();
+        return lockedPool.balance();
     }
 
     /**
      * @return Total number of unlocked distribution tokens.
      */
     function totalUnlocked() public view override returns (uint256) {
-        return _unlockedPool.balance();
+        return unlockedPool.balance();
     }
 
     /**
@@ -438,7 +438,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
 
         if (unlockedTokens > 0) {
             require(
-                _lockedPool.transfer(address(_unlockedPool), unlockedTokens),
+                lockedPool.transfer(address(unlockedPool), unlockedTokens),
                 "TokenGeyser: transfer out of locked pool failed"
             );
             emit TokensUnlocked(unlockedTokens, totalLocked());
@@ -459,7 +459,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      */
     function lockTokens(uint256 amount, uint256 durationSec) external onlyOwner {
         require(
-            unlockSchedules.length < _maxUnlockSchedules,
+            unlockSchedules.length < maxUnlockSchedules,
             "TokenGeyser: reached maximum unlock schedules"
         );
 
@@ -469,7 +469,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         uint256 lockedTokens = totalLocked();
         uint256 mintedLockedShares = (lockedTokens > 0)
             ? totalLockedShares.mul(amount).div(lockedTokens)
-            : amount.mul(_initialSharesPerToken);
+            : amount.mul(initialSharesPerToken);
 
         UnlockSchedule memory schedule;
         schedule.initialLockedShares = mintedLockedShares;
@@ -481,7 +481,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         totalLockedShares = totalLockedShares.add(mintedLockedShares);
 
         require(
-            _lockedPool.token().transferFrom(msg.sender, address(_lockedPool), amount),
+            lockedPool.token().transferFrom(msg.sender, address(lockedPool), amount),
             "TokenGeyser: transfer into locked pool failed"
         );
         emit TokensLocked(amount, durationSec, totalLocked());
@@ -499,7 +499,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         address to,
         uint256 amount
     ) public onlyOwner returns (bool) {
-        return _stakingPool.rescueFunds(tokenToRescue, to, amount);
+        return stakingPool.rescueFunds(tokenToRescue, to, amount);
     }
 
     //-------------------------------------------------------------------------

From c94981f63988ca57aa52c0687ec14230820b0373 Mon Sep 17 00:00:00 2001
From: aalavandhann <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 20:57:34 -0500
Subject: [PATCH 04/10] using clonable token pool and safe transfer

---
 contracts/ITokenPool.sol  | 15 +++++++++++
 contracts/TokenGeyser.sol | 56 +++++++++++++++++----------------------
 contracts/TokenPool.sol   | 27 ++++++++++++-------
 test/staking.ts           | 27 ++++++++++++++++---
 test/token_pool.ts        | 10 +++++--
 test/token_unlock.ts      | 18 ++++++++++---
 test/unstake.ts           |  4 +++
 7 files changed, 107 insertions(+), 50 deletions(-)
 create mode 100644 contracts/ITokenPool.sol

diff --git a/contracts/ITokenPool.sol b/contracts/ITokenPool.sol
new file mode 100644
index 0000000..67e5621
--- /dev/null
+++ b/contracts/ITokenPool.sol
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.24;
+
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+/**
+ * @title Token pool interface
+ */
+interface ITokenPool {
+    function init(IERC20 token_) external;
+    function token() external view returns (IERC20);
+    function balance() external view returns (uint256);
+    function transfer(address to, uint256 value) external;
+    function rescueFunds(address tokenToRescue, address to, uint256 amount) external;
+}
diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 9abd265..8c5d34b 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -2,8 +2,10 @@
 pragma solidity ^0.8.24;
 
 import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { Clones } from "@openzeppelin/contracts/proxy/Clones.sol";
+import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { SafeMathCompatibility } from "./_utils/SafeMathCompatibility.sol";
-import { TokenPool } from "./TokenPool.sol";
+import { ITokenPool } from "./ITokenPool.sol";
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { ITokenGeyser } from "./ITokenGeyser.sol";
@@ -28,6 +30,7 @@ import { ITokenGeyser } from "./ITokenGeyser.sol";
  */
 contract TokenGeyser is ITokenGeyser, Ownable {
     using SafeMathCompatibility for uint256;
+    using SafeERC20 for IERC20;
 
     //-------------------------------------------------------------------------
     // Events
@@ -42,9 +45,9 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     //-------------------------------------------------------------------------
     // Storage
 
-    TokenPool public stakingPool;
-    TokenPool public unlockedPool;
-    TokenPool public lockedPool;
+    ITokenPool public stakingPool;
+    ITokenPool public unlockedPool;
+    ITokenPool public lockedPool;
 
     //
     // Time-bonus params
@@ -112,6 +115,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      * @param initialSharesPerToken_ Number of shares to mint per staking token on first stake.
      */
     constructor(
+        address tokenPoolImpl,
         IERC20 stakingToken_,
         IERC20 distributionToken_,
         uint256 maxUnlockSchedules_,
@@ -126,10 +130,15 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         require(bonusPeriodSec_ != 0, "TokenGeyser: bonus period is zero");
         require(initialSharesPerToken_ > 0, "TokenGeyser: initialSharesPerToken is zero");
 
-        // TODO: use a factory here.
-        stakingPool = new TokenPool(stakingToken_);
-        unlockedPool = new TokenPool(distributionToken_);
-        lockedPool = new TokenPool(distributionToken_);
+        stakingPool = ITokenPool(Clones.clone(tokenPoolImpl));
+        stakingPool.init(stakingToken_);
+
+        unlockedPool = ITokenPool(Clones.clone(tokenPoolImpl));
+        unlockedPool.init(distributionToken_);
+
+        lockedPool = ITokenPool(Clones.clone(tokenPoolImpl));
+        lockedPool.init(distributionToken_);
+
         startBonus = startBonus_;
         bonusPeriodSec = bonusPeriodSec_;
         maxUnlockSchedules = maxUnlockSchedules_;
@@ -186,11 +195,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         // lastAccountingTimestampSec = block.timestamp;
 
         // interactions
-        require(
-            stakingPool.token().transferFrom(msg.sender, address(stakingPool), amount),
-            "TokenGeyser: transfer into staking pool failed"
-        );
-
+        stakingPool.token().safeTransferFrom(msg.sender, address(stakingPool), amount);
         emit Staked(msg.sender, amount, totalStakedFor(msg.sender));
     }
 
@@ -270,14 +275,8 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         // lastAccountingTimestampSec = block.timestamp;
 
         // interactions
-        require(
-            stakingPool.transfer(msg.sender, amount),
-            "TokenGeyser: transfer out of staking pool failed"
-        );
-        require(
-            unlockedPool.transfer(msg.sender, rewardAmount),
-            "TokenGeyser: transfer out of unlocked pool failed"
-        );
+        stakingPool.transfer(msg.sender, amount);
+        unlockedPool.transfer(msg.sender, rewardAmount);
 
         emit Unstaked(msg.sender, amount, totalStakedFor(msg.sender));
         emit TokensClaimed(msg.sender, rewardAmount);
@@ -437,10 +436,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         }
 
         if (unlockedTokens > 0) {
-            require(
-                lockedPool.transfer(address(unlockedPool), unlockedTokens),
-                "TokenGeyser: transfer out of locked pool failed"
-            );
+            lockedPool.transfer(address(unlockedPool), unlockedTokens);
             emit TokensUnlocked(unlockedTokens, totalLocked());
         }
 
@@ -480,10 +476,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
 
         totalLockedShares = totalLockedShares.add(mintedLockedShares);
 
-        require(
-            lockedPool.token().transferFrom(msg.sender, address(lockedPool), amount),
-            "TokenGeyser: transfer into locked pool failed"
-        );
+        lockedPool.token().safeTransferFrom(msg.sender, address(lockedPool), amount);
         emit TokensLocked(amount, durationSec, totalLocked());
     }
 
@@ -492,14 +485,13 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      * @param tokenToRescue Address of the token to be rescued.
      * @param to Address to which the rescued funds are to be sent.
      * @param amount Amount of tokens to be rescued.
-     * @return Transfer success.
      */
     function rescueFundsFromStakingPool(
         address tokenToRescue,
         address to,
         uint256 amount
-    ) public onlyOwner returns (bool) {
-        return stakingPool.rescueFunds(tokenToRescue, to, amount);
+    ) public onlyOwner {
+        stakingPool.rescueFunds(tokenToRescue, to, amount);
     }
 
     //-------------------------------------------------------------------------
diff --git a/contracts/TokenPool.sol b/contracts/TokenPool.sol
index 34ecd18..0bd5348 100644
--- a/contracts/TokenPool.sol
+++ b/contracts/TokenPool.sol
@@ -1,39 +1,48 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8.24;
 
-import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { ITokenPool } from "./ITokenPool.sol";
 
 /**
  * @title A simple holder of tokens.
  * This is a simple contract to hold tokens. It's useful in the case where a separate contract
  * needs to hold multiple distinct pools of the same token.
  */
-contract TokenPool is Ownable {
+contract TokenPool is ITokenPool, OwnableUpgradeable {
+    using SafeERC20 for IERC20;
     IERC20 public token;
 
-    constructor(IERC20 _token) Ownable(msg.sender) {
-        token = _token;
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
     }
 
-    function balance() public view returns (uint256) {
+    function init(IERC20 token_) public initializer {
+        __Ownable_init(msg.sender);
+        token = token_;
+    }
+
+    function balance() public view override returns (uint256) {
         return token.balanceOf(address(this));
     }
 
-    function transfer(address to, uint256 value) external onlyOwner returns (bool) {
-        return token.transfer(to, value);
+    function transfer(address to, uint256 value) external override onlyOwner {
+        token.safeTransfer(to, value);
     }
 
     function rescueFunds(
         address tokenToRescue,
         address to,
         uint256 amount
-    ) external onlyOwner returns (bool) {
+    ) external override onlyOwner {
         require(
             address(token) != tokenToRescue,
             "TokenPool: Cannot claim token held by the contract"
         );
 
-        return IERC20(tokenToRescue).transfer(to, amount);
+        IERC20(tokenToRescue).safeTransfer(to, amount);
     }
 }
diff --git a/test/staking.ts b/test/staking.ts
index cd362a2..37f8720 100644
--- a/test/staking.ts
+++ b/test/staking.ts
@@ -4,7 +4,11 @@ import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
 import { $AMPL, invokeRebase } from "../test/helper";
 import { SignerWithAddress } from "ethers";
 
-let ampl: any, dist: any, owner: SignerWithAddress, anotherAccount: SignerWithAddress;
+let ampl: any,
+  tokenPoolImpl: any,
+  dist: any,
+  owner: SignerWithAddress,
+  anotherAccount: SignerWithAddress;
 const InitialSharesPerToken = BigInt(10 ** 6);
 
 describe("staking", function () {
@@ -16,8 +20,12 @@ describe("staking", function () {
     await ampl.initialize(await owner.getAddress());
     await ampl.setMonetaryPolicy(await owner.getAddress());
 
+    const TokenPool = await ethers.getContractFactory("TokenPool");
+    const tokenPoolImpl = await TokenPool.deploy();
+
     const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
     dist = await TokenGeyser.deploy(
+      tokenPoolImpl.target,
       ampl.target,
       ampl.target,
       10,
@@ -26,11 +34,13 @@ describe("staking", function () {
       InitialSharesPerToken,
     );
 
-    return { ampl, dist, owner, anotherAccount };
+    return { ampl, tokenPoolImpl, dist, owner, anotherAccount };
   }
 
   beforeEach(async function () {
-    ({ ampl, dist, owner, anotherAccount } = await loadFixture(setupContracts));
+    ({ ampl, tokenPoolImpl, dist, owner, anotherAccount } = await loadFixture(
+      setupContracts,
+    ));
   });
 
   describe("when start bonus too high", function () {
@@ -38,6 +48,7 @@ describe("staking", function () {
       const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
       await expect(
         TokenGeyser.deploy(
+          tokenPoolImpl.target,
           ampl.target,
           ampl.target,
           10,
@@ -53,7 +64,15 @@ describe("staking", function () {
     it("should fail to construct", async function () {
       const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
       await expect(
-        TokenGeyser.deploy(ampl.target, ampl.target, 10, 50, 0, InitialSharesPerToken),
+        TokenGeyser.deploy(
+          tokenPoolImpl.target,
+          ampl.target,
+          ampl.target,
+          10,
+          50,
+          0,
+          InitialSharesPerToken,
+        ),
       ).to.be.revertedWith("TokenGeyser: bonus period is zero");
     });
   });
diff --git a/test/token_pool.ts b/test/token_pool.ts
index 4843ea5..08bb10e 100644
--- a/test/token_pool.ts
+++ b/test/token_pool.ts
@@ -1,4 +1,4 @@
-import { ethers } from "hardhat";
+import { ethers, upgrades } from "hardhat";
 import { expect } from "chai";
 import { SignerWithAddress } from "ethers";
 import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
@@ -14,7 +14,13 @@ describe("TokenPool", function () {
     const otherToken = await MockERC20.deploy(2000);
 
     const TokenPool = await ethers.getContractFactory("TokenPool");
-    const tokenPool = await TokenPool.deploy(token.target);
+    const tokenPool = await upgrades.deployProxy(
+      TokenPool.connect(owner),
+      [token.target],
+      {
+        initializer: "init(address)",
+      },
+    );
 
     return { token, otherToken, tokenPool, owner, anotherAccount };
   }
diff --git a/test/token_unlock.ts b/test/token_unlock.ts
index 459f67c..e739af1 100644
--- a/test/token_unlock.ts
+++ b/test/token_unlock.ts
@@ -10,7 +10,11 @@ import {
 } from "../test/helper";
 import { SignerWithAddress } from "ethers";
 
-let ampl: any, dist: any, owner: SignerWithAddress, anotherAccount: SignerWithAddress;
+let ampl: any,
+  tokenPoolImpl: any,
+  dist: any,
+  owner: SignerWithAddress,
+  anotherAccount: SignerWithAddress;
 const InitialSharesPerToken = BigInt(10 ** 6);
 const ONE_YEAR = 365 * 24 * 3600;
 const START_BONUS = 50;
@@ -24,8 +28,12 @@ async function setupContracts() {
   await ampl.initialize(await owner.getAddress());
   await ampl.setMonetaryPolicy(await owner.getAddress());
 
+  const TokenPool = await ethers.getContractFactory("TokenPool");
+  const tokenPoolImpl = await TokenPool.deploy();
+
   const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
   dist = await TokenGeyser.deploy(
+    tokenPoolImpl.target,
     ampl.target,
     ampl.target,
     10,
@@ -34,7 +42,7 @@ async function setupContracts() {
     InitialSharesPerToken,
   );
 
-  return { ampl, dist, owner, anotherAccount };
+  return { ampl, tokenPoolImpl, dist, owner, anotherAccount };
 }
 
 async function checkAvailableToUnlock(dist, v) {
@@ -46,7 +54,9 @@ async function checkAvailableToUnlock(dist, v) {
 
 describe("LockedPool", function () {
   beforeEach("setup contracts", async function () {
-    ({ ampl, dist, owner, anotherAccount } = await loadFixture(setupContracts));
+    ({ ampl, tokenPoolImpl, dist, owner, anotherAccount } = await loadFixture(
+      setupContracts,
+    ));
   });
 
   describe("distributionToken", function () {
@@ -60,6 +70,7 @@ describe("LockedPool", function () {
       it("should fail", async function () {
         const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
         const d = await TokenGeyser.deploy(
+          tokenPoolImpl.target,
           ampl.target,
           ampl.target,
           5n,
@@ -75,6 +86,7 @@ describe("LockedPool", function () {
       it("should fail", async function () {
         const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
         const d = await TokenGeyser.deploy(
+          tokenPoolImpl.target,
           ampl.target,
           ampl.target,
           5n,
diff --git a/test/unstake.ts b/test/unstake.ts
index cdbd06e..250c78a 100644
--- a/test/unstake.ts
+++ b/test/unstake.ts
@@ -16,10 +16,14 @@ async function setupContracts() {
   await ampl.initialize(await owner.getAddress());
   await ampl.setMonetaryPolicy(await owner.getAddress());
 
+  const TokenPool = await ethers.getContractFactory("TokenPool");
+  const tokenPoolImpl = await TokenPool.deploy();
+
   const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
   const startBonus = 50; // 50%
   const bonusPeriod = 86400; // 1 Day
   dist = await TokenGeyser.deploy(
+    tokenPoolImpl.target,
     ampl.target,
     ampl.target,
     10,

From 512e85301ad4449b5f490a683d0f039a51ea8fac Mon Sep 17 00:00:00 2001
From: aalavandhann <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 21:26:03 -0500
Subject: [PATCH 05/10] upgradable token geyser

---
 contracts/TokenGeyser.sol | 38 +++++++++++++++++++++++++-------------
 test/helper.ts            | 11 ++++++++++-
 test/staking.ts           | 23 +++++++++++++----------
 test/token_unlock.ts      | 16 +++++++---------
 test/unstake.ts           | 13 +++++++++----
 5 files changed, 64 insertions(+), 37 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 8c5d34b..ac9eb5a 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8.24;
 
-import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import { Clones } from "@openzeppelin/contracts/proxy/Clones.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { SafeMathCompatibility } from "./_utils/SafeMathCompatibility.sol";
@@ -28,7 +28,7 @@ import { ITokenGeyser } from "./ITokenGeyser.sol";
  *      More background and motivation available at:
  *      https://github.com/ampleforth/RFCs/blob/master/RFCs/rfc-1.md
  */
-contract TokenGeyser is ITokenGeyser, Ownable {
+contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
     using SafeMathCompatibility for uint256;
     using SafeERC20 for IERC20;
 
@@ -53,18 +53,18 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     // Time-bonus params
     //
     uint256 public constant BONUS_DECIMALS = 2;
-    uint256 public startBonus = 0;
-    uint256 public bonusPeriodSec = 0;
+    uint256 public startBonus;
+    uint256 public bonusPeriodSec;
 
     //
     // Global accounting state
     //
-    uint256 public totalLockedShares = 0;
-    uint256 public totalStakingShares = 0;
-    uint256 public totalStakingShareSeconds = 0;
-    uint256 public lastAccountingTimestampSec = block.timestamp;
-    uint256 public maxUnlockSchedules = 0;
-    uint256 public initialSharesPerToken = 0;
+    uint256 public totalLockedShares;
+    uint256 public totalStakingShares;
+    uint256 public totalStakingShareSeconds;
+    uint256 public lastAccountingTimestampSec;
+    uint256 public maxUnlockSchedules;
+    uint256 public initialSharesPerToken;
 
     //
     // User accounting state
@@ -103,7 +103,12 @@ contract TokenGeyser is ITokenGeyser, Ownable {
     UnlockSchedule[] public unlockSchedules;
 
     //-------------------------------------------------------------------------
-    // Constructor
+    // Construction
+
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
 
     /**
      * @param stakingToken_ The token users deposit as stake.
@@ -114,7 +119,7 @@ contract TokenGeyser is ITokenGeyser, Ownable {
      * @param bonusPeriodSec_ Length of time for bonus to increase linearly to max.
      * @param initialSharesPerToken_ Number of shares to mint per staking token on first stake.
      */
-    constructor(
+    function init(
         address tokenPoolImpl,
         IERC20 stakingToken_,
         IERC20 distributionToken_,
@@ -122,7 +127,9 @@ contract TokenGeyser is ITokenGeyser, Ownable {
         uint256 startBonus_,
         uint256 bonusPeriodSec_,
         uint256 initialSharesPerToken_
-    ) Ownable(msg.sender) {
+    ) public initializer {
+        __Ownable_init(msg.sender);
+
         // The start bonus must be some fraction of the max. (i.e. <= 100%)
         require(startBonus_ <= 10 ** BONUS_DECIMALS, "TokenGeyser: start bonus too high");
         // If no period is desired, instead set startBonus = 100%
@@ -141,6 +148,11 @@ contract TokenGeyser is ITokenGeyser, Ownable {
 
         startBonus = startBonus_;
         bonusPeriodSec = bonusPeriodSec_;
+
+        totalLockedShares = 0;
+        totalStakingShares = 0;
+        totalStakingShareSeconds = 0;
+        lastAccountingTimestampSec = block.timestamp;
         maxUnlockSchedules = maxUnlockSchedules_;
         initialSharesPerToken = initialSharesPerToken_;
     }
diff --git a/test/helper.ts b/test/helper.ts
index 61633a6..ceb5c54 100644
--- a/test/helper.ts
+++ b/test/helper.ts
@@ -1,4 +1,4 @@
-import hre, { ethers } from "hardhat";
+import hre, { ethers, upgrades } from "hardhat";
 import { expect } from "chai";
 
 const AMPL_DECIMALS = 9;
@@ -80,6 +80,14 @@ async function printStatus(dist) {
   // dist.unlockSchedules.staticCall(1)
 }
 
+async function deployGeyser(owner, params) {
+  const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
+  const dist = await upgrades.deployProxy(TokenGeyser.connect(owner), params, {
+    initializer: "init(address,address,address,uint256,uint256,uint256,uint256)",
+  });
+  return dist;
+}
+
 module.exports = {
   checkAmplAprox,
   checkSharesAprox,
@@ -88,4 +96,5 @@ module.exports = {
   TimeHelpers,
   printMethodOutput,
   printStatus,
+  deployGeyser,
 };
diff --git a/test/staking.ts b/test/staking.ts
index 37f8720..6a75081 100644
--- a/test/staking.ts
+++ b/test/staking.ts
@@ -1,7 +1,7 @@
 import { ethers } from "hardhat";
 import { expect } from "chai";
 import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
-import { $AMPL, invokeRebase } from "../test/helper";
+import { $AMPL, invokeRebase, deployGeyser } from "../test/helper";
 import { SignerWithAddress } from "ethers";
 
 let ampl: any,
@@ -23,8 +23,7 @@ describe("staking", function () {
     const TokenPool = await ethers.getContractFactory("TokenPool");
     const tokenPoolImpl = await TokenPool.deploy();
 
-    const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
-    dist = await TokenGeyser.deploy(
+    dist = await deployGeyser(owner, [
       tokenPoolImpl.target,
       ampl.target,
       ampl.target,
@@ -32,7 +31,7 @@ describe("staking", function () {
       50,
       86400,
       InitialSharesPerToken,
-    );
+    ]);
 
     return { ampl, tokenPoolImpl, dist, owner, anotherAccount };
   }
@@ -45,9 +44,8 @@ describe("staking", function () {
 
   describe("when start bonus too high", function () {
     it("should fail to construct", async function () {
-      const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
       await expect(
-        TokenGeyser.deploy(
+        deployGeyser(owner, [
           tokenPoolImpl.target,
           ampl.target,
           ampl.target,
@@ -55,16 +53,15 @@ describe("staking", function () {
           101,
           86400,
           InitialSharesPerToken,
-        ),
+        ]),
       ).to.be.revertedWith("TokenGeyser: start bonus too high");
     });
   });
 
   describe("when bonus period is 0", function () {
     it("should fail to construct", async function () {
-      const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
       await expect(
-        TokenGeyser.deploy(
+        deployGeyser(owner, [
           tokenPoolImpl.target,
           ampl.target,
           ampl.target,
@@ -72,11 +69,17 @@ describe("staking", function () {
           50,
           0,
           InitialSharesPerToken,
-        ),
+        ]),
       ).to.be.revertedWith("TokenGeyser: bonus period is zero");
     });
   });
 
+  describe("owner", function () {
+    it("should return the owner", async function () {
+      expect(await dist.owner()).to.equal(await owner.getAddress());
+    });
+  });
+
   describe("stakingToken", function () {
     it("should return the staking token", async function () {
       expect(await dist.stakingToken()).to.equal(ampl.target);
diff --git a/test/token_unlock.ts b/test/token_unlock.ts
index e739af1..b717931 100644
--- a/test/token_unlock.ts
+++ b/test/token_unlock.ts
@@ -7,6 +7,7 @@ import {
   invokeRebase,
   checkAmplAprox,
   checkSharesAprox,
+  deployGeyser,
 } from "../test/helper";
 import { SignerWithAddress } from "ethers";
 
@@ -31,8 +32,7 @@ async function setupContracts() {
   const TokenPool = await ethers.getContractFactory("TokenPool");
   const tokenPoolImpl = await TokenPool.deploy();
 
-  const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
-  dist = await TokenGeyser.deploy(
+  dist = await deployGeyser(owner, [
     tokenPoolImpl.target,
     ampl.target,
     ampl.target,
@@ -40,7 +40,7 @@ async function setupContracts() {
     START_BONUS,
     BONUS_PERIOD,
     InitialSharesPerToken,
-  );
+  ]);
 
   return { ampl, tokenPoolImpl, dist, owner, anotherAccount };
 }
@@ -68,8 +68,7 @@ describe("LockedPool", function () {
   describe("lockTokens", function () {
     describe("when not approved", function () {
       it("should fail", async function () {
-        const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
-        const d = await TokenGeyser.deploy(
+        const d = await deployGeyser(owner, [
           tokenPoolImpl.target,
           ampl.target,
           ampl.target,
@@ -77,15 +76,14 @@ describe("LockedPool", function () {
           START_BONUS,
           BONUS_PERIOD,
           InitialSharesPerToken,
-        );
+        ]);
         await expect(d.lockTokens($AMPL(10), ONE_YEAR)).to.be.reverted;
       });
     });
 
     describe("when number of unlock schedules exceeds the maxUnlockSchedules", function () {
       it("should fail", async function () {
-        const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
-        const d = await TokenGeyser.deploy(
+        const d = await deployGeyser(owner, [
           tokenPoolImpl.target,
           ampl.target,
           ampl.target,
@@ -93,7 +91,7 @@ describe("LockedPool", function () {
           START_BONUS,
           BONUS_PERIOD,
           InitialSharesPerToken,
-        );
+        ]);
         await ampl.approve(d.target, $AMPL(100));
         for (let i = 0; i < 5; i++) {
           await d.lockTokens($AMPL(10), ONE_YEAR);
diff --git a/test/unstake.ts b/test/unstake.ts
index 250c78a..1e35afd 100644
--- a/test/unstake.ts
+++ b/test/unstake.ts
@@ -1,7 +1,13 @@
 import { ethers } from "hardhat";
 import { expect } from "chai";
 import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers";
-import { $AMPL, invokeRebase, checkAmplAprox, TimeHelpers } from "../test/helper";
+import {
+  $AMPL,
+  invokeRebase,
+  checkAmplAprox,
+  TimeHelpers,
+  deployGeyser,
+} from "../test/helper";
 import { SignerWithAddress } from "ethers";
 
 let ampl: any, dist: any, owner: SignerWithAddress, anotherAccount: SignerWithAddress;
@@ -19,10 +25,9 @@ async function setupContracts() {
   const TokenPool = await ethers.getContractFactory("TokenPool");
   const tokenPoolImpl = await TokenPool.deploy();
 
-  const TokenGeyser = await ethers.getContractFactory("TokenGeyser");
   const startBonus = 50; // 50%
   const bonusPeriod = 86400; // 1 Day
-  dist = await TokenGeyser.deploy(
+  dist = await deployGeyser(owner, [
     tokenPoolImpl.target,
     ampl.target,
     ampl.target,
@@ -30,7 +35,7 @@ async function setupContracts() {
     startBonus,
     bonusPeriod,
     InitialSharesPerToken,
-  );
+  ]);
 
   await ampl.transfer(await anotherAccount.getAddress(), $AMPL(50000));
   await ampl.connect(anotherAccount).approve(dist.target, $AMPL(50000));

From d3b664bb3e3ad4831e19ec8a33894fd59a55fd29 Mon Sep 17 00:00:00 2001
From: aalavandhann <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 12 Nov 2024 21:44:13 -0500
Subject: [PATCH 06/10] added pausability

---
 contracts/TokenGeyser.sol | 26 +++++++++++++++----
 test/staking.ts           | 53 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index ac9eb5a..491148f 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -2,6 +2,7 @@
 pragma solidity ^0.8.24;
 
 import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import { PausableUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
 import { Clones } from "@openzeppelin/contracts/proxy/Clones.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { SafeMathCompatibility } from "./_utils/SafeMathCompatibility.sol";
@@ -28,7 +29,7 @@ import { ITokenGeyser } from "./ITokenGeyser.sol";
  *      More background and motivation available at:
  *      https://github.com/ampleforth/RFCs/blob/master/RFCs/rfc-1.md
  */
-contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
+contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     using SafeMathCompatibility for uint256;
     using SafeERC20 for IERC20;
 
@@ -129,6 +130,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
         uint256 initialSharesPerToken_
     ) public initializer {
         __Ownable_init(msg.sender);
+        __Pausable_init();
 
         // The start bonus must be some fraction of the max. (i.e. <= 100%)
         require(startBonus_ <= 10 ** BONUS_DECIMALS, "TokenGeyser: start bonus too high");
@@ -179,7 +181,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
      * @dev Transfers amount of deposit tokens from the user.
      * @param amount Number of deposit tokens to stake.
      */
-    function stake(uint256 amount) external {
+    function stake(uint256 amount) external whenNotPaused {
         require(amount > 0, "TokenGeyser: stake amount is zero");
         require(
             totalStakingShares == 0 || totalStaked() > 0,
@@ -216,7 +218,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
      * alotted number of distribution tokens.
      * @param amount Number of deposit tokens to unstake / withdraw.
      */
-    function unstake(uint256 amount) external returns (uint256) {
+    function unstake(uint256 amount) external whenNotPaused returns (uint256) {
         updateAccounting();
 
         // checks
@@ -367,6 +369,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
      */
     function updateAccounting()
         public
+        whenNotPaused
         returns (uint256, uint256, uint256, uint256, uint256, uint256)
     {
         unlockTokens();
@@ -432,7 +435,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
      *      previously defined unlock schedules. Publicly callable.
      * @return Number of newly unlocked distribution tokens.
      */
-    function unlockTokens() public returns (uint256) {
+    function unlockTokens() public whenNotPaused returns (uint256) {
         uint256 unlockedTokens = 0;
         uint256 lockedTokens = totalLocked();
 
@@ -458,6 +461,16 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
     //-------------------------------------------------------------------------
     // Admin only methods
 
+    /// @notice Pauses all user interactions.
+    function pause() external onlyOwner {
+        _pause();
+    }
+
+    /// @notice Unpauses all user interactions.
+    function unpause() external onlyOwner {
+        _unpause();
+    }
+
     /**
      * @dev This funcion allows the contract owner to add more locked distribution tokens, along
      *      with the associated "unlock schedule". These locked tokens immediately begin unlocking
@@ -465,7 +478,10 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable {
      * @param amount Number of distribution tokens to lock. These are transferred from the caller.
      * @param durationSec Length of time to linear unlock the tokens.
      */
-    function lockTokens(uint256 amount, uint256 durationSec) external onlyOwner {
+    function lockTokens(
+        uint256 amount,
+        uint256 durationSec
+    ) external onlyOwner {
         require(
             unlockSchedules.length < maxUnlockSchedules,
             "TokenGeyser: reached maximum unlock schedules"
diff --git a/test/staking.ts b/test/staking.ts
index 6a75081..7da5e5b 100644
--- a/test/staking.ts
+++ b/test/staking.ts
@@ -74,6 +74,59 @@ describe("staking", function () {
     });
   });
 
+  describe("#pause", function () {
+    describe("when triggered by non-owner", function () {
+      it("should revert", async function () {
+        await dist.transferOwnership(anotherAccount);
+        await expect(dist.pause()).to.be.revertedWithCustomError(
+          dist,
+          "OwnableUnauthorizedAccount",
+        );
+      });
+    });
+
+    describe("when already paused", function () {
+      it("should revert", async function () {
+        await dist.pause();
+        await expect(dist.pause()).to.be.revertedWithCustomError(dist, "EnforcedPause");
+      });
+    });
+
+    describe("when valid", function () {
+      it("should pause", async function () {
+        await dist.pause();
+        expect(await dist.paused()).to.eq(true);
+      });
+    });
+  });
+
+  describe("#unpause", function () {
+    describe("when triggered by non-owner", function () {
+      it("should revert", async function () {
+        await dist.pause();
+        await dist.transferOwnership(anotherAccount);
+        await expect(dist.unpause()).to.be.revertedWithCustomError(
+          dist,
+          "OwnableUnauthorizedAccount",
+        );
+      });
+    });
+
+    describe("when not paused", function () {
+      it("should revert", async function () {
+        await expect(dist.unpause()).to.be.revertedWithCustomError(dist, "ExpectedPause");
+      });
+    });
+
+    describe("when valid", function () {
+      it("should unpause", async function () {
+        await dist.pause();
+        await dist.unpause();
+        expect(await dist.paused()).to.eq(false);
+      });
+    });
+  });
+
   describe("owner", function () {
     it("should return the owner", async function () {
       expect(await dist.owner()).to.equal(await owner.getAddress());

From 9e7528b7cb38c4a08966c96b2050f0701d8b96b4 Mon Sep 17 00:00:00 2001
From: aalavandhann <6264334+aalavandhan@users.noreply.github.com>
Date: Wed, 13 Nov 2024 07:59:58 -0500
Subject: [PATCH 07/10] added reentrant checks

---
 contracts/TokenGeyser.sol | 176 ++++++++++++++++++++++----------------
 1 file changed, 100 insertions(+), 76 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 491148f..1342c69 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -3,6 +3,7 @@ pragma solidity ^0.8.24;
 
 import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import { PausableUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
+import { ReentrancyGuardUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
 import { Clones } from "@openzeppelin/contracts/proxy/Clones.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { SafeMathCompatibility } from "./_utils/SafeMathCompatibility.sol";
@@ -29,7 +30,12 @@ import { ITokenGeyser } from "./ITokenGeyser.sol";
  *      More background and motivation available at:
  *      https://github.com/ampleforth/RFCs/blob/master/RFCs/rfc-1.md
  */
-contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
+contract TokenGeyser is
+    ITokenGeyser,
+    OwnableUpgradeable,
+    PausableUpgradeable,
+    ReentrancyGuardUpgradeable
+{
     using SafeMathCompatibility for uint256;
     using SafeERC20 for IERC20;
 
@@ -131,11 +137,12 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     ) public initializer {
         __Ownable_init(msg.sender);
         __Pausable_init();
+        __ReentrancyGuard_init();
 
         // The start bonus must be some fraction of the max. (i.e. <= 100%)
         require(startBonus_ <= 10 ** BONUS_DECIMALS, "TokenGeyser: start bonus too high");
         // If no period is desired, instead set startBonus = 100%
-        // and bonusPeriod to a small value like 1sec.
+        // and bonusPeriod to a small value like 1 sec.
         require(bonusPeriodSec_ != 0, "TokenGeyser: bonus period is zero");
         require(initialSharesPerToken_ > 0, "TokenGeyser: initialSharesPerToken is zero");
 
@@ -178,10 +185,10 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     }
 
     /**
-     * @dev Transfers amount of deposit tokens from the user.
+     * @notice Transfers amount of deposit tokens from the user.
      * @param amount Number of deposit tokens to stake.
      */
-    function stake(uint256 amount) external whenNotPaused {
+    function stake(uint256 amount) external nonReentrant whenNotPaused {
         require(amount > 0, "TokenGeyser: stake amount is zero");
         require(
             totalStakingShares == 0 || totalStaked() > 0,
@@ -193,7 +200,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
             : amount.mul(initialSharesPerToken);
         require(mintedStakingShares > 0, "TokenGeyser: Stake amount is too small");
 
-        updateAccounting();
+        _updateAccounting();
 
         // 1. User Accounting
         UserTotals storage totals = userTotals[msg.sender];
@@ -205,7 +212,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
 
         // 2. Global Accounting
         totalStakingShares = totalStakingShares.add(mintedStakingShares);
-        // Already set in updateAccounting()
+        // Already set in _updateAccounting()
         // lastAccountingTimestampSec = block.timestamp;
 
         // interactions
@@ -214,12 +221,14 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     }
 
     /**
-     * @dev Unstakes a certain amount of previously deposited tokens. User also receives their
+     * @notice Unstakes a certain amount of previously deposited tokens. User also receives their
      * alotted number of distribution tokens.
      * @param amount Number of deposit tokens to unstake / withdraw.
      */
-    function unstake(uint256 amount) external whenNotPaused returns (uint256) {
-        updateAccounting();
+    function unstake(
+        uint256 amount
+    ) external nonReentrant whenNotPaused returns (uint256) {
+        _updateAccounting();
 
         // checks
         require(amount > 0, "TokenGeyser: unstake amount is zero");
@@ -303,7 +312,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     }
 
     /**
-     * @dev Applies an additional time-bonus to a distribution amount. This is necessary to
+     * @notice Applies an additional time-bonus to a distribution amount. This is necessary to
      *      encourage long-term deposits instead of constant unstake/restakes.
      *      The bonus-multiplier is the result of a linear function that starts at startBonus and
      *      ends at 100% over bonusPeriodSec, then stays at 100% thereafter.
@@ -358,7 +367,7 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     }
 
     /**
-     * @dev A globally callable function to update the accounting state of the system.
+     * @notice A globally callable function to update the accounting state of the system.
      *      Global state and state for the caller are updated.
      * @return [0] balance of the locked pool
      * @return [1] balance of the unlocked pool
@@ -368,45 +377,12 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
      * @return [5] block timestamp
      */
     function updateAccounting()
-        public
+        external
+        nonReentrant
         whenNotPaused
         returns (uint256, uint256, uint256, uint256, uint256, uint256)
     {
-        unlockTokens();
-
-        // Global accounting
-        uint256 newStakingShareSeconds = block
-            .timestamp
-            .sub(lastAccountingTimestampSec)
-            .mul(totalStakingShares);
-        totalStakingShareSeconds = totalStakingShareSeconds.add(newStakingShareSeconds);
-        lastAccountingTimestampSec = block.timestamp;
-
-        // User Accounting
-        UserTotals storage totals = userTotals[msg.sender];
-        uint256 newUserStakingShareSeconds = block
-            .timestamp
-            .sub(totals.lastAccountingTimestampSec)
-            .mul(totals.stakingShares);
-        totals.stakingShareSeconds = totals.stakingShareSeconds.add(
-            newUserStakingShareSeconds
-        );
-        totals.lastAccountingTimestampSec = block.timestamp;
-
-        uint256 totalUserRewards = (totalStakingShareSeconds > 0)
-            ? totalUnlocked().mul(totals.stakingShareSeconds).div(
-                totalStakingShareSeconds
-            )
-            : 0;
-
-        return (
-            totalLocked(),
-            totalUnlocked(),
-            totals.stakingShareSeconds,
-            totalStakingShareSeconds,
-            totalUserRewards,
-            block.timestamp
-        );
+        return _updateAccounting();
     }
 
     /**
@@ -426,36 +402,17 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
     /**
      * @return Number of unlock schedules.
      */
-    function unlockScheduleCount() public view returns (uint256) {
+    function unlockScheduleCount() external view returns (uint256) {
         return unlockSchedules.length;
     }
 
     /**
-     * @dev Moves distribution tokens from the locked pool to the unlocked pool, according to the
+     * @notice Moves distribution tokens from the locked pool to the unlocked pool, according to the
      *      previously defined unlock schedules. Publicly callable.
      * @return Number of newly unlocked distribution tokens.
      */
-    function unlockTokens() public whenNotPaused returns (uint256) {
-        uint256 unlockedTokens = 0;
-        uint256 lockedTokens = totalLocked();
-
-        if (totalLockedShares == 0) {
-            unlockedTokens = lockedTokens;
-        } else {
-            uint256 unlockedShares = 0;
-            for (uint256 s = 0; s < unlockSchedules.length; s++) {
-                unlockedShares = unlockedShares.add(_unlockScheduleShares(s));
-            }
-            unlockedTokens = unlockedShares.mul(lockedTokens).div(totalLockedShares);
-            totalLockedShares = totalLockedShares.sub(unlockedShares);
-        }
-
-        if (unlockedTokens > 0) {
-            lockedPool.transfer(address(unlockedPool), unlockedTokens);
-            emit TokensUnlocked(unlockedTokens, totalLocked());
-        }
-
-        return unlockedTokens;
+    function unlockTokens() external nonReentrant whenNotPaused returns (uint256) {
+        return _unlockTokens();
     }
 
     //-------------------------------------------------------------------------
@@ -478,17 +435,14 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
      * @param amount Number of distribution tokens to lock. These are transferred from the caller.
      * @param durationSec Length of time to linear unlock the tokens.
      */
-    function lockTokens(
-        uint256 amount,
-        uint256 durationSec
-    ) external onlyOwner {
+    function lockTokens(uint256 amount, uint256 durationSec) external onlyOwner {
         require(
             unlockSchedules.length < maxUnlockSchedules,
             "TokenGeyser: reached maximum unlock schedules"
         );
 
         // Update lockedTokens amount before using it in computations after.
-        updateAccounting();
+        _updateAccounting();
 
         uint256 lockedTokens = totalLocked();
         uint256 mintedLockedShares = (lockedTokens > 0)
@@ -518,13 +472,83 @@ contract TokenGeyser is ITokenGeyser, OwnableUpgradeable, PausableUpgradeable {
         address tokenToRescue,
         address to,
         uint256 amount
-    ) public onlyOwner {
+    ) external onlyOwner {
         stakingPool.rescueFunds(tokenToRescue, to, amount);
     }
 
     //-------------------------------------------------------------------------
     // Private methods
 
+    /**
+     * @dev Updates time-dependent global storage state.
+     */
+    function _updateAccounting()
+        private
+        returns (uint256, uint256, uint256, uint256, uint256, uint256)
+    {
+        _unlockTokens();
+
+        // Global accounting
+        uint256 newStakingShareSeconds = block
+            .timestamp
+            .sub(lastAccountingTimestampSec)
+            .mul(totalStakingShares);
+        totalStakingShareSeconds = totalStakingShareSeconds.add(newStakingShareSeconds);
+        lastAccountingTimestampSec = block.timestamp;
+
+        // User Accounting
+        UserTotals storage totals = userTotals[msg.sender];
+        uint256 newUserStakingShareSeconds = block
+            .timestamp
+            .sub(totals.lastAccountingTimestampSec)
+            .mul(totals.stakingShares);
+        totals.stakingShareSeconds = totals.stakingShareSeconds.add(
+            newUserStakingShareSeconds
+        );
+        totals.lastAccountingTimestampSec = block.timestamp;
+
+        uint256 totalUserRewards = (totalStakingShareSeconds > 0)
+            ? totalUnlocked().mul(totals.stakingShareSeconds).div(
+                totalStakingShareSeconds
+            )
+            : 0;
+
+        return (
+            totalLocked(),
+            totalUnlocked(),
+            totals.stakingShareSeconds,
+            totalStakingShareSeconds,
+            totalUserRewards,
+            block.timestamp
+        );
+    }
+
+    /**
+     * @dev Unlocks distribution tokens based on reward schedule.
+     */
+    function _unlockTokens() private returns (uint256) {
+        uint256 unlockedTokens = 0;
+        uint256 lockedTokens = totalLocked();
+
+        if (totalLockedShares == 0) {
+            unlockedTokens = lockedTokens;
+        } else {
+            uint256 unlockedShares = 0;
+            for (uint256 s = 0; s < unlockSchedules.length; s++) {
+                unlockedShares = unlockedShares.add(_unlockScheduleShares(s));
+            }
+            unlockedTokens = unlockedShares.mul(lockedTokens).div(totalLockedShares);
+            totalLockedShares = totalLockedShares.sub(unlockedShares);
+        }
+
+        if (unlockedTokens > 0) {
+            lockedPool.transfer(address(unlockedPool), unlockedTokens);
+            emit TokensUnlocked(unlockedTokens, totalLocked());
+        }
+
+        return unlockedTokens;
+    }
+
     /**
      * @dev Returns the number of unlockable shares from a given schedule. The returned value
      *      depends on the time since the last unlock. This function updates schedule accounting,

From b8e2cb8363b34343ab266bb185e61ebd4dcd2093 Mon Sep 17 00:00:00 2001
From: aalavandhann <6264334+aalavandhan@users.noreply.github.com>
Date: Wed, 13 Nov 2024 13:29:33 -0500
Subject: [PATCH 08/10] typo

---
 contracts/TokenGeyser.sol | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 1342c69..90016b5 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -222,7 +222,7 @@ contract TokenGeyser is
 
     /**
      * @notice Unstakes a certain amount of previously deposited tokens. User also receives their
-     * alotted number of distribution tokens.
+     * allotted number of distribution tokens.
      * @param amount Number of deposit tokens to unstake / withdraw.
      */
     function unstake(
@@ -348,7 +348,7 @@ contract TokenGeyser is
 
     /**
      * @param addr The user to look up staking information for.
-     * @return The number of staking tokens deposited for addr.
+     * @return The number of staking tokens deposited for address.
      */
     function totalStakedFor(address addr) public view returns (uint256) {
         return
@@ -429,9 +429,9 @@ contract TokenGeyser is
     }
 
     /**
-     * @dev This funcion allows the contract owner to add more locked distribution tokens, along
+     * @dev This function allows the contract owner to add more locked distribution tokens, along
      *      with the associated "unlock schedule". These locked tokens immediately begin unlocking
-     *      linearly over the duraction of durationSec timeframe.
+     *      linearly over the duration of durationSec time frame.
      * @param amount Number of distribution tokens to lock. These are transferred from the caller.
      * @param durationSec Length of time to linear unlock the tokens.
      */
@@ -550,7 +550,7 @@ contract TokenGeyser is
     }
 
     /**
-     * @dev Returns the number of unlockable shares from a given schedule. The returned value
+     * @dev Returns the number of unlock-able shares from a given schedule. The returned value
      *      depends on the time since the last unlock. This function updates schedule accounting,
      *      but does not actually transfer any tokens.
      * @param s Index of the unlock schedule.

From 03e4bec747b8bb37924ddeb144e1f60e2e6e3792 Mon Sep 17 00:00:00 2001
From: aalavandhan1984 <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 19 Nov 2024 12:17:15 -0500
Subject: [PATCH 09/10] renamed user totals to user

---
 contracts/TokenGeyser.sol | 30 ++++++++++++++----------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 90016b5..1dd044b 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -203,9 +203,9 @@ contract TokenGeyser is
         _updateAccounting();
 
         // 1. User Accounting
-        UserTotals storage totals = userTotals[msg.sender];
-        totals.stakingShares = totals.stakingShares.add(mintedStakingShares);
-        totals.lastAccountingTimestampSec = block.timestamp;
+        UserTotals storage user = userTotals[msg.sender];
+        user.stakingShares = user.stakingShares.add(mintedStakingShares);
+        user.lastAccountingTimestampSec = block.timestamp;
 
         Stake memory newStake = Stake(mintedStakingShares, block.timestamp);
         userStakes[msg.sender].push(newStake);
@@ -243,7 +243,7 @@ contract TokenGeyser is
         );
 
         // 1. User Accounting
-        UserTotals storage totals = userTotals[msg.sender];
+        UserTotals storage user = userTotals[msg.sender];
         Stake[] storage accountStakes = userStakes[msg.sender];
 
         // Redeem from most recent stake and go backwards in time.
@@ -282,12 +282,12 @@ contract TokenGeyser is
                 sharesLeftToBurn = 0;
             }
         }
-        totals.stakingShareSeconds = totals.stakingShareSeconds.sub(
+        user.stakingShareSeconds = user.stakingShareSeconds.sub(
             stakingShareSecondsToBurn
         );
-        totals.stakingShares = totals.stakingShares.sub(stakingSharesToBurn);
+        user.stakingShares = user.stakingShares.sub(stakingSharesToBurn);
         // Already set in updateAccounting
-        // totals.lastAccountingTimestampSec = block.timestamp;
+        // user.lastAccountingTimestampSec = block.timestamp;
 
         // 2. Global Accounting
         totalStakingShareSeconds = totalStakingShareSeconds.sub(
@@ -497,26 +497,24 @@ contract TokenGeyser is
         lastAccountingTimestampSec = block.timestamp;
 
         // User Accounting
-        UserTotals storage totals = userTotals[msg.sender];
+        UserTotals storage user = userTotals[msg.sender];
         uint256 newUserStakingShareSeconds = block
             .timestamp
-            .sub(totals.lastAccountingTimestampSec)
-            .mul(totals.stakingShares);
-        totals.stakingShareSeconds = totals.stakingShareSeconds.add(
+            .sub(user.lastAccountingTimestampSec)
+            .mul(user.stakingShares);
+        user.stakingShareSeconds = user.stakingShareSeconds.add(
             newUserStakingShareSeconds
         );
-        totals.lastAccountingTimestampSec = block.timestamp;
+        user.lastAccountingTimestampSec = block.timestamp;
 
         uint256 totalUserRewards = (totalStakingShareSeconds > 0)
-            ? totalUnlocked().mul(totals.stakingShareSeconds).div(
-                totalStakingShareSeconds
-            )
+            ? totalUnlocked().mul(user.stakingShareSeconds).div(totalStakingShareSeconds)
             : 0;
 
         return (
             totalLocked(),
             totalUnlocked(),
-            totals.stakingShareSeconds,
+            user.stakingShareSeconds,
             totalStakingShareSeconds,
             totalUserRewards,
             block.timestamp

From 0249dfa42378b5509b50ef4cef1ee18f2341cb0e Mon Sep 17 00:00:00 2001
From: aalavandhan1984 <6264334+aalavandhan@users.noreply.github.com>
Date: Tue, 19 Nov 2024 12:18:33 -0500
Subject: [PATCH 10/10] renamed totalStakedFor to totalStakedBy

---
 contracts/ITokenGeyser.sol |  2 +-
 contracts/TokenGeyser.sol  |  8 ++++----
 test/staking.ts            | 14 +++++++-------
 test/unstake.ts            | 18 +++++++++---------
 4 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/contracts/ITokenGeyser.sol b/contracts/ITokenGeyser.sol
index 8bfb847..febc7cd 100644
--- a/contracts/ITokenGeyser.sol
+++ b/contracts/ITokenGeyser.sol
@@ -9,7 +9,7 @@ import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 interface ITokenGeyser {
     function stake(uint256 amount) external;
     function unstake(uint256 amount) external returns (uint256);
-    function totalStakedFor(address addr) external view returns (uint256);
+    function totalStakedBy(address addr) external view returns (uint256);
     function totalStaked() external view returns (uint256);
     function totalLocked() external view returns (uint256);
     function totalUnlocked() external view returns (uint256);
diff --git a/contracts/TokenGeyser.sol b/contracts/TokenGeyser.sol
index 1dd044b..b47df57 100644
--- a/contracts/TokenGeyser.sol
+++ b/contracts/TokenGeyser.sol
@@ -217,7 +217,7 @@ contract TokenGeyser is
 
         // interactions
         stakingPool.token().safeTransferFrom(msg.sender, address(stakingPool), amount);
-        emit Staked(msg.sender, amount, totalStakedFor(msg.sender));
+        emit Staked(msg.sender, amount, totalStakedBy(msg.sender));
     }
 
     /**
@@ -233,7 +233,7 @@ contract TokenGeyser is
         // checks
         require(amount > 0, "TokenGeyser: unstake amount is zero");
         require(
-            totalStakedFor(msg.sender) >= amount,
+            totalStakedBy(msg.sender) >= amount,
             "TokenGeyser: unstake amount is greater than total user stakes"
         );
         uint256 stakingSharesToBurn = totalStakingShares.mul(amount).div(totalStaked());
@@ -301,7 +301,7 @@ contract TokenGeyser is
         stakingPool.transfer(msg.sender, amount);
         unlockedPool.transfer(msg.sender, rewardAmount);
 
-        emit Unstaked(msg.sender, amount, totalStakedFor(msg.sender));
+        emit Unstaked(msg.sender, amount, totalStakedBy(msg.sender));
         emit TokensClaimed(msg.sender, rewardAmount);
 
         require(
@@ -350,7 +350,7 @@ contract TokenGeyser is
      * @param addr The user to look up staking information for.
      * @return The number of staking tokens deposited for address.
      */
-    function totalStakedFor(address addr) public view returns (uint256) {
+    function totalStakedBy(address addr) public view returns (uint256) {
         return
             totalStakingShares > 0
                 ? totalStaked().mul(userTotals[addr].stakingShares).div(
diff --git a/test/staking.ts b/test/staking.ts
index 7da5e5b..042ae8e 100644
--- a/test/staking.ts
+++ b/test/staking.ts
@@ -163,7 +163,7 @@ describe("staking", function () {
       it("should update the total staked", async function () {
         await dist.stake($AMPL(100));
         expect(await dist.totalStaked()).to.equal($AMPL(100));
-        expect(await dist.totalStakedFor(await owner.getAddress())).to.equal($AMPL(100));
+        expect(await dist.totalStakedBy(await owner.getAddress())).to.equal($AMPL(100));
         expect(await dist.totalStakingShares()).to.equal(
           $AMPL(100) * InitialSharesPerToken,
         );
@@ -187,10 +187,10 @@ describe("staking", function () {
       });
       it("should update the total staked", async function () {
         expect(await dist.totalStaked()).to.equal($AMPL(200));
-        expect(await dist.totalStakedFor(await anotherAccount.getAddress())).to.equal(
+        expect(await dist.totalStakedBy(await anotherAccount.getAddress())).to.equal(
           $AMPL(50),
         );
-        expect(await dist.totalStakedFor(await owner.getAddress())).to.equal($AMPL(150));
+        expect(await dist.totalStakedBy(await owner.getAddress())).to.equal($AMPL(150));
         expect(await dist.totalStakingShares()).to.equal(
           $AMPL(200) * InitialSharesPerToken,
         );
@@ -210,10 +210,10 @@ describe("staking", function () {
       });
       it("should updated the total staked shares", async function () {
         expect(await dist.totalStaked()).to.equal($AMPL(250));
-        expect(await dist.totalStakedFor(await anotherAccount.getAddress())).to.equal(
+        expect(await dist.totalStakedBy(await anotherAccount.getAddress())).to.equal(
           $AMPL(100),
         );
-        expect(await dist.totalStakedFor(await owner.getAddress())).to.equal($AMPL(150));
+        expect(await dist.totalStakedBy(await owner.getAddress())).to.equal($AMPL(150));
         expect(await dist.totalStakingShares()).to.equal(
           $AMPL(125) * InitialSharesPerToken,
         );
@@ -246,10 +246,10 @@ describe("staking", function () {
       });
       it("should updated the total staked shares", async function () {
         expect(await dist.totalStaked()).to.equal($AMPL(175));
-        expect(await dist.totalStakedFor(await anotherAccount.getAddress())).to.equal(
+        expect(await dist.totalStakedBy(await anotherAccount.getAddress())).to.equal(
           $AMPL(25),
         );
-        expect(await dist.totalStakedFor(await owner.getAddress())).to.equal($AMPL(150));
+        expect(await dist.totalStakedBy(await owner.getAddress())).to.equal($AMPL(150));
         expect(await dist.totalStakingShares()).to.equal(
           $AMPL(350) * InitialSharesPerToken,
         );
diff --git a/test/unstake.ts b/test/unstake.ts
index 1e35afd..cf389a7 100644
--- a/test/unstake.ts
+++ b/test/unstake.ts
@@ -126,7 +126,7 @@ describe("unstaking", function () {
         await dist.connect(anotherAccount).unstake($AMPL(30));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(20));
         expect(
-          await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
+          await dist.totalStakedBy.staticCall(await anotherAccount.getAddress()),
         ).to.eq($AMPL(20));
         checkAmplAprox(await totalRewardsFor(anotherAccount), 40);
       });
@@ -173,7 +173,7 @@ describe("unstaking", function () {
         await dist.connect(anotherAccount).unstake($AMPL(250));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(250));
         expect(
-          await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
+          await dist.totalStakedBy.staticCall(await anotherAccount.getAddress()),
         ).to.eq($AMPL(250));
         checkAmplAprox(await totalRewardsFor(anotherAccount), 625); // (.5 * .75 * 1000) + 250
       });
@@ -223,7 +223,7 @@ describe("unstaking", function () {
         await dist.connect(anotherAccount).unstake($AMPL(30));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(70));
         expect(
-          await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
+          await dist.totalStakedBy.staticCall(await anotherAccount.getAddress()),
         ).to.eq($AMPL(70));
         checkAmplAprox(await totalRewardsFor(anotherAccount), 40.8);
       });
@@ -303,9 +303,9 @@ describe("unstaking", function () {
         await dist.connect(anotherAccount).unstake($AMPL(30));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(70));
         expect(
-          await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
+          await dist.totalStakedBy.staticCall(await anotherAccount.getAddress()),
         ).to.eq($AMPL(20));
-        expect(await dist.totalStakedFor.staticCall(await owner.getAddress())).to.eq(
+        expect(await dist.totalStakedBy.staticCall(await owner.getAddress())).to.eq(
           $AMPL(50),
         );
         checkAmplAprox(await totalRewardsFor(anotherAccount), 18.24);
@@ -348,8 +348,8 @@ describe("unstaking", function () {
       it("should update the total staked and rewards", async function () {
         await dist.connect(anotherAccount).unstake($AMPL(10000));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(8000));
-        expect(await dist.totalStakedFor.staticCall(ethers.ZeroAddress)).to.eq($AMPL(0));
-        expect(await dist.totalStakedFor.staticCall(await owner.getAddress())).to.eq(
+        expect(await dist.totalStakedBy.staticCall(ethers.ZeroAddress)).to.eq($AMPL(0));
+        expect(await dist.totalStakedBy.staticCall(await owner.getAddress())).to.eq(
           $AMPL(8000),
         );
         checkAmplAprox(await totalRewardsFor(anotherAccount), 0);
@@ -357,9 +357,9 @@ describe("unstaking", function () {
         await dist.unstake($AMPL(8000));
         expect(await dist.totalStaked.staticCall()).to.eq($AMPL(0));
         expect(
-          await dist.totalStakedFor.staticCall(await anotherAccount.getAddress()),
+          await dist.totalStakedBy.staticCall(await anotherAccount.getAddress()),
         ).to.eq($AMPL(0));
-        expect(await dist.totalStakedFor.staticCall(await owner.getAddress())).to.eq(
+        expect(await dist.totalStakedBy.staticCall(await owner.getAddress())).to.eq(
           $AMPL(0),
         );
         checkAmplAprox(await totalRewardsFor(anotherAccount), 0);