refactor: support multiple tokens

zama-ai · Dec 23, 2024 · 73be915 · 73be915
1 parent 1476a51
commit 73be915
Show file tree

Hide file tree

Showing 8 changed files with 48 additions and 68 deletions.
diff --git a/contracts/finance/ConfidentialVestingWallet.sol b/contracts/finance/ConfidentialVestingWallet.sol
@@ -14,15 +14,13 @@ import { IConfidentialERC20 } from "../token/ERC20/IConfidentialERC20.sol";
  *         To use with the native asset, it is necessary to wrap the native asset to a ConfidentialERC20-like token.
  */
 abstract contract ConfidentialVestingWallet {
-    /// @notice Emitted when tokens are released to the beneficiary address.
-    event ConfidentialERC20Released();
+    /// @notice      Emitted when tokens are released to the beneficiary address.
+    /// @param token Address of the token being released.
+    event ConfidentialERC20Released(address indexed token);
 
     /// @notice Beneficiary address.
     address public immutable BENEFICIARY;
 
-    /// @notice Confidential ERC20.
-    IConfidentialERC20 public immutable CONFIDENTIAL_ERC20;
-
     /// @notice Duration (in seconds).
     uint128 public immutable DURATION;
 
@@ -38,24 +36,21 @@ abstract contract ConfidentialVestingWallet {
     euint64 internal immutable _EUINT64_ZERO;
 
     /// @notice Total encrypted amount released (to the beneficiary).
-    euint64 internal _amountReleased;
+    mapping(address token => euint64 amountReleased) internal _amountReleased;
 
     /**
      * @param beneficiary_      Beneficiary address.
-     * @param token_            Confidential token address.
      * @param startTimestamp_   Start timestamp.
      * @param duration_         Duration (in seconds).
      */
-    constructor(address beneficiary_, address token_, uint128 startTimestamp_, uint128 duration_) {
+    constructor(address beneficiary_, uint128 startTimestamp_, uint128 duration_) {
         START_TIMESTAMP = startTimestamp_;
-        CONFIDENTIAL_ERC20 = IConfidentialERC20(token_);
         DURATION = duration_;
         END_TIMESTAMP = startTimestamp_ + duration_;
         BENEFICIARY = beneficiary_;
 
         /// @dev Store this constant variable in the storage.
         _EUINT64_ZERO = TFHE.asEuint64(0);
-        _amountReleased = _EUINT64_ZERO;
 
         TFHE.allow(_EUINT64_ZERO, beneficiary_);
         TFHE.allowThis(_EUINT64_ZERO);
@@ -65,43 +60,44 @@ abstract contract ConfidentialVestingWallet {
      * @notice  Release the tokens that have already vested.
      * @dev     Anyone can call this function but the beneficiary receives the tokens.
      */
-    function release() public virtual {
-        euint64 amount = _releasable();
-        euint64 amountReleased = TFHE.add(_amountReleased, amount);
-        _amountReleased = amountReleased;
+    function release(address token) public virtual {
+        euint64 amount = _releasable(token);
+        euint64 amountReleased = TFHE.add(_amountReleased[token], amount);
+        _amountReleased[token] = amountReleased;
 
         TFHE.allow(amountReleased, BENEFICIARY);
         TFHE.allowThis(amountReleased);
-        TFHE.allowTransient(amount, address(CONFIDENTIAL_ERC20));
-        CONFIDENTIAL_ERC20.transfer(BENEFICIARY, amount);
+        TFHE.allowTransient(amount, token);
+        IConfidentialERC20(token).transfer(BENEFICIARY, amount);
 
-        emit ConfidentialERC20Released();
+        emit ConfidentialERC20Released(token);
     }
 
     /**
      * @notice                  Return the encrypted amount of total tokens released.
      * @dev                     It is only reencryptable by the owner.
      * @return amountReleased   Total amount of tokens released.
      */
-    function released() public view virtual returns (euint64 amountReleased) {
-        return _amountReleased;
+    function released(address token) public view virtual returns (euint64 amountReleased) {
+        return _amountReleased[token];
     }
 
     /**
      * @notice                  Calculate the amount of tokens that can be released.
      * @return releasableAmount Releasable amount.
      */
-    function _releasable() internal virtual returns (euint64 releasableAmount) {
-        return TFHE.sub(_vestedAmount(uint128(block.timestamp)), released());
+    function _releasable(address token) internal virtual returns (euint64 releasableAmount) {
+        return TFHE.sub(_vestedAmount(token, uint128(block.timestamp)), released(token));
     }
 
     /**
      * @notice                  Calculate the amount of tokens that has already vested.
      * @param timestamp         Current timestamp.
      * @return vestedAmount     Vested amount.
      */
-    function _vestedAmount(uint128 timestamp) internal virtual returns (euint64 vestedAmount) {
-        return _vestingSchedule(TFHE.add(CONFIDENTIAL_ERC20.balanceOf(address(this)), released()), timestamp);
+    function _vestedAmount(address token, uint128 timestamp) internal virtual returns (euint64 vestedAmount) {
+        return
+            _vestingSchedule(TFHE.add(IConfidentialERC20(token).balanceOf(address(this)), released(token)), timestamp);
     }
 
     /**

diff --git a/contracts/finance/ConfidentialVestingWalletCliff.sol b/contracts/finance/ConfidentialVestingWalletCliff.sol
@@ -22,18 +22,16 @@ abstract contract ConfidentialVestingWalletCliff is ConfidentialVestingWallet {
 
     /**
      * @param beneficiary_      Beneficiary address.
-     * @param token_            Confidential token address.
      * @param startTimestamp_   Start timestamp.
      * @param duration_         Duration (in seconds).
      * @param cliffSeconds_     Cliff (in seconds).
      */
     constructor(
         address beneficiary_,
-        address token_,
         uint128 startTimestamp_,
         uint128 duration_,
         uint128 cliffSeconds_
-    ) ConfidentialVestingWallet(beneficiary_, token_, startTimestamp_, duration_) {
+    ) ConfidentialVestingWallet(beneficiary_, startTimestamp_, duration_) {
         if (cliffSeconds_ > duration_) {
             revert InvalidCliffDuration(cliffSeconds_, duration_);
         }

diff --git a/contracts/test/finance/TestConfidentialVestingWallet.sol b/contracts/test/finance/TestConfidentialVestingWallet.sol
@@ -7,10 +7,9 @@ import { SepoliaZamaFHEVMConfig } from "fhevm/config/ZamaFHEVMConfig.sol";
 contract TestConfidentialVestingWallet is SepoliaZamaFHEVMConfig, ConfidentialVestingWallet {
     constructor(
         address beneficiary_,
-        address token_,
         uint64 startTimestamp_,
         uint64 duration_
-    ) ConfidentialVestingWallet(beneficiary_, token_, startTimestamp_, duration_) {
+    ) ConfidentialVestingWallet(beneficiary_, startTimestamp_, duration_) {
         //
     }
 }
diff --git a/contracts/test/finance/TestConfidentialVestingWalletCliff.sol b/contracts/test/finance/TestConfidentialVestingWalletCliff.sol
@@ -7,11 +7,10 @@ import { SepoliaZamaFHEVMConfig } from "fhevm/config/ZamaFHEVMConfig.sol";
 contract TestConfidentialVestingWalletCliff is SepoliaZamaFHEVMConfig, ConfidentialVestingWalletCliff {
     constructor(
         address beneficiary_,
-        address token_,
         uint64 startTimestamp_,
         uint64 duration_,
         uint64 cliff_
-    ) ConfidentialVestingWalletCliff(beneficiary_, token_, startTimestamp_, duration_, cliff_) {
+    ) ConfidentialVestingWalletCliff(beneficiary_, startTimestamp_, duration_, cliff_) {
         //
     }
 }
diff --git a/test/finance/ConfidentialVestingWallet.fixture.ts b/test/finance/ConfidentialVestingWallet.fixture.ts
@@ -8,23 +8,23 @@ import { reencryptEuint64 } from "../reencrypt";
 export async function deployConfidentialVestingWalletFixture(
   account: Signer,
   beneficiaryAddress: string,
-  token: string,
   startTimestamp: bigint,
   duration: bigint,
 ): Promise<TestConfidentialVestingWallet> {
   const contractFactory = await ethers.getContractFactory("TestConfidentialVestingWallet");
-  const contract = await contractFactory.connect(account).deploy(beneficiaryAddress, token, startTimestamp, duration);
+  const contract = await contractFactory.connect(account).deploy(beneficiaryAddress, startTimestamp, duration);
   await contract.waitForDeployment();
   return contract;
 }
 
 export async function reencryptReleased(
   account: Signer,
   instance: FhevmInstance,
+  tokenAddress: string,
   vestingWallet: ConfidentialVestingWallet,
   vestingWalletAddress: string,
 ): Promise<bigint> {
-  const releasedHandled = await vestingWallet.released();
+  const releasedHandled = await vestingWallet.released(tokenAddress);
   const releasedAmount = await reencryptEuint64(account, instance, releasedHandled, vestingWalletAddress);
   return releasedAmount;
 }
diff --git a/test/finance/ConfidentialVestingWallet.test.ts b/test/finance/ConfidentialVestingWallet.test.ts
@@ -35,7 +35,6 @@ describe("ConfidentialVestingWallet", function () {
     const contractConfidentialVestingWallet = await deployConfidentialVestingWalletFixture(
       this.signers.alice,
       this.beneficiaryAddress,
-      this.confidentialERC20Address,
       this.startTimestamp,
       this.duration,
     );
@@ -46,18 +45,9 @@ describe("ConfidentialVestingWallet", function () {
 
   it("post-deployment state", async function () {
     expect(await this.confidentialVestingWallet.BENEFICIARY()).to.equal(this.beneficiaryAddress);
-    expect(await this.confidentialVestingWallet.CONFIDENTIAL_ERC20()).to.equal(this.confidentialERC20);
     expect(await this.confidentialVestingWallet.DURATION()).to.equal(this.duration);
     expect(await this.confidentialVestingWallet.END_TIMESTAMP()).to.be.eq(this.startTimestamp + this.duration);
     expect(await this.confidentialVestingWallet.START_TIMESTAMP()).to.be.eq(this.startTimestamp);
-    expect(
-      await reencryptReleased(
-        this.beneficiary,
-        this.instance,
-        this.confidentialVestingWallet,
-        this.confidentialVestingWalletAddress,
-      ),
-    ).to.be.eq(0n);
   });
 
   it("can release", async function () {
@@ -82,14 +72,17 @@ describe("ConfidentialVestingWallet", function () {
     let nextTimestamp = this.startTimestamp;
     await ethers.provider.send("evm_setNextBlockTimestamp", [nextTimestamp.toString()]);
 
-    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release();
-    await expect(tx).to.emit(this.confidentialVestingWallet, "ConfidentialERC20Released");
+    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release(this.confidentialERC20Address);
+    await expect(tx)
+      .to.emit(this.confidentialVestingWallet, "ConfidentialERC20Released")
+      .withArgs(this.confidentialERC20Address);
 
     // It should be equal to 0 because the vesting has not started.
     expect(
       await reencryptReleased(
         this.beneficiary,
         this.instance,
+        this.confidentialERC20Address,
         this.confidentialVestingWallet,
         this.confidentialVestingWalletAddress,
       ),
@@ -98,14 +91,15 @@ describe("ConfidentialVestingWallet", function () {
     nextTimestamp = this.startTimestamp + this.duration / BigInt(4);
     await ethers.provider.send("evm_setNextBlockTimestamp", [nextTimestamp.toString()]);
 
-    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release();
+    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release(this.confidentialERC20Address);
     await tx.wait();
 
     // It should be equal to 1/4 of the amount vested.
     expect(
       await reencryptReleased(
         this.beneficiary,
         this.instance,
+        this.confidentialERC20Address,
         this.confidentialVestingWallet,
         this.confidentialVestingWalletAddress,
       ),
@@ -118,14 +112,15 @@ describe("ConfidentialVestingWallet", function () {
     nextTimestamp = this.startTimestamp + this.duration / BigInt(2);
     await ethers.provider.send("evm_setNextBlockTimestamp", [nextTimestamp.toString()]);
 
-    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release();
+    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release(this.confidentialERC20Address);
     await tx.wait();
 
     // It should be equal to 1/4 of the amount vested since 1/4 was already collected.
     expect(
       await reencryptReleased(
         this.beneficiary,
         this.instance,
+        this.confidentialERC20Address,
         this.confidentialVestingWallet,
         this.confidentialVestingWalletAddress,
       ),
@@ -138,14 +133,15 @@ describe("ConfidentialVestingWallet", function () {
     nextTimestamp = this.startTimestamp + this.duration;
     await ethers.provider.send("evm_setNextBlockTimestamp", [nextTimestamp.toString()]);
 
-    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release();
+    tx = await this.confidentialVestingWallet.connect(this.beneficiary).release(this.confidentialERC20Address);
     await tx.wait();
 
     // It should be equal to 1/2 of the amount vested since 2/4 was already collected.
     expect(
       await reencryptReleased(
         this.beneficiary,
         this.instance,
+        this.confidentialERC20Address,
         this.confidentialVestingWallet,
         this.confidentialVestingWalletAddress,
       ),

diff --git a/test/finance/ConfidentialVestingWalletCliff.fixture.ts b/test/finance/ConfidentialVestingWalletCliff.fixture.ts
@@ -6,15 +6,14 @@ import type { TestConfidentialVestingWalletCliff } from "../../types";
 export async function deployConfidentialVestingWalletCliffFixture(
   account: Signer,
   beneficiaryAddress: string,
-  token: string,
   startTimestamp: bigint,
   duration: bigint,
   cliffSeconds: bigint,
 ): Promise<TestConfidentialVestingWalletCliff> {
   const contractFactory = await ethers.getContractFactory("TestConfidentialVestingWalletCliff");
   const contract = await contractFactory
     .connect(account)
-    .deploy(beneficiaryAddress, token, startTimestamp, duration, cliffSeconds);
+    .deploy(beneficiaryAddress, startTimestamp, duration, cliffSeconds);
   await contract.waitForDeployment();
   return contract;
 }