refactor: wrapper contracts

zama-ai · Nov 29, 2024 · 1466317 · 1466317
1 parent 557c7da
commit 1466317
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 58 deletions.
diff --git a/contracts/token/ERC20/EncryptedERC20Wrapped.sol b/contracts/token/ERC20/EncryptedERC20Wrapped.sol
@@ -4,11 +4,12 @@ pragma solidity ^0.8.24;
 import { IERC20Metadata } from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
-import { EncryptedERC20 } from "./EncryptedERC20.sol";
-
 import "fhevm/lib/TFHE.sol";
 import "fhevm/gateway/GatewayCaller.sol";
 
+import { IEncryptedERC20Wrapped } from "./IEncryptedERC20Wrapped.sol";
+import { EncryptedERC20 } from "./EncryptedERC20.sol";
+
 /**
  * @title             EncryptedERC20Wrapped
  * @notice            This contract allows users to wrap/unwrap trustlessly
@@ -17,34 +18,9 @@ import "fhevm/gateway/GatewayCaller.sol";
  *                    tokens with a fee on transfer. All ERC20 tokens must have decimals
  *                    inferior or equal to 18 decimals but superior or equal to 6 decimals.
  */
-abstract contract EncryptedERC20Wrapped is EncryptedERC20, GatewayCaller {
+abstract contract EncryptedERC20Wrapped is EncryptedERC20, IEncryptedERC20Wrapped, GatewayCaller {
     using SafeERC20 for IERC20Metadata;
 
-    /// @notice Returned if the amount is greater than 2**64.
-    error AmountTooHigh();
-
-    /// @notice Returned if user cannot transfer or mint.
-    error CannotTransferOrUnwrap();
-
-    /// @notice Emitted when token is unwrapped.
-    event Unwrap(address indexed to, uint64 amount);
-
-    /// @notice Emitted if unwrap fails.
-    event UnwrapFail(address account, uint64 amount);
-
-    /// @notice Emitted when token is wrapped.
-    event Wrap(address indexed to, uint64 amount);
-
-    /**
-     * @notice          Keeps track of unwrap information.
-     * @param account   Account that initiates the unwrap request.
-     * @param amount    Amount to be unwrapped.
-     */
-    struct UnwrapRequest {
-        address account;
-        uint64 amount;
-    }
-
     /// @notice ERC20 token that is wrapped.
     IERC20Metadata public immutable ERC20_TOKEN;
 
@@ -113,7 +89,7 @@ abstract contract EncryptedERC20Wrapped is EncryptedERC20, GatewayCaller {
 
         uint64 amountUint64 = uint64(amountAdjusted);
 
-        _unsafeMint(msg.sender, TFHE.asEuint64(amountUint64));
+        _unsafeMint(msg.sender, amountUint64);
         _totalSupply += amountUint64;
 
         emit Wrap(msg.sender, amountUint64);
@@ -129,7 +105,7 @@ abstract contract EncryptedERC20Wrapped is EncryptedERC20, GatewayCaller {
         delete unwrapRequests[requestId];
 
         if (canUnwrap) {
-            _unsafeBurn(unwrapRequest.account, TFHE.asEuint64(unwrapRequest.amount));
+            _unsafeBurn(unwrapRequest.account, unwrapRequest.amount);
             _totalSupply -= unwrapRequest.amount;
 
             /// @dev It does a supply adjustment.
@@ -160,4 +136,12 @@ abstract contract EncryptedERC20Wrapped is EncryptedERC20, GatewayCaller {
         _canTransferOrUnwrap(from);
         super._transferNoEvent(from, to, amount, isTransferable);
     }
+
+    function _unsafeBurn(address account, uint64 amount) internal {
+        euint64 newBalanceAccount = TFHE.sub(_balances[account], amount);
+        _balances[account] = newBalanceAccount;
+        TFHE.allowThis(newBalanceAccount);
+        TFHE.allow(newBalanceAccount, account);
+        emit Transfer(account, address(0), _PLACEHOLDER);
+    }
 }
diff --git a/contracts/token/ERC20/EncryptedWETH.sol b/contracts/token/ERC20/EncryptedWETH.sol
@@ -6,40 +6,17 @@ import { EncryptedERC20 } from "./EncryptedERC20.sol";
 import "fhevm/lib/TFHE.sol";
 import "fhevm/gateway/GatewayCaller.sol";
 
+import { IEncryptedERC20Wrapped } from "./IEncryptedERC20Wrapped.sol";
+
 /**
  * @title             EncryptedWETH
  * @notice            This contract allows users to wrap/unwrap trustlessly
  *                    ETH (or other native tokens) to EncryptedERC20 tokens.
  */
-abstract contract EncryptedWETH is EncryptedERC20, GatewayCaller {
-    /// @notice Returned if the amount is greater than 2**64.
-    error AmountTooHigh();
-
-    /// @notice Returned if user cannot transfer or mint.
-    error CannotTransferOrUnwrap();
-
+abstract contract EncryptedWETH is EncryptedERC20, IEncryptedERC20Wrapped, GatewayCaller {
     /// @notice Returned if ETH transfer fails.
     error ETHTransferFail();
 
-    /// @notice Emitted when encrypted wrapped ether is unwrapped.
-    event Unwrap(address indexed to, uint64 amount);
-
-    /// @notice Emitted if unwrap fails.
-    event UnwrapFail(address account, uint64 amount);
-
-    /// @notice Emitted when ether is wrapped.
-    event Wrap(address indexed to, uint64 amount);
-
-    /**
-     * @notice          Keeps track of unwrap information.
-     * @param account   Account that initiates the unwrap request.
-     * @param amount    Amount to be unwrapped.
-     */
-    struct UnwrapRequest {
-        address account;
-        uint64 amount;
-    }
-
     /// @notice Tracks whether the account can move funds.
     mapping(address account => bool isRestricted) public isAccountRestricted;
 
@@ -105,7 +82,7 @@ abstract contract EncryptedWETH is EncryptedERC20, GatewayCaller {
 
         uint64 amountUint64 = uint64(amountAdjusted);
 
-        _unsafeMint(msg.sender, TFHE.asEuint64(amountUint64));
+        _unsafeMint(msg.sender, amountUint64);
         _totalSupply += amountUint64;
 
         emit Wrap(msg.sender, amountUint64);
@@ -121,7 +98,7 @@ abstract contract EncryptedWETH is EncryptedERC20, GatewayCaller {
         delete unwrapRequests[requestId];
 
         if (canUnwrap) {
-            _unsafeBurn(unwrapRequest.account, TFHE.asEuint64(unwrapRequest.amount));
+            _unsafeBurn(unwrapRequest.account, unwrapRequest.amount);
             _totalSupply -= unwrapRequest.amount;
 
             /// @dev It does a supply adjustment.
@@ -158,4 +135,12 @@ abstract contract EncryptedWETH is EncryptedERC20, GatewayCaller {
         _canTransferOrUnwrap(from);
         super._transferNoEvent(from, to, amount, isTransferable);
     }
+
+    function _unsafeBurn(address account, uint64 amount) internal {
+        euint64 newBalanceAccount = TFHE.sub(_balances[account], amount);
+        _balances[account] = newBalanceAccount;
+        TFHE.allowThis(newBalanceAccount);
+        TFHE.allow(newBalanceAccount, account);
+        emit Transfer(account, address(0), _PLACEHOLDER);
+    }
 }
diff --git a/contracts/token/ERC20/IEncryptedERC20Wrapped.sol b/contracts/token/ERC20/IEncryptedERC20Wrapped.sol
@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity ^0.8.24;
+
+/**
+ * @title       IEncryptedERC20Wrapped/
+ * @notice      Interface that defines events, errors, and structs for
+ *              contrats that wrap native asset or ERC20 tokens.
+ */
+interface IEncryptedERC20Wrapped {
+    /// @notice Returned if the amount is greater than 2**64.
+    error AmountTooHigh();
+
+    /// @notice Returned if user cannot transfer or mint.
+    error CannotTransferOrUnwrap();
+
+    /// @notice Emitted when token is unwrapped.
+    event Unwrap(address indexed to, uint64 amount);
+
+    /// @notice Emitted if unwrap fails.
+    event UnwrapFail(address account, uint64 amount);
+
+    /// @notice Emitted when token is wrapped.
+    event Wrap(address indexed to, uint64 amount);
+
+    /**
+     * @notice          Keeps track of unwrap information.
+     * @param account   Account that initiates the unwrap request.
+     * @param amount    Amount to be unwrapped.
+     */
+    struct UnwrapRequest {
+        address account;
+        uint64 amount;
+    }
+}