diff --git a/contracts/EncryptedDEXPair.sol b/contracts/EncryptedDEXPair.sol index 19eb83e..2fc25c8 100644 --- a/contracts/EncryptedDEXPair.sol +++ b/contracts/EncryptedDEXPair.sol @@ -122,20 +122,25 @@ contract EncryptedDEXPair is EncryptedERC20 { if (firstBlockPerEpoch[currentEpoch] == 0) { firstBlockPerEpoch[currentEpoch] = block.number; } - reserve0PendingAdd = reserve0PendingAdd + amount0; reserve1PendingAdd = reserve1PendingAdd + amount1; - euint64 liquidity; if (totalSupply() == 0) { // this condition is equivalent to currentEpoch==0 (see batchSettlement logic) liquidity = TFHE.shr(amount0, 1) + TFHE.shr(amount1, 1); + ebool isBelowMinimum = TFHE.lt(liquidity, MINIMIMUM_LIQUIDITY); + liquidity = TFHE.cmux(isBelowMinimum, ZERO, liquidity); + euint64 amount0Back = TFHE.cmux(isBelowMinimum, amount0, ZERO); + euint64 amount1Back = TFHE.cmux(isBelowMinimum, amount1, ZERO); + token0.transfer(msg.sender, amount0Back); // refund first liquidity if it is below the minimal amount + token1.transfer(msg.sender, amount1Back); // refund first liquidity if it is below the minimal amount + reserve0PendingAdd = reserve0PendingAdd - amount0Back; + reserve1PendingAdd = reserve1PendingAdd - amount1Back; } else { euint64 liquidity0 = TFHE.div(TFHE.mul(TFHE.shr(amount0, 32), _totalSupply >> 32), reserve0 >> 32); euint64 liquidity1 = TFHE.div(TFHE.mul(TFHE.shr(amount1, 32), _totalSupply >> 32), reserve1 >> 32); liquidity = TFHE.shl(TFHE.min(liquidity0, liquidity1), 32); } - pendingMints[currentEpoch][to] = pendingMints[currentEpoch][to] + liquidity; pendingTotalMints[currentEpoch] = pendingTotalMints[currentEpoch] + liquidity; } diff --git a/test/EncryptedDEX/EncryptedDEX.ts b/test/EncryptedDEX/EncryptedDEX.ts index c3288b4..1724593 100644 --- a/test/EncryptedDEX/EncryptedDEX.ts +++ b/test/EncryptedDEX/EncryptedDEX.ts @@ -12,7 +12,7 @@ describe("Private DEX", function () { this.signers = await getSigners(); }); - it("Encrypted DEX Pool", async function () { + it("Encrypted DEX Pool - multiple epochs", async function () { const COIN = 2n ** 32n; let token0 = await deployEncryptedERC20Fixture(); let token0Address = await token0.getAddress(); @@ -272,4 +272,213 @@ describe("Private DEX", function () { console.log("Carol's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "carol")) / COIN); console.log("Carol's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "carol")) / COIN); }); + + it("Encrypted DEX Pool - single epoch (addLiquidity+swap)", async function () { + const COIN = 2n ** 32n; + let token0 = await deployEncryptedERC20Fixture(); + let token0Address = await token0.getAddress(); + let token1 = await deployEncryptedERC20Fixture(); + let token1Address = await token1.getAddress(); + + BigInt(token0Address) > BigInt(token1Address) ? ([token0, token1] = [token1, token0]) : null; // sort tokens according to addresses + token0Address = await token0.getAddress(); + token1Address = await token1.getAddress(); + + const tx0 = await token0.mint(2_000_000_000n * COIN); + await tx0.wait(); + const instances0 = await createInstances(token0Address, ethers, this.signers); + const tx1 = await token1.mint(2_000_000_000n * COIN); + await tx1.wait(); + const instances1 = await createInstances(token1Address, ethers, this.signers); + let balance = await getPrivateBalanceERC20(token0Address, "alice"); + expect(balance).to.equal(2_000_000_000n * COIN); + const totalSupply = await token0.totalSupply(); + expect(totalSupply).to.equal(2_000_000_000n * COIN); + + const tx2 = await token0["transfer(address,bytes)"]( + this.signers.bob.address, + instances0.alice.encrypt64(100_000_000n * COIN), + ); + await tx2.wait(); + + balance = await getPrivateBalanceERC20(token0Address, "alice"); + expect(balance).to.equal(1_900_000_000n * COIN); + balance = await getPrivateBalanceERC20(token0Address, "bob"); + expect(balance).to.equal(100_000_000n * COIN); + + const tx3 = await token0["transfer(address,bytes)"]( + this.signers.carol.address, + instances0.alice.encrypt64(200_000_000n * COIN), + ); + await tx3.wait(); + + const tx4 = await token1["transfer(address,bytes)"]( + this.signers.bob.address, + instances1.alice.encrypt64(200_000_000n * COIN), + ); + await tx4.wait(); + + const tx5 = await token1["transfer(address,bytes)"]( + this.signers.carol.address, + instances1.alice.encrypt64(400_000_000n * COIN), + ); + await tx5.wait(); + + balance = await getPrivateBalanceERC20(token1Address, "carol"); + expect(balance).to.equal(400_000_000n * COIN); + // BOB and CAROL are market makers, BOB starts with 100M token0 and 200M token1, CAROL starts with 200M token0 and 400M token1 + + console.log("Initial balances of market makers (Bob and Carol) : "); + console.log("Bob's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "bob")) / COIN); + console.log("Bob's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "bob")) / COIN); + console.log("Carol's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "carol")) / COIN); + console.log("Carol's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "carol")) / COIN); + + const tx6 = await token0["transfer(address,bytes)"]( + this.signers.dave.address, + instances0.alice.encrypt64(1_000_000n * COIN), + ); + await tx6.wait(); + const tx6bis = await token1["transfer(address,bytes)"]( + this.signers.dave.address, + instances1.alice.encrypt64(0n * COIN), // to obfuscate the direction of the swap later, needed to initialize dave's token1 balance + ); + await tx6bis.wait(); + + const tx7 = await token1["transfer(address,bytes)"]( + this.signers.eve.address, + instances1.alice.encrypt64(1_000_000n * COIN), + ); + await tx7.wait(); + const tx7bis = await token0["transfer(address,bytes)"]( + this.signers.eve.address, + instances0.alice.encrypt64(0n * COIN), // to obfuscate the direction of the swap later, needed to initialize eve's token0 balance + ); + await tx7bis.wait(); + + console.log("Initial balances of traders (Dave and Eve) : "); + console.log("Dave's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "dave")) / COIN); + console.log("Dave's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "dave")) / COIN); + console.log("Eve's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "eve")) / COIN); + console.log("Eve's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "eve")) / COIN, "\n"); + + const dexFactory = await deployEncryptedDEXFactoryFixture(); + const tx8 = await dexFactory.createPair(token0Address, token1Address); + await tx8.wait(); + + const pairAddress = await dexFactory.getPair(token0Address, token1Address); + const pair = await ethers.getContractAt("EncryptedDEXPair", pairAddress); + console.log("DEX contract was deployed for pair token0/token1 \n"); + + const instancesPair = await createInstances(pairAddress, ethers, this.signers); + + const tx9 = await token0 + .connect(this.signers.bob) + ["approve(address,bytes)"](pairAddress, instances0.bob.encrypt64(100_000_000n * COIN)); + await tx9.wait(); + const tx10 = await token1 + .connect(this.signers.bob) + ["approve(address,bytes)"](pairAddress, instances1.bob.encrypt64(200_000_000n * COIN)); + await tx10.wait(); + + const tx11 = await pair + .connect(this.signers.bob) + .addLiquidity( + instancesPair.bob.encrypt64(100_000_000n * COIN), + instancesPair.bob.encrypt64(200_000_000n * COIN), + this.signers.bob.address, + 0n, + ); + await tx11.wait(); + console.log("Bob submitted an addLiquidity order at tradingEpoch ", await pair.currentTradingEpoch()); + + const tx12 = await token0 + .connect(this.signers.carol) + ["approve(address,bytes)"](pairAddress, instances0.carol.encrypt64(200_000_000n * COIN)); + await tx12.wait(); + const tx13 = await token1 + .connect(this.signers.carol) + ["approve(address,bytes)"](pairAddress, instances0.carol.encrypt64(400_000_000n * COIN)); + await tx13.wait(); + const tx14 = await pair + .connect(this.signers.carol) + .addLiquidity( + instancesPair.carol.encrypt64(200_000_000n * COIN), + instancesPair.carol.encrypt64(400_000_000n * COIN), + this.signers.carol.address, + 0n, + ); + await tx14.wait(); + console.log("Carol submitted an addLiquidity order at tradingEpoch ", await pair.currentTradingEpoch(), "\n"); + + const tx18 = await token0 + .connect(this.signers.dave) + ["approve(address,bytes)"](pairAddress, instances0.bob.encrypt64(100_000_000n * COIN)); + await tx18.wait(); + const tx19 = await token1 + .connect(this.signers.dave) + ["approve(address,bytes)"](pairAddress, instances1.bob.encrypt64(100_000_000n * COIN)); + await tx19.wait(); + const tx20 = await pair + .connect(this.signers.dave) + .swapTokens( + instancesPair.dave.encrypt64(1_000_000n * COIN), + instancesPair.dave.encrypt64(0n * COIN), + this.signers.dave.address, + 1n, + ); + await tx20.wait(); + console.log("Dave submitted a swap order at tradingEpoch ", await pair.currentTradingEpoch(), "\n"); + + const tx21 = await token0 + .connect(this.signers.eve) + ["approve(address,bytes)"](pairAddress, instances0.bob.encrypt64(100_000_000n * COIN)); + await tx21.wait(); + const tx22 = await token1 + .connect(this.signers.eve) + ["approve(address,bytes)"](pairAddress, instances1.bob.encrypt64(100_000_000n * COIN)); + await tx22.wait(); + const tx23 = await pair + .connect(this.signers.eve) + .swapTokens( + instancesPair.eve.encrypt64(0n * COIN), + instancesPair.eve.encrypt64(1_000_000n * COIN), + this.signers.eve.address, + 1n, + ); + await tx23.wait(); + console.log("Eve submitted a swap order at tradingEpoch ", await pair.currentTradingEpoch(), "\n"); + + const tx24 = await pair.batchSettlement({ gasLimit: 10_000_000 }); + await tx24.wait(); + + console.log( + "Batch Settlement was confirmed with threshold decryptions for tradingEpoch ", + (await pair.currentTradingEpoch()) - 1n, + ); + console.log("New reserves for tradingEpoch ", await pair.currentTradingEpoch(), " are now publicly revealed : "); + let [reserve0, reserve1] = await pair.getReserves(); + console.log("Reserve token0 ", reserve0 / COIN); + console.log("Reserve token1 ", reserve1 / COIN, "\n"); + + const tx25 = await pair.claimSwap(0n, this.signers.dave.address); + await tx25.wait(); + const tx26 = await pair.claimSwap(0n, this.signers.eve.address); + await tx26.wait(); + + const tx16 = await pair.claimMint(0n, this.signers.bob.address); + await tx16.wait(); + const tx17 = await pair.claimMint(0n, this.signers.carol.address); + await tx17.wait(); + balance = await getPrivateBalanceERC20(pairAddress, "bob"); + console.log("Bob now owns a private balance of ", balance / COIN, " liquidity tokens"); + balance = await getPrivateBalanceERC20(pairAddress, "carol"); + console.log("Carol now owns a private balance of ", balance / COIN, " liquidity tokens \n"); + + console.log("New balances of traders (Dave and Eve) : "); + console.log("Dave's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "dave")) / COIN); + console.log("Dave's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "dave")) / COIN); + console.log("Eve's token0 balance : ", (await getPrivateBalanceERC20(token0Address, "eve")) / COIN); + console.log("Eve's token1 balance : ", (await getPrivateBalanceERC20(token1Address, "eve")) / COIN, "\n"); + }); });