[Loopring AMM] Add more tests and fix a few issues (#1835)

Author: Brechtpd

packages/loopring_v3/contracts/amm/libamm/AmmBlockReceiver.sol

@@ -42,7 +42,7 @@ library AmmBlockReceiver
         S.approveAmmUpdates(ctx, false);
 
         // Update state
-        S.poolTokenBurnedSupply = ctx.poolTokenBurnedSupply;
+        S._totalSupply = ctx.totalSupply;
 
         return ctx.txIdx - txIdx;
     }
@@ -65,7 +65,7 @@ library AmmBlockReceiver
             domainSeparator: S.domainSeparator,
             accountID: S.accountID,
             poolTokenID: S.poolTokenID,
-            poolTokenBurnedSupply: S.poolTokenBurnedSupply,
+            totalSupply: S._totalSupply,
             size: size,
             tokens: S.tokens,
             tokenBalancesL2: new uint96[](size)

packages/loopring_v3/contracts/amm/libamm/AmmData.sol

@@ -81,7 +81,7 @@ library AmmData
         uint32  accountID;
 
         uint16  poolTokenID;
-        uint    poolTokenBurnedSupply;
+        uint    totalSupply;
 
         uint     size; // == token.length;
         Token[]  tokens;
@@ -92,7 +92,7 @@ library AmmData
         // Pool token state variables
         string poolName;
         string symbol;
-        uint poolTokenBurnedSupply;
+        uint _totalSupply;
 
         mapping(address => uint) balanceOf;
         mapping(address => mapping(address => uint)) allowance;

packages/loopring_v3/contracts/amm/libamm/AmmExitProcess.sol

@@ -65,7 +65,7 @@ library AmmExitProcess
                 return;
             }
 
-            ctx.poolTokenBurnedSupply = ctx.poolTokenBurnedSupply.add(exit.burnAmount);
+            ctx.totalSupply = ctx.totalSupply.sub(exit.burnAmount);
         } else {
             require(slippageOK, "EXIT_SLIPPAGE_INVALID");
             _burnL2(ctx, exit.burnAmount, exit.owner, exit.burnStorageID);
@@ -124,29 +124,24 @@ library AmmExitProcess
         ctx.approveTransfer(transfer);
 
         // Update pool balance
-        ctx.poolTokenBurnedSupply = ctx.poolTokenBurnedSupply.add(transfer.amount);
+        ctx.totalSupply = ctx.totalSupply.sub(transfer.amount);
     }
 
     function _calculateExitAmounts(
         AmmData.Context  memory  ctx,
         AmmData.PoolExit memory  exit
         )
         private
-        view
+        pure
         returns(
             bool /* slippageOK */,
             uint96[] memory amounts
         )
     {
         amounts = new uint96[](ctx.size);
 
-        // Check if we can still use this exit
-        if (block.timestamp > exit.validUntil) {
-            return (false, amounts);
-        }
-
         // Calculate how much will be withdrawn
-        uint ratio = uint(AmmData.POOL_TOKEN_BASE()).mul(exit.burnAmount) / ctx.totalSupply();
+        uint ratio = uint(AmmData.POOL_TOKEN_BASE()).mul(exit.burnAmount) / ctx.totalSupply;
 
         for (uint i = 0; i < ctx.size; i++) {
             amounts[i] = (ratio.mul(ctx.tokenBalancesL2[i]) / AmmData.POOL_TOKEN_BASE()).toUint96();

packages/loopring_v3/contracts/amm/libamm/AmmJoinProcess.sol

@@ -99,15 +99,15 @@ library AmmJoinProcess
         ctx.approveTransfer(transfer);
 
         // Update pool balance
-        ctx.poolTokenBurnedSupply = ctx.poolTokenBurnedSupply.sub(transfer.amount);
+        ctx.totalSupply = ctx.totalSupply.add(transfer.amount);
     }
 
     function _calculateJoinAmounts(
         AmmData.Context  memory  ctx,
         AmmData.PoolJoin memory  join
         )
         private
-        view
+        pure
         returns(
             bool            slippageOK,
             uint96          mintAmount,
@@ -117,11 +117,7 @@ library AmmJoinProcess
         // Check if we can still use this join
         amounts = new uint96[](ctx.size);
 
-        if (block.timestamp > join.validUntil) {
-            return (false, 0, amounts);
-        }
-
-        if (ctx.totalSupply() == 0) {
+        if (ctx.totalSupply == 0) {
             return(true, AmmData.POOL_TOKEN_BASE().toUint96(), join.joinAmounts);
         }
 
@@ -130,7 +126,7 @@ library AmmJoinProcess
         for (uint i = 0; i < ctx.size; i++) {
             if (ctx.tokenBalancesL2[i] > 0) {
                 uint amountOut = uint(join.joinAmounts[i])
-                    .mul(ctx.totalSupply()) / uint(ctx.tokenBalancesL2[i]);
+                    .mul(ctx.totalSupply) / uint(ctx.tokenBalancesL2[i]);
 
                 if (!initialized) {
                     initialized = true;
@@ -146,10 +142,10 @@ library AmmJoinProcess
         }
 
         // Calculate the amounts to deposit
-        uint ratio = uint(AmmData.POOL_TOKEN_BASE()).mul(mintAmount) / ctx.totalSupply();
+        uint ratio = uint(AmmData.POOL_TOKEN_BASE()).mul(mintAmount) / ctx.totalSupply;
 
         for (uint i = 0; i < ctx.size; i++) {
-            amounts[i] = ratio.mul(ctx.tokenBalancesL2[i] / AmmData.POOL_TOKEN_BASE()).toUint96();
+            amounts[i] = (ratio.mul(ctx.tokenBalancesL2[i]) / AmmData.POOL_TOKEN_BASE()).toUint96();
         }
 
         slippageOK = (mintAmount >= join.mintMinAmount);

packages/loopring_v3/contracts/amm/libamm/AmmPoolToken.sol

@@ -28,7 +28,7 @@ library AmmPoolToken
         view
         returns (uint)
     {
-        return AmmData.POOL_TOKEN_MINTED_SUPPLY().sub(S.poolTokenBurnedSupply);
+        return S._totalSupply;
     }
 
     function approve(

packages/loopring_v3/contracts/amm/libamm/AmmStatus.sol

@@ -73,7 +73,6 @@ library AmmStatus
 
         // Mint all liquidity tokens to the pool account on L2
         S.balanceOf[address(this)] = AmmData.POOL_TOKEN_MINTED_SUPPLY();
-        S.poolTokenBurnedSupply = AmmData.POOL_TOKEN_MINTED_SUPPLY();
         S.allowance[address(this)][address(exchange.getDepositContract())] = uint(-1);
         exchange.deposit(
             address(this), // from
@@ -91,11 +90,13 @@ library AmmStatus
         )
         internal
     {
-        uint64 validUntil = S.forcedExit[exitOwner].validUntil;
-        require(validUntil > 0 && validUntil < block.timestamp, "INVALID_CHALLENGE");
-
-        uint size = S.tokens.length;
+        // If the exchange is in withdrawal mode allow the pool to be shutdown immediately
         if (!S.exchange.isInWithdrawalMode()) {
+            uint64 validUntil = S.forcedExit[exitOwner].validUntil;
+            require(validUntil > 0 && validUntil < block.timestamp, "INVALID_CHALLENGE");
+
+            uint size = S.tokens.length;
+
             for (uint i = 0; i < size; i++) {
                 S.exchange.forceWithdraw{value: msg.value / size}(
                     address(this),

packages/loopring_v3/contracts/amm/libamm/AmmUtil.sol

@@ -28,16 +28,6 @@ library AmmUtil
         ctx.exchange.approveTransaction(transfer.from, hash);
     }
 
-    function totalSupply(
-        AmmData.Context  memory  ctx
-        )
-        internal
-        pure
-        returns (uint)
-    {
-        return AmmData.POOL_TOKEN_MINTED_SUPPLY().sub(ctx.poolTokenBurnedSupply);
-    }
-
     function isAlmostEqualAmount(
         uint96 amount,
         uint96 targetAmount

packages/loopring_v3/contracts/amm/libamm/AmmWithdrawal.sol

@@ -48,8 +48,8 @@ library AmmWithdrawal
         // Burn any additional pool tokens stuck in forced exits
         AmmData.PoolExit storage exit = S.forcedExit[msg.sender];
         if (exit.burnAmount > 0) {
-            delete S.forcedExit[msg.sender];
             poolAmount = poolAmount.add(exit.burnAmount);
+            delete S.forcedExit[msg.sender];
         }
 
         require(poolAmount > 0, "ZERO_POOL_AMOUNT");
@@ -66,7 +66,7 @@ library AmmWithdrawal
             AmmUtil.transferOut(token, amount, msg.sender);
         }
 
-        S.poolTokenBurnedSupply = S.poolTokenBurnedSupply.add(poolAmount);
+        S._totalSupply = S._totalSupply.sub(poolAmount);
     }
 
     function _checkWithdrawalConditionInShutdown(
@@ -90,7 +90,7 @@ library AmmWithdrawal
             // Check that nothing is withdrawable anymore.
             require(
                 exchange.getAmountWithdrawable(address(this), token) == 0,
-                "MORE_TO_WITHDRAWAL"
+                "MORE_TO_WITHDRAW"
             );
         }
     }