Merge pull request #119 from primitivefinance/fix/post-review

Fix/post review

src/ConstantSum/ConstantSum.sol

@@ -5,7 +5,8 @@ import {
     FixedPointMathLib,
     computeTradingFunction,
     computeSwapDeltaLiquidity,
-    computeDeltaLiquidity
+    computeDeltaLiquidityRoundDown,
+    computeDeltaLiquidityRoundUp
 } from "./ConstantSumMath.sol";
 import {
     decodePriceUpdate,
@@ -67,8 +68,9 @@ contract ConstantSum is PairStrategy {
         ConstantSumParams memory params;
 
         (reserves, params) = abi.decode(data, (uint256[], ConstantSumParams));
-        totalLiquidity =
-            computeDeltaLiquidity(reserves[0], reserves[1], params.price);
+        totalLiquidity = computeDeltaLiquidityRoundDown(
+            reserves[0], reserves[1], params.price
+        );
 
         if (pool.reserves.length != 2 || reserves.length != 2) {
             revert InvalidReservesLength();
@@ -106,8 +108,9 @@ contract ConstantSum is PairStrategy {
         (uint256 deltaX, uint256 deltaY, uint256 minDeltaL) =
             abi.decode(data, (uint256, uint256, uint256));
 
-        deltaLiquidity =
-            computeDeltaLiquidity(deltaX, deltaY, internalParams[poolId].price);
+        deltaLiquidity = computeDeltaLiquidityRoundDown(
+            deltaX, deltaY, internalParams[poolId].price
+        );
         if (deltaLiquidity < minDeltaL) revert InvalidDeltaLiquidity();
 
         deltas = new uint256[](2);
@@ -145,9 +148,10 @@ contract ConstantSum is PairStrategy {
         (uint256 deltaX, uint256 deltaY, uint256 maxDeltaL) =
             abi.decode(data, (uint256, uint256, uint256));
 
-        deltaLiquidity =
-            computeDeltaLiquidity(deltaX, deltaY, internalParams[poolId].price);
-        if (maxDeltaL > deltaLiquidity) revert InvalidDeltaLiquidity();
+        deltaLiquidity = computeDeltaLiquidityRoundUp(
+            deltaX, deltaY, internalParams[poolId].price
+        );
+        if (deltaLiquidity > maxDeltaL) revert InvalidDeltaLiquidity();
 
         deltas = new uint256[](2);
         deltas[0] = deltaX;

src/ConstantSum/ConstantSumMath.sol

@@ -31,22 +31,30 @@ function computeInitialPoolData(
     return abi.encode(reserves, params);
 }
 
-function computeDeltaLiquidity(
+function computeDeltaLiquidityRoundUp(
     uint256 deltaX,
     uint256 deltaY,
     uint256 price
 ) pure returns (uint256) {
     return price.mulWadUp(deltaX) + deltaY;
 }
 
+function computeDeltaLiquidityRoundDown(
+    uint256 deltaX,
+    uint256 deltaY,
+    uint256 price
+) pure returns (uint256) {
+    return price.mulWadDown(deltaX) + deltaY;
+}
+
 function computeSwapDeltaLiquidity(
     uint256 delta,
     ConstantSumParams memory params,
     bool isSwapXForY
 ) pure returns (uint256) {
     if (isSwapXForY) {
-        return (params.swapFee).mulWadUp(delta);
+        return params.swapFee.mulWadUp(delta.mulWadUp(params.price));
     } else {
-        return (params.swapFee).mulDivUp(delta, params.price);
+        return params.swapFee.mulWadUp(delta);
     }
 }

src/ConstantSum/ConstantSumSolver.sol

@@ -41,11 +41,6 @@ contract ConstantSumSolver {
         return computeInitialPoolData(rx, ry, params);
     }
 
-    struct SimulateSwapState {
-        uint256 amountOut;
-        uint256 deltaLiquidity;
-    }
-
     function simulateSwap(
         uint256 poolId,
         bool swapXIn,
@@ -56,39 +51,37 @@ contract ConstantSumSolver {
             IStrategy(strategy).getPoolParams(poolId), (ConstantSumParams)
         );
 
-        SimulateSwapState memory state;
+        uint256 amountOut;
 
         if (swapXIn) {
-            computeSwapDeltaLiquidity(amountIn, poolParams, true);
-            state.amountOut = amountIn.mulWadDown(poolParams.price).mulWadDown(
+            amountOut = amountIn.mulWadDown(poolParams.price).mulWadDown(
                 ONE - poolParams.swapFee
             );
 
-            if (pool.reserves[1] < state.amountOut) {
+            if (pool.reserves[1] < amountOut) {
                 revert NotEnoughLiquidity();
             }
         } else {
-            computeSwapDeltaLiquidity(amountIn, poolParams, false);
-            state.amountOut = (ONE - poolParams.swapFee).mulWadDown(amountIn)
+            amountOut = (ONE - poolParams.swapFee).mulWadDown(amountIn)
                 .divWadDown(poolParams.price);
 
-            if (pool.reserves[0] < state.amountOut) {
+            if (pool.reserves[0] < amountOut) {
                 revert NotEnoughLiquidity();
             }
         }
 
         bytes memory swapData;
 
         if (swapXIn) {
-            swapData = abi.encode(0, 1, amountIn, state.amountOut);
+            swapData = abi.encode(0, 1, amountIn, amountOut);
         } else {
-            swapData = abi.encode(1, 0, amountIn, state.amountOut);
+            swapData = abi.encode(1, 0, amountIn, amountOut);
         }
 
         (bool valid,,,,,,) = IStrategy(strategy).validateSwap(
             address(this), poolId, pool, swapData
         );
-        return (valid, state.amountOut, swapData);
+        return (valid, amountOut, swapData);
     }
 
     function preparePriceUpdate(uint256 newPrice)

src/DFMM.sol

@@ -110,15 +110,17 @@ contract DFMM is IDFMM {
 
         for (uint256 i = 0; i < tokensLength; i++) {
             address token = params.tokens[i];
-            uint256 decimals = ERC20(params.tokens[i]).decimals();
+            uint256 decimals = ERC20(token).decimals();
 
             if (decimals > 18 || decimals < 6) {
                 revert InvalidTokenDecimals();
             }
 
-            for (uint256 j = i + 1; j < tokensLength; j++) {
-                if (token == params.tokens[j]) {
-                    revert InvalidDuplicateTokens();
+            unchecked {
+                for (uint256 j = i + 1; j < tokensLength; j++) {
+                    if (token == params.tokens[j]) {
+                        revert InvalidDuplicateTokens();
+                    }
                 }
             }
         }
@@ -281,9 +283,13 @@ contract DFMM is IDFMM {
     // Internals
 
     /**
-     * @dev Transfers `amounts` of `tokens` from the sender to the contract. Note
-     * that if any ETH is present in the contract, it will be wrapped to WETH and
-     * used if sufficient. Any excess of ETH will be sent back to the sender.
+     * @notice Transfers `amounts` of `tokens` from the sender to the contract.
+     * @dev Note that for pools with `WETH` as a token, the contract will accept
+     * full payments in native ether. If the contract receives more ether than
+     * the amount of `WETH` needed, the contract will refund the remaining ether.
+     * If the contract receives less ether than the amount of `WETH` needed, the
+     * contract will refund the native ether and request the full amount of `WETH`
+     * needed instead.
      * @param tokens An array of token addresses to transfer.
      * @param amounts An array of amounts to transfer expressed in WAD.
      */
@@ -301,23 +307,30 @@ contract DFMM is IDFMM {
                 downscaleUp(amount, computeScalingFactor(token));
             uint256 preBalance = ERC20(token).balanceOf(address(this));
 
-            if (token == weth && address(this).balance >= amount) {
+            // note: `msg.value` can be used safely in a loop because `weth` is a unique token,
+            // therefore we only enter this branch once.
+            if (token == weth && msg.value >= amount) {
                 WETH(payable(weth)).deposit{ value: amount }();
             } else {
                 SafeTransferLib.safeTransferFrom(
                     ERC20(token), msg.sender, address(this), downscaledAmount
                 );
             }
 
+            // If not enough native ether was sent as payment
+            // or too much ether was sent,
+            // refund all the remaining ether back to the sender.
+            if (token == weth && msg.value != 0) {
+                SafeTransferLib.safeTransferETH(
+                    msg.sender, address(this).balance
+                );
+            }
+
             uint256 postBalance = ERC20(token).balanceOf(address(this));
             if (postBalance < preBalance + downscaledAmount) {
                 revert InvalidTransfer();
             }
         }
-
-        if (address(this).balance > 0) {
-            SafeTransferLib.safeTransferETH(msg.sender, address(this).balance);
-        }
     }
 
     /**
@@ -334,6 +347,7 @@ contract DFMM is IDFMM {
         } else {
             uint256 downscaledAmount =
                 downscaleDown(amount, computeScalingFactor(token));
+            if (downscaledAmount == 0) return;
             uint256 preBalance = ERC20(token).balanceOf(address(this));
             SafeTransferLib.safeTransfer(ERC20(token), to, downscaledAmount);
             uint256 postBalance = ERC20(token).balanceOf(address(this));

src/GeometricMean/G3MMath.sol

@@ -18,7 +18,7 @@ function computeTradingFunction(
     uint256 a = uint256(int256(rX.divWadUp(L)).powWad(int256(params.wX)));
     uint256 b = uint256(int256(rY.divWadUp(L)).powWad(int256(params.wY)));
 
-    return int256(a.mulWadUp(b)) - int256(1 ether);
+    return int256(a.mulWadUp(b)) - int256(ONE);
 }
 
 function computeDeltaGivenDeltaLRoundUp(

src/LogNormal/LogNormal.sol

@@ -88,7 +88,7 @@ contract LogNormal is PairStrategy {
         (reserves, totalLiquidity, params) =
             abi.decode(data, (uint256[], uint256, LogNormalParams));
 
-        if (params.mean < MIN_WIDTH || params.mean > MAX_MEAN) {
+        if (params.mean < MIN_MEAN || params.mean > MAX_MEAN) {
             revert InvalidMean();
         }
 

src/NTokenGeometricMean/NTokenGeometricMeanMath.sol

@@ -17,7 +17,7 @@ function computeTradingFunction(
     uint256 accumulator = ONE;
     for (uint256 i = 0; i < reserves.length; i++) {
         uint256 a = uint256(
-            int256(reserves[i].divWadDown(L)).powWad(int256(params.weights[i]))
+            int256(reserves[i].divWadUp(L)).powWad(int256(params.weights[i]))
         );
         accumulator = accumulator.mulWadUp(a);
     }
@@ -30,15 +30,15 @@ function computeDeltaGivenDeltaLRoundUp(
     uint256 deltaLiquidity,
     uint256 totalLiquidity
 ) pure returns (uint256) {
-    return reserve.mulWadUp(deltaLiquidity.divWadUp(totalLiquidity));
+    return reserve.mulDivUp(deltaLiquidity, totalLiquidity);
 }
 
 function computeDeltaGivenDeltaLRoundDown(
     uint256 reserve,
     uint256 deltaLiquidity,
     uint256 totalLiquidity
 ) pure returns (uint256) {
-    return reserve.mulWadDown(deltaLiquidity.divWadDown(totalLiquidity));
+    return reserve.mulDivDown(deltaLiquidity, totalLiquidity);
 }
 
 function computeL(
@@ -137,7 +137,7 @@ function computeSwapDeltaLiquidity(
     uint256 weight,
     uint256 swapFee
 ) pure returns (uint256) {
-    return weight.mulWadDown(swapFee).mulWadDown(totalLiquidity).mulWadDown(
-        amountIn.divWadDown(reserve)
+    return weight.mulWadUp(swapFee).mulWadUp(totalLiquidity).mulWadUp(
+        amountIn.divWadUp(reserve)
     );
 }

test/ConstantSum/unit/Allocate.t.sol

@@ -3,7 +3,7 @@ pragma solidity ^0.8.13;
 
 import { ConstantSumSetUp } from "./SetUp.sol";
 import {
-    computeDeltaLiquidity,
+    computeDeltaLiquidityRoundDown,
     ConstantSumParams
 } from "src/ConstantSum/ConstantSumMath.sol";
 
@@ -15,7 +15,8 @@ contract ConstantSumAllocateTest is ConstantSumSetUp {
         ConstantSumParams memory params =
             abi.decode(constantSum.getPoolParams(POOL_ID), (ConstantSumParams));
 
-        uint256 deltaL = computeDeltaLiquidity(deltaX, deltaY, params.price);
+        uint256 deltaL =
+            computeDeltaLiquidityRoundDown(deltaX, deltaY, params.price);
         dfmm.allocate(POOL_ID, abi.encode(deltaX, deltaY, deltaL));
     }
 }

test/DFMM/unit/Init.t.sol

@@ -17,6 +17,9 @@ contract DFMMInit is DFMMSetUp, Script {
     bytes defaultData =
         abi.encode(valid, initialInvariant, defaultReserves, initialLiquidity);
 
+    /// @notice for handling ether refunds
+    receive() external payable { }
+
     function test_DFMM_init_StoresStrategyInitialReservesAndLiquidity()
         public
     {