Feat: update wrap

src/Pool/Modules/PoolEuler.sol

@@ -66,7 +66,7 @@ contract PoolEuler is Pool {
         return IEToken(address(sharesToken)).convertBalanceToUnderlying(1e18);
     }
 
-    /// Returns the shares balance TODO: lots of notes
+    /// Returns the shares balance
     /// The decimals of the shares amount returned is adjusted to match the decimals of the baseToken
     function _getSharesBalance() internal view virtual override returns (uint104) {
         return (sharesToken.balanceOf(address(this)) / scaleFactor).u104();
@@ -75,14 +75,13 @@ contract PoolEuler is Pool {
     /// Internal function for wrapping base asset tokens.
     /// @param receiver The address the wrapped tokens should be sent.
     /// @return shares The amount of wrapped tokens that are sent to the receiver.
-    function _wrap(address receiver) internal virtual override returns (uint256 shares) {
-        uint256 baseOut = baseToken.balanceOf(address(this));
-        if (baseOut == 0) return 0;
+    function _wrap(uint256 assets, address receiver) internal virtual override returns (uint256 shares) {
+        if (assets == 0) return 0;
 
-        IEToken(address(sharesToken)).deposit(0, baseOut); // first param is subaccount, 0 for primary
-        uint256 sharesReceived = _getSharesBalance() - sharesCached;  // this includes any shares in pool previously
+        IEToken(address(sharesToken)).deposit(0, assets); // first param is subaccount, 0 for primary
+        shares = _getSharesBalance() - sharesCached;  // this includes any shares in pool previously
         if (receiver != address(this)) {
-            sharesToken.safeTransfer(receiver, sharesReceived);
+            sharesToken.safeTransfer(receiver, shares);
         }
     }
 
@@ -96,11 +95,16 @@ contract PoolEuler is Pool {
     /// Internal function for unwrapping unaccounted for base in this contract.
     /// @param receiver The address the wrapped tokens should be sent.
     /// @return assets The amount of assets sent to the receiver in native decimals.
-    function _unwrap(address receiver) internal virtual override returns (uint256 assets) {
+    function _unwrap(uint256 shares, address receiver) internal virtual override returns (uint256 assets) {
+        if (shares == 0) return 0;
+
         uint256 surplus = _getSharesBalance() - sharesCached;
-        if (surplus == 0) return 0;
+        if (shares > surplus) {
+            revert CannotUnwrapMoreThanSurplus(shares, surplus);
+        }
+
         // convert to base
-        assets = _unwrapPreview(surplus);
+        assets = _unwrapPreview(shares);
         IEToken(address(sharesToken)).withdraw(0, assets); // first param is subaccount, 0 for primary
 
         if (receiver != address(this)) {

src/Pool/Modules/PoolNonTv.sol

@@ -95,10 +95,15 @@ contract PoolNonTv is Pool {
     /// Internal function for wrapping base asset tokens.
     /// Since there is nothing to unwrap, we return the surplus balance.
     /// @return shares The amount of wrapped tokens that are sent to the receiver.
-    function _wrap(address receiver) internal virtual override returns (uint256 shares) {
-        shares = _getSharesBalance() - sharesCached;
+    function _wrap(uint256 assets, address receiver) internal virtual override returns (uint256 shares) {
+        uint256 surplus = _getSharesBalance() - sharesCached;
+
+        if (shares > surplus) {
+            revert CannotWrapMoreThanSurplus(assets, surplus);
+        }
+
         if (receiver != address(this)) {
-            sharesToken.safeTransfer(receiver, shares);
+            sharesToken.safeTransfer(receiver, assets);
         }
     }
 
@@ -112,8 +117,12 @@ contract PoolNonTv is Pool {
     /// Internal function for unwrapping unaccounted for base in this contract.
     /// Since there is nothing to unwrap, we return the surplus balance.
     /// @return assets The amount of base assets sent to the receiver.
-    function _unwrap(address receiver) internal virtual override returns (uint256 assets) {
-        assets = _getSharesBalance() - sharesCached;
+    function _unwrap(uint256 shares, address receiver) internal virtual override returns (uint256 assets) {
+        uint256 surplus = _getSharesBalance() - sharesCached;
+        if (shares > surplus) {
+            revert CannotUnwrapMoreThanSurplus(shares, surplus);
+        }
+
         if (receiver != address(this)) {
             sharesToken.safeTransfer(receiver, assets);
         }

src/Pool/Modules/PoolYearnVault.sol

@@ -56,10 +56,9 @@ contract PoolYearnVault is Pool {
     /// Internal function for wrapping base tokens.
     /// @param receiver The address the wrapped tokens should be sent.
     /// @return shares The amount of wrapped tokens that are sent to the receiver.
-    function _wrap(address receiver) internal virtual override returns (uint256 shares) {
-        uint256 baseOut = baseToken.balanceOf(address(this));
-        if (baseOut == 0) return 0;
-        shares = IYVToken(address(sharesToken)).deposit(baseOut, receiver);
+    function _wrap(uint256 assets, address receiver) internal virtual override returns (uint256 shares) {
+        if (assets == 0) return 0;
+        shares = IYVToken(address(sharesToken)).deposit(assets, receiver);
     }
 
     /// Internal function to preview how many shares will be received when depositing a given amount of base.
@@ -73,10 +72,14 @@ contract PoolYearnVault is Pool {
     /// Internal function for unwrapping unaccounted for base in this contract.
     /// @param receiver The address the wrapped tokens should be sent.
     /// @return base_ The amount of base base sent to the receiver.
-    function _unwrap(address receiver) internal virtual override returns (uint256 base_) {
+    function _unwrap(uint256 shares, address receiver) internal virtual override returns (uint256 base_) {
+        if (shares == 0) return 0;
+
         uint256 surplus = _getSharesBalance() - sharesCached;
-        if (surplus == 0) return 0;
-        base_ = IYVToken(address(sharesToken)).withdraw(surplus, receiver);
+        if (shares > surplus) {
+            revert CannotUnwrapMoreThanSurplus(shares, surplus);
+        }
+        base_ = IYVToken(address(sharesToken)).withdraw(shares, receiver);
     }
 
     /// Internal function to preview how many base tokens will be received when unwrapping a given amount of shares.

src/Pool/Pool.sol

@@ -390,7 +390,7 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
         // Wrap all base found in this contract.
         baseIn = baseToken.balanceOf(address(this));
 
-        _wrap(address(this));
+        _wrap(baseIn, address(this));
 
         // Gather data
         uint256 supply = _totalSupply;
@@ -454,7 +454,7 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
         _mint(to, lpTokensMinted);
 
         // Return any unused base tokens
-        if (sharesBalance > cache.sharesCached + sharesIn) _unwrap(remainder);
+        if (sharesBalance > cache.sharesCached + sharesIn) _unwrap(baseToken.balanceOf(address(this)), remainder);
 
         // confirm new virtual fyToken balance is not less than new supply
         if ((cache.fyTokenCached + fyTokenIn + lpTokensMinted) < supply + lpTokensMinted) {
@@ -627,7 +627,7 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
 
         // Burn and transfer
         _burn(address(this), lpTokensBurned); // This is calling the actual ERC20 _burn.
-        baseOut = _unwrap(baseTo);
+        baseOut = _unwrap(sharesOut, baseTo);
 
         if (fyTokenOut != 0) fyToken.safeTransfer(fyTokenTo, fyTokenOut);
 
@@ -719,7 +719,7 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
         );
 
         // Transfer
-        _unwrap(to);
+        _unwrap(sharesOut, to);
 
         emit Trade(maturity, msg.sender, to, baseOut.i128(), -(fyTokenIn.i128()));
     }
@@ -802,19 +802,21 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
         uint128 fyTokenOut,
         uint128 max
     ) external virtual override returns (uint128 baseIn) {
-        // Wrap any base assets found in contract.
-        _wrap(address(this));
 
         // Calculate trade
-        uint128 sharesBalance = _getSharesBalance();
         Cache memory cache = _getCache();
         uint128 sharesIn = _buyFYTokenPreview(
             fyTokenOut,
             cache.sharesCached,
             cache.fyTokenCached,
             _computeG1(cache.g1Fee)
         );
+
+        // convert base to shares
         baseIn = _unwrapPreview(sharesIn).u128();
+        _wrap(baseIn, address(this));
+
+        uint128 sharesBalance = _getSharesBalance();
 
         // Checks
         if (sharesBalance - cache.sharesCached < sharesIn)
@@ -922,7 +924,8 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
     /// @return fyTokenOut Amount of fyToken that will be deposited on `to` wallet.
     function sellBase(address to, uint128 min) external virtual override returns (uint128 fyTokenOut) {
         // Wrap any base assets found in contract.
-        _wrap(address(this));
+        uint256 baseIn = baseToken.balanceOf(address(this));
+        _wrap(baseIn, address(this));
 
         // Calculate trade
         Cache memory cache = _getCache();
@@ -944,7 +947,7 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
             revert FYTokenCachedBadState();
         }
 
-        emit Trade(maturity, msg.sender, to, -(_unwrapPreview(sharesIn).u128().i128()), fyTokenOut.i128());
+        emit Trade(maturity, msg.sender, to, -(baseIn.u128().i128()), fyTokenOut.i128());
     }
 
     /// Returns how much fyToken would be obtained by selling `baseIn`.
@@ -1041,7 +1044,7 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
         _update(cache.sharesCached - sharesOut, fyTokenBalance, cache.sharesCached, cache.fyTokenCached);
 
         // Transfer
-        baseOut = _unwrap(to).u128();
+        baseOut = _unwrap(baseToken.balanceOf(address(this)), to).u128();
 
         // Check slippage
         if (baseOut < min) revert SlippageDuringSellFYToken(baseOut, min);
@@ -1087,15 +1090,15 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
     /// @param receiver The address to which the wrapped tokens will be sent.
     /// @return shares The amount of wrapped tokens sent to the receiver.
     function wrap(address receiver) external returns (uint256 shares) {
-        shares = _wrap(receiver);
+        uint256 assets = baseToken.balanceOf(address(this));
+        shares = _wrap(assets, receiver);
     }
 
     /// Internal function for wrapping base tokens whichwraps the entire balance of base found in this contract.
     /// @dev This should be overridden by modules.
     /// @param receiver The address the wrapped tokens should be sent.
     /// @return shares The amount of wrapped tokens that are sent to the receiver.
-    function _wrap(address receiver) internal virtual returns (uint256 shares) {
-        uint256 assets = baseToken.balanceOf(address(this));
+    function _wrap(uint256 assets, address receiver) internal virtual returns (uint256 shares) {
         if (assets == 0) {
             shares = 0;
         } else {
@@ -1127,15 +1130,19 @@ contract Pool is PoolEvents, IPool, ERC20Permit, AccessControl {
     /// @param receiver The address to which the assets will be sent.
     /// @return assets The amount of asset tokens sent to the receiver.
     function unwrap(address receiver) external returns (uint256 assets) {
-        assets = _unwrap(receiver);
+        uint256 surplus = _getSharesBalance() - sharesCached;
+        assets = _unwrap(surplus, receiver);
     }
 
     /// Internal function for unwrapping unaccounted for base in this contract.
     /// @dev This should be overridden by modules.
     /// @param receiver The address the wrapped tokens should be sent.
     /// @return assets The amount of base assets sent to the receiver.
-    function _unwrap(address receiver) internal virtual returns (uint256 assets) {
+    function _unwrap(uint256 amount, address receiver) internal virtual returns (uint256 assets) {
         uint256 surplus = _getSharesBalance() - sharesCached;
+        if (amount > surplus) {
+            revert CannotUnwrapMoreThanSurplus(amount, surplus);
+        }
         if (surplus == 0) {
             assets = 0;
         } else {

src/Pool/PoolErrors.sol

@@ -10,6 +10,12 @@ error AfterMaturity();
 /// The approval of the sharesToken failed miserably.
 error ApproveFailed();
 
+/// Cannot unwrap more than the surplus over the reserves. Wrap your head around that one.
+error CannotUnwrapMoreThanSurplus(uint256 amount, uint256 surplus);
+
+/// For Non-TV pools, cannot wrap more than the surplus over the reserves. No exceptions!
+error CannotWrapMoreThanSurplus(uint256 amount, uint256 surplus);
+
 /// The update would cause the FYToken cached to be less than the total supply. This should never happen but may
 /// occur due to unexpected rounding errors.  We cannot allow this to happen as it could have many unexpected and
 /// side effects which may pierce the fabric of the space-time continuum.

src/YieldMath.sol

@@ -491,15 +491,15 @@ library YieldMath {
         return uint128(a);
     }
 
-    /// Calculate a YieldSpace pool invariant according to the whitepaper
+    /// Calculate a YieldSpace pool hoagies according to the whitepaper
     /// @dev Implemented using base reserves and uint128 to be backwards compatible with yieldspace-v2
     /// @param baseReserves base reserve amount
     /// @param fyTokenReserves fyToken reserves amount
     /// @param totalSupply pool token total amount
     /// @param timeTillMaturity time till maturity in seconds e.g. 90 days in seconds
     /// @param k time till maturity coefficient, multiplied by 2^64.  e.g. 25 years in seconds
-    /// @return result the invariant value
-    function invariant(uint128 baseReserves, uint128 fyTokenReserves, uint256 totalSupply, uint128 timeTillMaturity, int128 k)
+    /// @return result the hoagies value
+    function hoagies(uint128 baseReserves, uint128 fyTokenReserves, uint256 totalSupply, uint128 timeTillMaturity, int128 k)
         public pure returns(uint128 result)
     {
         if (totalSupply == 0) return 0;

src/YieldMathExtensions.sol

@@ -7,11 +7,11 @@ import "./YieldMath.sol";
 
 library YieldMathExtensions {
 
-    /// @dev Calculate the invariant for this pool
-    function invariant(IPool pool) external view returns (uint128) {
+    /// @dev Calculate the hoagies for this pool
+    function hoagies(IPool pool) external view returns (uint128) {
         uint32 maturity = pool.maturity();
         uint32 timeToMaturity = (maturity > uint32(block.timestamp)) ? maturity - uint32(block.timestamp) : 0;
-        return YieldMath.invariant(
+        return YieldMath.hoagies(
             pool.getBaseBalance(),
             pool.getFYTokenBalance(),
             pool.totalSupply(),