Skip to content

Commit

Permalink
fix: first liquidity fix and more tests
Browse files Browse the repository at this point in the history
  • Loading branch information
jatZama committed Mar 5, 2024
1 parent ed4d059 commit 5be9fa6
Show file tree
Hide file tree
Showing 2 changed files with 218 additions and 4 deletions.
11 changes: 8 additions & 3 deletions contracts/EncryptedDEXPair.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
211 changes: 210 additions & 1 deletion test/EncryptedDEX/EncryptedDEX.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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");
});
});

0 comments on commit 5be9fa6

Please sign in to comment.