feat: max deposit (#99)

* feat: max uint deposit * fix: invariant test * fix: rebase * fix: dont max mint * fix: uint 256 Co-authored-by: FP <83050944+fp-crypto@users.noreply.github.com> * fix: uint 256 Co-authored-by: FP <83050944+fp-crypto@users.noreply.github.com> * test: check balances --------- Co-authored-by: FP <83050944+fp-crypto@users.noreply.github.com>
yearn · Sep 9, 2024 · bbabc72 · bbabc72
1 parent 59ece7f
commit bbabc72
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 0 deletions.
diff --git a/src/TokenizedStrategy.sol b/src/TokenizedStrategy.sol
@@ -497,6 +497,12 @@ contract TokenizedStrategy {
     ) external nonReentrant returns (uint256 shares) {
         // Get the storage slot for all following calls.
         StrategyData storage S = _strategyStorage();
+
+        // Deposit full balance if using max uint.
+        if (assets == type(uint256).max) {
+            assets = S.asset.balanceOf(msg.sender);
+        }
+
         // Checking max deposit will also check if shutdown.
         require(
             assets <= _maxDeposit(S, receiver),
@@ -524,6 +530,7 @@ contract TokenizedStrategy {
     ) external nonReentrant returns (uint256 assets) {
         // Get the storage slot for all following calls.
         StrategyData storage S = _strategyStorage();
+
         // Checking max mint will also check if shutdown.
         require(shares <= _maxMint(S, receiver), "ERC4626: mint more than max");
         // Check for rounding error.

diff --git a/src/test/Accounting.t.sol b/src/test/Accounting.t.sol
@@ -582,4 +582,35 @@ contract AccountingTest is Setup {
         assertEq(strategy.pricePerShare(), wad);
         checkStrategyTotals(strategy, 0, 0, 0, 0);
     }
+
+    function test_maxUintDeposit_depositsBalance(
+        address _address,
+        uint256 _amount
+    ) public {
+        _amount = bound(_amount, minFuzzAmount, maxFuzzAmount);
+        vm.assume(
+            _address != address(0) &&
+                _address != address(strategy) &&
+                _address != address(yieldSource)
+        );
+
+        asset.mint(_address, _amount);
+
+        vm.prank(_address);
+        asset.approve(address(strategy), _amount);
+
+        assertEq(asset.balanceOf(_address), _amount);
+
+        vm.prank(_address);
+        strategy.deposit(type(uint256).max, _address);
+
+        // Should just deposit the available amount.
+        checkStrategyTotals(strategy, _amount, _amount, 0, _amount);
+
+        assertEq(asset.balanceOf(_address), 0);
+        assertEq(strategy.balanceOf(_address), _amount);
+        assertEq(asset.balanceOf(address(strategy)), 0);
+
+        assertEq(asset.balanceOf(address(yieldSource)), _amount);
+    }
 }
diff --git a/src/test/ERC4626Std.t.sol b/src/test/ERC4626Std.t.sol
@@ -15,4 +15,22 @@ contract ERC4626StdTest is ERC4626Test, Setup {
         _vaultMayBeEmpty = true;
         _unlimitedAmount = true;
     }
+
+    //Avoid special case for deposits of uint256 max
+    function test_previewDeposit(
+        Init memory init,
+        uint256 assets
+    ) public override {
+        if (assets == type(uint256).max) assets -= 1;
+        super.test_previewDeposit(init, assets);
+    }
+
+    function test_deposit(
+        Init memory init,
+        uint256 assets,
+        uint256 allowance
+    ) public override {
+        if (assets == type(uint256).max) assets -= 1;
+        super.test_deposit(init, assets, allowance);
+    }
 }