refactor: EncryptedErrors for EncryptedERC20

contracts/token/ERC20/IEncryptedERC20.sol

@@ -1,4 +1,4 @@
-// SPDX-License-Identifier: UNLICENSED
+// SPDX-License-Identifier: BSD-3-Clause-Clear
 pragma solidity ^0.8.24;
 
 import "fhevm/lib/TFHE.sol";

contracts/token/ERC20/extensions/EncryptedERC20WithErrors.sol

@@ -0,0 +1,154 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity ^0.8.24;
+
+import "fhevm/lib/TFHE.sol";
+import { EncryptedERC20 } from "../EncryptedERC20.sol";
+import { EncryptedErrors } from "../../../utils/EncryptedErrors.sol";
+
+/**
+ * @title EncryptedERC20WithErrors
+ * @notice      This contract implements an encrypted ERC20-like token with confidential balances using
+ *              Zama's FHE (Fully Homomorphic Encryption) library.
+ * @dev         It supports standard ERC20 functions such as transferring tokens, minting,
+ *              and setting allowances, but uses encrypted data types.
+ *              The total supply is not encrypted.
+ *              It also supports error handling for encrypted errors.
+ */
+
+abstract contract EncryptedERC20WithErrors is EncryptedERC20, EncryptedErrors {
+    /**
+     * @notice Emitted when tokens are moved from one account (`from`) to
+     * another (`to`).
+     */
+    event TransferWithErrorHandling(address indexed from, address indexed to, uint256 transferId);
+
+    /**
+     * @notice Error codes allow tracking (in the storage) whether a transfer worked.
+     * @dev    NO_ERROR: the transfer worked as expected
+     *         UNSUFFICIENT_BALANCE: the transfer failed because the
+     *         from balances were strictly inferior to the amount to transfer.
+     *         UNSUFFICIENT_APPROVAL: the transfer failed because the sender allowance
+     *         was strictly lower than the amount to transfer.
+     */
+    enum ErrorCodes {
+        NO_ERROR,
+        UNSUFFICIENT_BALANCE,
+        UNSUFFICIENT_APPROVAL
+    }
+
+    /// @notice Keeps track of the current transferId.
+    uint256 private _transferIdCounter;
+
+    /// @notice A mapping from transferId to the error code.
+    mapping(uint256 transferId => euint8 errorCode) internal _errorCodeForTransferId;
+
+    /**
+     * @param name_ Name of the token.
+     * @param symbol_ Symbol.
+     */
+    constructor(
+        string memory name_,
+        string memory symbol_
+    ) EncryptedERC20(name_, symbol_) EncryptedErrors(uint8(type(ErrorCodes).max)) {}
+
+    /**
+     * @notice See {IEncryptedERC20-transfer}.
+     */
+    function transfer(address to, euint64 amount) public virtual override returns (bool) {
+        _isSenderAllowedForAmount(amount);
+
+        // Make sure the owner has enough tokens
+        ebool canTransfer = TFHE.le(amount, _balances[msg.sender]);
+
+        euint8 errorCode = TFHE.select(
+            canTransfer,
+            _errorCodes[uint8(ErrorCodes.NO_ERROR)],
+            _errorCodes[uint8(ErrorCodes.UNSUFFICIENT_BALANCE)]
+        );
+
+        _transferWithErrorCode(msg.sender, to, amount, canTransfer, errorCode);
+        return true;
+    }
+
+    /**
+     * @notice See {IEncryptedERC20-transferFrom}.
+     */
+    function transferFrom(address from, address to, euint64 amount) public virtual override returns (bool) {
+        _isSenderAllowedForAmount(amount);
+        address spender = msg.sender;
+        (ebool isTransferable, euint8 errorCode) = _updateAllowanceWithErrorCode(from, spender, amount);
+        _transferWithErrorCode(from, to, amount, isTransferable, errorCode);
+        return true;
+    }
+
+    /**
+     * @notice Returns the error code corresponding to `transferId`.
+     */
+    function getErrorCodeForTransferId(uint256 transferId) external view virtual returns (euint8 errorCode) {
+        return _errorCodeForTransferId[transferId];
+    }
+
+    function _transferWithErrorCode(
+        address from,
+        address to,
+        euint64 amount,
+        ebool isTransferable,
+        euint8 errorCode
+    ) internal virtual {
+        // Add to the balance of `to` and subract from the balance of `from`.
+        euint64 transferValue = TFHE.select(isTransferable, amount, TFHE.asEuint64(0));
+        euint64 newBalanceTo = TFHE.add(_balances[to], transferValue);
+        _balances[to] = newBalanceTo;
+
+        TFHE.allow(newBalanceTo, address(this));
+        TFHE.allow(newBalanceTo, to);
+
+        euint64 newBalanceFrom = TFHE.sub(_balances[from], transferValue);
+        _balances[from] = newBalanceFrom;
+
+        TFHE.allow(newBalanceFrom, address(this));
+        TFHE.allow(newBalanceFrom, from);
+
+        emit TransferWithErrorHandling(from, to, _transferIdCounter);
+
+        // Set error code in the storage and increment
+        _errorCodeForTransferId[_transferIdCounter++] = errorCode;
+
+        TFHE.allowThis(errorCode);
+        TFHE.allow(errorCode, from);
+        TFHE.allow(errorCode, to);
+    }
+
+    function _updateAllowanceWithErrorCode(
+        address owner,
+        address spender,
+        euint64 amount
+    ) internal virtual returns (ebool isTransferable, euint8 errorCode) {
+        euint64 currentAllowance = _allowance(owner, spender);
+
+        // Make sure sure the allowance suffices
+        ebool allowedTransfer = TFHE.le(amount, currentAllowance);
+
+        errorCode = TFHE.select(
+            allowedTransfer,
+            _errorCodes[uint8(ErrorCodes.UNSUFFICIENT_APPROVAL)],
+            _errorCodes[uint8(ErrorCodes.NO_ERROR)]
+        );
+
+        // Make sure the owner has enough tokens
+        ebool canTransfer = TFHE.le(amount, _balances[owner]);
+
+        errorCode = TFHE.select(
+            TFHE.eq(errorCode, 0),
+            TFHE.select(
+                canTransfer,
+                _errorCodes[uint8(ErrorCodes.UNSUFFICIENT_BALANCE)],
+                _errorCodes[uint8(ErrorCodes.NO_ERROR)]
+            ),
+            errorCode
+        );
+
+        isTransferable = TFHE.and(canTransfer, allowedTransfer);
+        _approve(owner, spender, TFHE.select(isTransferable, TFHE.sub(currentAllowance, amount), currentAllowance));
+    }
+}

contracts/token/ERC20/extensions/EncryptedERC20WithErrorsMintable.sol

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity ^0.8.24;
+
+import "fhevm/lib/TFHE.sol";
+import { Ownable2Step, Ownable } from "@openzeppelin/contracts/access/Ownable2Step.sol";
+
+import { EncryptedERC20WithErrors } from "./EncryptedERC20WithErrors.sol";
+
+/**
+ * @title       EncryptedERC20WithErrorsMintable
+ * @notice      This contract inherits EncryptedERC20WithErrors.
+ * @dev         It allows an owner to mint tokens. Mint amounts are public.
+ */
+contract EncryptedERC20WithErrorsMintable is Ownable2Step, EncryptedERC20WithErrors {
+    /**
+     * @notice Emitted when `amount` tokens are minted to one account (`to`).
+     */
+    event Mint(address indexed to, uint64 amount);
+
+    /**
+     * @param name_ Name of the token.
+     * @param symbol_ Symbol.
+     * @param owner_ Owner address.
+     */
+    constructor(
+        string memory name_,
+        string memory symbol_,
+        address owner_
+    ) Ownable(owner_) EncryptedERC20WithErrors(name_, symbol_) {}
+
+    /**
+     * @notice Mint tokens.
+     * @param amount Amount of tokens to mint.
+     */
+    function mint(uint64 amount) public onlyOwner {
+        _balances[msg.sender] = TFHE.add(_balances[msg.sender], amount);
+        TFHE.allow(_balances[msg.sender], address(this));
+        TFHE.allow(_balances[msg.sender], msg.sender);
+        /// @dev Since _totalSupply is not encrypted and _totalSupply >= balances[msg.sender],
+        /// the next line contains an overflow check for the encrypted operation above.
+        _totalSupply = _totalSupply + amount;
+        emit Mint(msg.sender, amount);
+    }
+}

contracts/utils/EncryptedErrors.sol

@@ -1,140 +1,49 @@
 // SPDX-License-Identifier: BSD-3-Clause-Clear
-
 pragma solidity ^0.8.24;
 
 import "fhevm/lib/TFHE.sol";
 
 /**
- * This abstract contract is used for error handling in the fhEVM.
- *
- * Error codes are trivially encrypted during construction inside the `errorCodes` array.
- *
- * WARNING: `errorCodes[0]` should always refer to the `NO_ERROR` code, by default.
- *
+ * @notice This abstract contract is used for error handling in the fhEVM.
+ *         Error codes are encrypted in the constructor inside the `errorCodes` mapping.
+ * @dev    `errorCodes[0]` should always refer to the `NO_ERROR` code, by default.
  */
 abstract contract EncryptedErrors {
-    uint8 private immutable totalNumErrors;
-    euint8[] private errorCodes;
-    uint256 private counterErrors; // used to keep track of each error index
+    /// @notice The total number of errors is equal to zero.
+    error TotalNumberErrorsEqualToZero();
+
+    /// @notice Total number of errors.
+    uint8 private immutable _TOTAL_NUMBER_ERRORS;
 
-    // A mapping from errorId to the errorCode
-    mapping(uint256 => euint8) private errorCodesMapping;
+    /// @notice Mapping of error codes.
+    /// @dev    It does not use arrays they are more expensive than mappings.
+    mapping(uint8 errorCode => euint8 encryptedErrorCode) internal _errorCodes;
 
     /**
      * @notice Sets the non-null value for `numErrors` corresponding to the total number of errors.
-     * @param numErrors the total number of different errors.
-     * @dev `numErrors` must be non-null, note that `errorCodes[0]` corresponds to the `NO_ERROR` code.
+     * @param totalNumberErrors_ total number of different errors.
+     * @dev `numErrors` must be non-null (`errorCodes[0]` corresponds to the `NO_ERROR` code).
      */
-    constructor(uint8 numErrors) {
-        require(numErrors != 0, "numErrors must be greater than 0");
-        for (uint256 i = 0; i <= numErrors; i++) {
-            errorCodes.push(TFHE.asEuint8(i));
+    constructor(uint8 totalNumberErrors_) {
+        if (totalNumberErrors_ == 0) {
+            revert TotalNumberErrorsEqualToZero();
         }
-        totalNumErrors = numErrors;
-    }
 
-    /**
-     * @notice Returns the encrypted error code at index `indexCode`.
-     * @param indexCode the index of the requested error code.
-     * @return the encrypted error code located at `indexCode`.
-     */
-    function getErrorCode(uint8 indexCode) internal view returns (euint8) {
-        return errorCodes[indexCode];
-    }
-
-    /**
-     * @notice Returns the total number of error codes currently stored in `errorCodesMapping`.
-     * @return the number of error codes stored in the `errorCodesMapping` mapping.
-     */
-    function getErrorCounter() internal view returns (uint256) {
-        return counterErrors;
-    }
-
-    /**
-     * @notice Returns the total number of the possible errors.
-     * @return the total number of the different possible errors.
-     */
-    function getNumErrors() internal view returns (uint8) {
-        return totalNumErrors;
-    }
-
-    /**
-     * @notice Returns the encrypted error code which was stored in the mapping at key `errorId`.
-     * @param errorId the requested key stored in the `errorCodesMapping` mapping.
-     * @return the encrypted error code located at the `errorId` key.
-     * @dev `errorId` must be a valid id, i.e below the error counter.
-     */
-    function getError(uint256 errorId) internal view returns (euint8) {
-        require(errorId < counterErrors, "errorId must be a valid id");
-        return errorCodesMapping[errorId];
-    }
-
-    /**
-     * @notice Computes an encrypted error code, result will be either a reencryption of
-     * `errorCodes[indexCode]` if `condition` is an encrypted `true` or of `NO_ERROR` otherwise.
-     * @param condition the encrypted boolean used in the cmux.
-     * @param indexCode the index of the selected error code if `condition` encrypts `true`.
-     * @return the reencrypted error code depending on `condition` value.
-     * @dev `indexCode` must be non-null and below the total number of error codes.
-     */
-    function defineErrorIf(ebool condition, uint8 indexCode) internal view returns (euint8) {
-        require(indexCode != 0, "indexCode must be greater than 0");
-        require(indexCode <= totalNumErrors, "indexCode must be a valid error code");
-        euint8 errorCode = TFHE.select(condition, errorCodes[indexCode], errorCodes[0]);
-        return errorCode;
-    }
-
-    /**
-     * @notice Does the opposite of `defineErrorIf`, i.e result will be either a reencryption of
-     * `errorCodes[indexCode]` if `condition` is an encrypted `false` or of `NO_ERROR` otherwise.
-     * @param condition the encrypted boolean used in the cmux.
-     * @param indexCode the index of the selected error code if `condition` encrypts `false`.
-     * @return the reencrypted error code depending on `condition` value.
-     * @dev `indexCode` must be non-null and below the total number of error codes.
-     */
-    function defineErrorIfNot(ebool condition, uint8 indexCode) internal view returns (euint8) {
-        require(indexCode != 0, "indexCode must be greater than 0");
-        require(indexCode <= totalNumErrors, "indexCode must be a valid error code");
-        euint8 errorCode = TFHE.select(condition, errorCodes[0], errorCodes[indexCode]);
-        return errorCode;
-    }
-
-    /**
-     * @notice Computes an encrypted error code, result will be either a reencryption of
-     * `errorCodes[indexCode]` if `condition` is an encrypted `true` or of `errorCode` otherwise.
-     * @param condition the encrypted boolean used in the cmux.
-     * @param errorCode the selected error code if `condition` encrypts `true`.
-     * @return the reencrypted error code depending on `condition` value.
-     * @dev `indexCode` must be below the total number of error codes.
-     */
-    function changeErrorIf(ebool condition, uint8 indexCode, euint8 errorCode) internal view returns (euint8) {
-        require(indexCode <= totalNumErrors, "indexCode must be a valid error code");
-        return TFHE.select(condition, errorCodes[indexCode], errorCode);
-    }
+        for (uint8 i; i <= totalNumberErrors_; i++) {
+            euint8 errorCode = TFHE.asEuint8(i);
+            _errorCodes[i] = errorCode;
+            TFHE.allowThis(errorCode);
+        }
 
-    /**
-     * @notice Does the opposite of `changeErrorIf`, i.e result will be either a reencryption of
-     * `errorCodes[indexCode]` if `condition` is an encrypted `false` or of `errorCode` otherwise.
-     * @param condition the encrypted boolean used in the cmux.
-     * @param errorCode the selected error code if `condition` encrypts `false`.
-     * @return the reencrypted error code depending on `condition` value.
-     * @dev `indexCode` must be below the total number of error codes.
-     */
-    function changeErrorIfNot(ebool condition, uint8 indexCode, euint8 errorCode) internal view returns (euint8) {
-        require(indexCode <= totalNumErrors, "indexCode must be a valid error code");
-        return TFHE.select(condition, errorCode, errorCodes[indexCode]);
+        _TOTAL_NUMBER_ERRORS = totalNumberErrors_;
     }
 
     /**
-     * @notice Saves `errorCode` in storage, in the `errorCodesMapping` mapping, at the lowest unused key.
-     * This is the only stateful function of `EncryptedErrors` abstract contract.
-     * @param errorCode the encrypted error code to be saved in storage.
-     * @return the `errorId` key in `errorCodesMapping` where `errorCode` is stored.
+     * @notice Returns the total number of errors.
+     * @return totalNumberErrors total number of errors.
+     * @dev    It does not count `NO_ERROR` as one of the errors.
      */
-    function saveError(euint8 errorCode) internal returns (uint256) {
-        uint256 errorId = counterErrors;
-        counterErrors++;
-        errorCodesMapping[errorId] = errorCode;
-        return errorId;
+    function getTotalNumberErrors() external view returns (uint8 totalNumberErrors) {
+        return _TOTAL_NUMBER_ERRORS;
     }
 }