add custom error and fix tests

Author: Amxx

contracts/token/ERC7984/extensions/ERC7984ERC20Wrapper.sol

@@ -30,6 +30,7 @@ abstract contract ERC7984ERC20Wrapper is ERC7984, IERC1363Receiver {
     event UnwrapFinalized(address indexed receiver, euint64 encryptedAmount, uint64 cleartextAmount);
 
     error InvalidUnwrapRequest(euint64 amount);
+    error ERC7984TotalSupplyOverflow();
 
     constructor(IERC20 underlying_) {
         _underlying = underlying_;
@@ -150,17 +151,24 @@ abstract contract ERC7984ERC20Wrapper is ERC7984, IERC1363Receiver {
     }
 
     /**
-     * @dev Returns the underlying balance divided by the {rate}, closely approximating the total supply of wrapped tokens.
+     * @dev Returns the underlying balance divided by the {rate}, over-approximating the total supply of wrapped tokens.
      *
      * NOTE: The return value of this function can be inflated by directly sending underlying tokens to the wrapper contract.
      */
-    function totalSupply() public view virtual returns (uint64) {
-        return SafeCast.toUint64(underlying().balanceOf(address(this)) / rate());
+    function totalSupply() public view virtual returns (uint256) {
+        return underlying().balanceOf(address(this)) / rate();
+    }
+
+    /// @dev Returns the maximum total supply of wrapped tokens supported by the encrypted datatype.
+    function maxTotalSupply() public view virtual returns (uint256) {
+        return type(uint64).max;
     }
 
     /// @dev This function must revert if the new {ERC7984-confidentialTotalSupply} state is invalid.
     function _checkTotalSupply() internal virtual {
-        SafeCast.toUint64(underlying().balanceOf(address(this)) / rate());
+        if (totalSupply() > maxTotalSupply()) {
+            revert ERC7984TotalSupplyOverflow();
+        }
     }
 
     function _update(address from, address to, euint64 amount) internal virtual override returns (euint64) {

test/token/ERC7984/extensions/ERC7984Wrapper.test.ts

@@ -83,14 +83,16 @@ describe('ERC7984Wrapper', function () {
         });
 
         it('max amount works', async function () {
-          const maxAmount = (2n ** 64n - 1n) / 10n ** 6n;
-          await this.token.$_mint(this.holder.address, ethers.parseUnits(maxAmount.toString(), 18));
-          const amountToWrap = ethers.parseUnits(maxAmount.toString(), 18);
+          await this.token.$_mint(this.holder.address, ethers.MaxUint256 / 2n); // mint a lot of tokens
+
+          const rate = await this.wrapper.rate();
+          const maxConfidentialSupply = await this.wrapper.maxTotalSupply();
+          const maxUnderlyingSupply = maxConfidentialSupply * rate;
 
           if (viaCallback) {
-            await this.token.connect(this.holder).transferAndCall(this.wrapper, amountToWrap);
+            await this.token.connect(this.holder).transferAndCall(this.wrapper, maxUnderlyingSupply);
           } else {
-            await this.wrapper.connect(this.holder).wrap(this.holder.address, amountToWrap);
+            await this.wrapper.connect(this.holder).wrap(this.holder.address, maxUnderlyingSupply);
           }
 
           await expect(
@@ -100,22 +102,25 @@ describe('ERC7984Wrapper', function () {
               this.wrapper.target,
               this.holder,
             ),
-          ).to.eventually.equal(ethers.parseUnits(maxAmount.toString(), 6));
+          ).to.eventually.equal(maxConfidentialSupply);
         });
 
         it('amount exceeding max fails', async function () {
-          const maxAmount = (2n ** 64n - 1n) / 10n ** 6n;
-          await this.token.$_mint(this.holder.address, ethers.parseUnits(maxAmount.toString(), 18));
-          const amountToWrap = ethers.parseUnits((maxAmount + 1n).toString(), 18);
+          await this.token.$_mint(this.holder.address, ethers.MaxUint256 / 2n); // mint a lot of tokens
 
-          let tx;
-          if (viaCallback) {
-            tx = this.token.connect(this.holder).transferAndCall(this.wrapper, amountToWrap);
-          } else {
-            tx = this.wrapper.connect(this.holder).wrap(this.holder.address, amountToWrap);
-          }
+          const rate = await this.wrapper.rate();
+          const maxConfidentialSupply = await this.wrapper.maxTotalSupply();
+          const maxUnderlyingSupply = maxConfidentialSupply * rate;
 
-          await expect(tx).to.be.revertedWithCustomError(this.wrapper, 'SafeCastOverflowedUintDowncast');
+          // first deposit close to the max
+          await this.wrapper.connect(this.holder).wrap(this.holder.address, maxUnderlyingSupply);
+
+          // try to deposit more, causing the total supply to exceed the max supported amount
+          await expect(
+            viaCallback
+              ? this.token.connect(this.holder).transferAndCall(this.wrapper, rate)
+              : this.wrapper.connect(this.holder).wrap(this.holder.address, rate),
+          ).to.be.revertedWithCustomError(this.wrapper, 'ERC7984TotalSupplyOverflow');
         });
 
         if (viaCallback) {