Merge branch 'fix/spearbit-audit' into fix/returned-emitted-init-liqu…

…idity
primitivefinance · Mar 25, 2024 · dbdbe65 · dbdbe65
2 parents 8040789 + 687a0e3
commit dbdbe65
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 5 deletions.
diff --git a/src/DFMM.sol b/src/DFMM.sol
@@ -109,8 +109,8 @@ contract DFMM is IDFMM {
         for (uint256 i = 0; i < tokensLength; i++) {
             address token = params.tokens[i];
 
-            for (uint256 j = 0; j < tokensLength; j++) {
-                if (i != j && token == params.tokens[j]) {
+            for (uint256 j = i + 1; j < tokensLength; j++) {
+                if (token == params.tokens[j]) {
                     revert InvalidDuplicateTokens();
                 }
             }
@@ -161,8 +161,8 @@ contract DFMM is IDFMM {
             _pools[poolId].reserves[i] += deltas[i];
         }
 
-        _pools[poolId].totalLiquidity += deltaLiquidity;
         _manageTokens(msg.sender, poolId, true, deltaLiquidity);
+        _pools[poolId].totalLiquidity += deltaLiquidity;
 
         for (uint256 i = 0; i < length; i++) {
             _transferFrom(_pools[poolId].tokens[i], deltas[i]);

diff --git a/test/G3M/unit/Allocate.t.sol b/test/G3M/unit/Allocate.t.sol
@@ -3,6 +3,7 @@ pragma solidity ^0.8.13;
 
 import "./SetUp.sol";
 import "solmate/utils/FixedPointMathLib.sol";
+import { LPToken } from "src/LPToken.sol";
 
 contract G3MAllocateTest is G3MSetUp {
     using FixedPointMathLib for uint256;
@@ -107,8 +108,10 @@ contract G3MAllocateTest is G3MSetUp {
             solver.allocateGivenDeltaY(POOL_ID, maxDeltaY);
         (uint256[] memory reserves, uint256 liquidity) =
             getReservesAndLiquidity(POOL_ID);
+        console2.log("liquidity", liquidity);
 
         uint256 preLiquidityBalance = liquidityOf(address(this), POOL_ID);
+        console2.log(preLiquidityBalance);
 
         bytes memory data = abi.encode(maxDeltaX, maxDeltaY, deltaLiquidity);
         console2.log(maxDeltaX);
@@ -124,11 +127,30 @@ contract G3MAllocateTest is G3MSetUp {
         assertEq(adjustedReserves[1], reserves[1] + deltas[1]);
         assertEq(adjustedLiquidity, liquidity + deltaLiquidity);
 
-        /*
         assertEq(
             preLiquidityBalance + deltaLiquidity,
             liquidityOf(address(this), POOL_ID)
         );
-        */
+    }
+
+    function test_G3M_allocate_ReceiveAppropriateLpTokens() public init_100 {
+      (, uint256 initialL) = getReservesAndLiquidity(POOL_ID);
+      Pool memory pool = dfmm.pools(POOL_ID);
+      LPToken liquidityToken = LPToken(pool.liquidityToken);
+
+      uint256 startBalance = liquidityToken.balanceOf(address(this));
+
+      uint256 dyMax = 100 ether;
+      (uint256 dxMax, uint256 dL) = solver.allocateGivenDeltaY(POOL_ID, dyMax);
+      bytes memory data = abi.encode(dxMax, dyMax, dL);
+
+      dfmm.allocate(POOL_ID, data);
+
+      (, uint256 nextL) = getReservesAndLiquidity(POOL_ID);
+      uint256 endBalance = liquidityToken.balanceOf(address(this));
+
+      // Add 1_000 wei to account for liquidity that was burnt on init
+      assertEq(startBalance + 1_000, initialL);
+      assertEq(endBalance + 1_000, nextL);
     }
 }
diff --git a/test/G3M/unit/SetUp.sol b/test/G3M/unit/SetUp.sol
@@ -31,6 +31,10 @@ contract G3MSetUp is SetUp {
         defaultReserveX, defaultStrikePrice, defaultParams
     );
 
+    bytes default100InitialPoolData = computeInitialPoolData(
+      defaultReserveX * 100, defaultStrikePrice, defaultParams
+    );
+
     function setUp() public override {
         SetUp.setUp();
         g3m = new GeometricMean(address(dfmm));
@@ -58,4 +62,24 @@ contract G3MSetUp is SetUp {
 
         _;
     }
+
+    modifier init_100() {
+      address[] memory tokens = new address[](2);
+        tokens[0] = address(tokenX);
+        tokens[1] = address(tokenY);
+
+        InitParams memory defaultInitParams = InitParams({
+            name: "",
+            symbol: "",
+            strategy: address(g3m),
+            tokens: tokens,
+            data: default100InitialPoolData,
+            feeCollector: address(0),
+            controllerFee: 0
+        });
+
+        (POOL_ID,,) = dfmm.init(defaultInitParams);
+
+        _;
+    }
 }