Merge pull request #144 from morpho-dao/refactor/move-adding-to-colla…

…teral-list Refactor/move adding to collateral list
morpho-org · Jan 14, 2023 · 7b06160 · 7b06160
2 parents d5cacec + 95901e4
commit 7b06160
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 17 deletions.
diff --git a/src/MorphoInternal.sol b/src/MorphoInternal.sol
@@ -333,6 +333,8 @@ abstract contract MorphoInternal is MorphoStorage {
             onPool,
             inP2P
         );
+        if (onPool == 0 && inP2P == 0) _userBorrows[user].remove(underlying);
+        else _userBorrows[user].add(underlying);
     }
 
     function _setPauseStatus(address underlying, bool isPaused) internal {

diff --git a/src/PositionsManager.sol b/src/PositionsManager.sol
@@ -15,6 +15,7 @@ import {PercentageMath} from "@morpho-utils/math/PercentageMath.sol";
 
 import {Permit2Lib} from "./libraries/Permit2Lib.sol";
 import {ERC20, SafeTransferLib} from "@solmate/utils/SafeTransferLib.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 
 import {MorphoStorage} from "./MorphoStorage.sol";
 import {PositionsManagerInternal} from "./PositionsManagerInternal.sol";
@@ -24,6 +25,7 @@ contract PositionsManager is IPositionsManager, PositionsManagerInternal {
     using Permit2Lib for ERC20;
     using SafeTransferLib for ERC20;
     using MarketBalanceLib for Types.MarketBalances;
+    using EnumerableSet for EnumerableSet.AddressSet;
 
     using WadRayMath for uint256;
     using PercentageMath for uint256;
@@ -66,6 +68,7 @@ contract PositionsManager is IPositionsManager, PositionsManagerInternal {
         ERC20(underlying).transferFrom2(from, address(this), amount);
 
         _marketBalances[underlying].collateral[onBehalf] += amount.rayDiv(indexes.supply.poolIndex);
+        _userCollaterals[onBehalf].add(underlying);
 
         _POOL.supplyToPool(underlying, amount);
 
@@ -131,7 +134,9 @@ contract PositionsManager is IPositionsManager, PositionsManagerInternal {
         // The following check requires storage indexes to be up-to-date.
         _validateWithdrawCollateral(underlying, amount, supplier);
 
-        _marketBalances[underlying].collateral[supplier] -= amount.rayDiv(indexes.supply.poolIndex);
+        uint256 newBalance = _marketBalances[underlying].collateral[supplier] - amount.rayDiv(indexes.supply.poolIndex);
+        _marketBalances[underlying].collateral[supplier] = newBalance;
+        if (newBalance == 0) _userCollaterals[supplier].remove(underlying);
 
         _POOL.withdrawFromPool(underlying, _market[underlying].aToken, amount);
         ERC20(underlying).safeTransfer(receiver, amount);

diff --git a/src/PositionsManagerInternal.sol b/src/PositionsManagerInternal.sol
@@ -182,8 +182,6 @@ abstract contract PositionsManagerInternal is MatchingEngine {
         onPool = marketBalances.scaledPoolSupplyBalance(user);
         inP2P = marketBalances.scaledP2PSupplyBalance(user);
 
-        _userCollaterals[user].add(underlying);
-
         /// Peer-to-peer supply ///
 
         // Match the peer-to-peer borrow delta.
@@ -236,7 +234,6 @@ abstract contract PositionsManagerInternal is MatchingEngine {
         Types.MarketBalances storage marketBalances = _marketBalances[underlying];
         Types.Deltas storage deltas = market.deltas;
 
-        _userBorrows[user].add(underlying);
         vars.onPool = marketBalances.scaledPoolBorrowBalance(user);
         vars.inP2P = marketBalances.scaledP2PBorrowBalance(user);
 
@@ -319,11 +316,6 @@ abstract contract PositionsManagerInternal is MatchingEngine {
 
             if (amount == 0) {
                 _updateBorrowerInDS(underlying, user, onPool, inP2P);
-
-                if (inP2P == 0 && onPool == 0) {
-                    _userBorrows[user].remove(underlying);
-                }
-
                 return (onPool, inP2P, 0, toRepay);
             }
         }
@@ -392,8 +384,6 @@ abstract contract PositionsManagerInternal is MatchingEngine {
             /// Note: Only used in breaking repay. Suppliers should not be able to supply if the pool is supply capped.
             toSupply = _handleSupplyCap(underlying, amount);
         }
-
-        if (inP2P == 0 && onPool == 0) _userBorrows[user].remove(underlying);
     }
 
     function _executeWithdraw(
@@ -421,10 +411,6 @@ abstract contract PositionsManagerInternal is MatchingEngine {
             if (amount == 0) {
                 _updateSupplierInDS(underlying, user, vars.onPool, vars.inP2P);
 
-                if (vars.inP2P == 0 && vars.onPool == 0) {
-                    _userCollaterals[user].remove(underlying);
-                }
-
                 return vars;
             }
         }
@@ -480,8 +466,6 @@ abstract contract PositionsManagerInternal is MatchingEngine {
             emit Events.P2PAmountsUpdated(underlying, deltas.p2pSupplyAmount, deltas.p2pBorrowAmount);
             vars.toBorrow = amount;
         }
-
-        if (vars.inP2P == 0 && vars.onPool == 0) _userCollaterals[user].remove(underlying);
     }
 
     function _calculateAmountToSeize(