fixing tests in real fhevm mode (non mocked mode is also passing now)

jatZama · Feb 16, 2024 · 7db4765 · 7db4765
1 parent d483662
commit 7db4765
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 28 deletions.
diff --git a/contracts/EncryptedDEXPair.sol b/contracts/EncryptedDEXPair.sol
@@ -234,12 +234,12 @@ contract EncryptedDEXPair is EncryptedERC20 {
             uint32 burnedTotal
         )
     {
-        reserve0PendingAddDec = TFHE.decrypt(reserve0PendingAdd);
-        reserve1PendingAddDec = TFHE.decrypt(reserve1PendingAdd);
-        mintedTotal = TFHE.decrypt(pendingTotalMints[currentTradingEpoch]);
-        amount0In = TFHE.decrypt(pendingTotalToken0In[currentTradingEpoch]);
-        amount1In = TFHE.decrypt(pendingTotalToken1In[currentTradingEpoch]);
-        burnedTotal = TFHE.decrypt(pendingTotalBurns[currentTradingEpoch]);
+        if(TFHE.isInitialized(reserve0PendingAdd)) reserve0PendingAddDec = TFHE.decrypt(reserve0PendingAdd);
+        if(TFHE.isInitialized(reserve1PendingAdd)) reserve1PendingAddDec = TFHE.decrypt(reserve1PendingAdd);
+        if(TFHE.isInitialized(pendingTotalMints[currentTradingEpoch])) mintedTotal = TFHE.decrypt(pendingTotalMints[currentTradingEpoch]);
+        if(TFHE.isInitialized(pendingTotalToken0In[currentTradingEpoch])) amount0In = TFHE.decrypt(pendingTotalToken0In[currentTradingEpoch]);
+        if(TFHE.isInitialized(pendingTotalToken1In[currentTradingEpoch])) amount1In = TFHE.decrypt(pendingTotalToken1In[currentTradingEpoch]);
+        if(TFHE.isInitialized(pendingTotalBurns[currentTradingEpoch])) burnedTotal = TFHE.decrypt(pendingTotalBurns[currentTradingEpoch]);
     }
 
     function batchSettlement() external {
@@ -258,51 +258,48 @@ contract EncryptedDEXPair is EncryptedERC20 {
         ) = requestAllDecryptions();
 
         // update reserves after new liquidity deposits
-
         reserve0 += reserve0PendingAddDec;
         reserve1 += reserve1PendingAddDec;
         reserve0PendingAdd = ZERO;
         reserve1PendingAdd = ZERO;
 
         // Liquidity Mints
-        if (mintedTotal > 0) {
-            _mint(mintedTotal);
-            decryptedTotalMints[currentTradingEpoch] = mintedTotal;
-        }
         require(
             currentTradingEpoch != 0 || mintedTotal >= 100,
             "Initial minted liquidity amount should be greater than 100"
         ); // this is to lock forever at least 100 liquidity tokens inside the pool, so totalSupply of liquidity
         // would remain above 100 to avoid security issues,  for instance if a single market maker wants to burn the whole liquidity in a single transaction, making the pool unusable
+        if (mintedTotal > 0) {
+            _mint(mintedTotal);
+            decryptedTotalMints[currentTradingEpoch] = mintedTotal;
+        }
 
         // Token Swaps
         decryptedTotalToken0In[currentTradingEpoch] = amount0In;
         decryptedTotalToken1In[currentTradingEpoch] = amount1In;
+        uint32 amount0InMinusFee = uint32((99 * uint64(amount0In)) / 100); // 1% fee for liquidity providers
+        uint32 amount1InMinusFee = uint32((99 * uint64(amount1In)) / 100); // 1% fee for liquidity providers
         bool priceToken1Increasing = (uint64(amount0In) * uint64(reserve1) > uint64(amount1In) * uint64(reserve0));
         uint32 amount0Out;
         uint32 amount1Out;
         if (priceToken1Increasing) {
             // in this case, first sell all amount1In at current fixed token1 price to get amount0Out, then swap remaining (amount0In-amount0Out) to get amount1out_remaining according to AMM formula
-            amount0Out = uint32((uint64(amount1In) * uint64(reserve0)) / uint64(reserve1));
+            amount0Out = uint32((uint64(amount1InMinusFee) * uint64(reserve0)) / uint64(reserve1));
             amount1Out =
-                amount1In +
+                amount1InMinusFee +
                 reserve1 -
                 uint32(
-                    (uint64(reserve1) * uint64(reserve0)) / (uint64(reserve0) + uint64(amount0In) - uint64(amount0Out))
+                    (uint64(reserve1) * uint64(reserve0)) / (uint64(reserve0) + uint64(amount0InMinusFee) - uint64(amount0Out))
                 );
-            amount0Out = uint32((99 * uint64(amount0Out)) / 100); // 1% fee for liquidity providers
-            amount1Out = uint32((99 * uint64(amount1Out)) / 100); // 1% fee for liquidity providers
         } else {
             // here we do the opposite, first sell token0 at current token0 price then swap remaining token1 according to AMM formula
-            amount1Out = uint32((uint64(amount0In) * uint64(reserve1)) / uint64(reserve0));
+            amount1Out = uint32((uint64(amount0InMinusFee) * uint64(reserve1)) / uint64(reserve0));
             amount0Out =
-                amount0In +
+                amount0InMinusFee +
                 reserve0 -
                 uint32(
-                    (uint64(reserve0) * uint64(reserve1)) / (uint64(reserve1) + uint64(amount1In) - uint64(amount1Out))
+                    (uint64(reserve0) * uint64(reserve1)) / (uint64(reserve1) + uint64(amount1InMinusFee) - uint64(amount1Out))
                 );
-            amount0Out = uint32((99 * uint64(amount0Out)) / 100); // 1% fee for liquidity providers
-            amount1Out = uint32((99 * uint64(amount1Out)) / 100); // 1% fee for liquidity providers
         }
         totalToken0ClaimableSwap[currentTradingEpoch] = amount0Out;
         totalToken1ClaimableSwap[currentTradingEpoch] = amount1Out;

diff --git a/package.json b/package.json
@@ -28,7 +28,7 @@
     "eslint": "^8.28.0",
     "eslint-config-prettier": "^8.5.0",
     "ethers": "^6.8.0",
-    "fhevm": "../fhevm",
+    "fhevm": "0.4.0-0",
     "fhevmjs": "^0.4.0-3",
     "fs-extra": "^10.1.0",
     "hardhat": "^2.19.4",

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/test/EncryptedDEX/EncryptedDEX.ts b/test/EncryptedDEX/EncryptedDEX.ts
@@ -203,7 +203,7 @@ describe("Private DEX", function () {
     await tx23.wait();
     console.log("Eve submitted a swap order at tradingEpoch ", await pair.currentTradingEpoch(), "\n");
 
-    const tx24 = await pair.batchSettlement();
+    const tx24 = await pair.batchSettlement({ gasLimit : 10_000_000 });
     await tx24.wait();
 
     console.log(
@@ -243,7 +243,7 @@ describe("Private DEX", function () {
     balance = await getPrivateBalanceERC20(pairAddress, "carol");
     console.log("Carol now owns a private balance of ", balance, " liquidity tokens \n");
 
-    const tx29 = await pair.batchSettlement();
+    const tx29 = await pair.batchSettlement({ gasLimit : 10_000_000 });
     await tx29.wait();
 
     console.log(

diff --git a/test/EncryptedDEX/EncryptedERC20.fixture.ts b/test/EncryptedDEX/EncryptedERC20.fixture.ts
@@ -25,6 +25,6 @@ export async function getPrivateBalanceERC20(erc20Address: string, userName: str
   const encryptedBalance = await erc20
     .connect(signers[userName])
     .balanceOf(signers[userName], token.publicKey, token.signature);
-  const balance = instances.alice.decrypt(erc20Address, encryptedBalance);
+  const balance = instances[userName].decrypt(erc20Address, encryptedBalance);
   return balance;
 }