From 61729c23b9292f6b0cb7006fdca718d5bd88a307 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 12 Aug 2025 12:34:10 +0200
Subject: [PATCH 01/52] Init `ERC7984Rwa` extension.

---
 .changeset/new-crews-boil.md                |   5 +
 contracts/interfaces/draft-IERCXXXXCRwa.sol |  65 ++++++++++++
 contracts/token/extensions/ERC7984Rwa.sol   | 109 ++++++++++++++++++++
 3 files changed, 179 insertions(+)
 create mode 100644 .changeset/new-crews-boil.md
 create mode 100644 contracts/interfaces/draft-IERCXXXXCRwa.sol
 create mode 100644 contracts/token/extensions/ERC7984Rwa.sol

diff --git a/.changeset/new-crews-boil.md b/.changeset/new-crews-boil.md
new file mode 100644
index 00000000..e8c18888
--- /dev/null
+++ b/.changeset/new-crews-boil.md
@@ -0,0 +1,5 @@
+---
+'openzeppelin-confidential-contracts': minor
+---
+
+Add `ERC7984Rwa` extension.
diff --git a/contracts/interfaces/draft-IERCXXXXCRwa.sol b/contracts/interfaces/draft-IERCXXXXCRwa.sol
new file mode 100644
index 00000000..58348480
--- /dev/null
+++ b/contracts/interfaces/draft-IERCXXXXCRwa.sol
@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import {IConfidentialFungibleToken} from "./IConfidentialFungibleToken.sol";
+
+/// @dev Interface for confidential RWA contracts.
+interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
+    /// @dev Emmited when the ownership of the contract changes.
+    event OwnershipTransferred(address indexed previousOwner, address indexed newOwner);
+    /// @dev Emmited when the contract is paused.
+    event Paused(address account);
+    /// @dev Emmited when the contract is unpaused.
+    event Unpaused(address account);
+
+    /// @dev The caller account is not authorized to perform an operation.
+    error OwnableUnauthorizedAccount(address account);
+    /// @dev The owner is not a valid owner account. (eg. `address(0)`)
+    error OwnableInvalidOwner(address owner);
+    /// @dev The operation failed because the contract is paused.
+    error EnforcedPause();
+    /// @dev The operation failed because the contract is not paused.
+    error ExpectedPause();
+
+    /// @dev Gets the address of the owner
+    function owner() external view returns (address);
+    /// @dev Transfers contract ownership to an account.
+    function transferOwnership(address _newOwner) external;
+    /// @dev Returns true if the contract is paused, and false otherwise.
+    function paused() external view returns (bool);
+    /// @dev Pauses contract.
+    function pause() external;
+    /// @dev Unpauses contract.
+    function unpause() external;
+    /// @dev Returns the confidential frozen balance of an account.
+    function confidentialFrozen(address acount) external view returns (euint64);
+    /// @dev Sets confidential amount of token for an account as frozen with proof.
+    function setConfidentialFrozen(address acount, externalEuint64 encryptedAmount, bytes calldata inputProof) external;
+    /// @dev Sets confidential amount of token for an account as frozen.
+    function setConfidentialFrozen(address acount, euint64 encryptedAmount) external;
+    /// @dev Receives and executes a batch of function calls on this contract.
+    function multicall(bytes[] calldata data) external returns (bytes[] memory results);
+    /// @dev Mints confidential amount of tokens to account with proof.
+    function mint(address to, externalEuint64 encryptedAmount, bytes calldata inputProof) external returns (euint64);
+    /// @dev Mints confidential amount of tokens to account.
+    function mint(address to, euint64 encryptedAmount) external;
+    /// @dev Burns confidential amount of tokens from account with proof.
+    function burn(
+        address account,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) external returns (euint64);
+    /// @dev Burns confidential amount of tokens from account.
+    function burn(address account, euint64 encryptedAmount) external returns (euint64);
+    /// @dev Forces transfer of confidential amount of tokens from account to account with proof by skipping compliance checks.
+    function forceTransfer(
+        address from,
+        address to,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) external returns (euint64);
+    /// @dev Forces transfer of confidential amount of tokens from account to account by skipping compliance checks.
+    function forceTransfer(address from, address to, euint64 encryptedAmount) external returns (euint64);
+}
diff --git a/contracts/token/extensions/ERC7984Rwa.sol b/contracts/token/extensions/ERC7984Rwa.sol
new file mode 100644
index 00000000..2df8eb8b
--- /dev/null
+++ b/contracts/token/extensions/ERC7984Rwa.sol
@@ -0,0 +1,109 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {ERC165, IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
+import {Pausable} from "@openzeppelin/contracts/utils/Pausable.sol";
+import {IERCXXXXCRwa} from "./../../interfaces/draft-IERCXXXXCRwa.sol";
+import {ConfidentialFungibleToken} from "./../ConfidentialFungibleToken.sol";
+
+/**
+ * @dev Extension of {ConfidentialFungibleToken} supporting confidential Real World Assets.
+ */
+abstract contract ERC7984Rwa is ConfidentialFungibleToken, Ownable, Pausable, Multicall, ERC165 {
+    /// @dev The caller account is not authorized to perform the operation.
+    error UnauthorizedSender(address account);
+    /// @dev The transfer does not follow token compliance.
+    error UncompliantTransfer(address from, address to, euint64 encryptedAmount);
+
+    constructor() {}
+
+    /// @dev Checks the sender is the owner or an authorized agent.
+    modifier onlyOwnerOrAgent() {
+        require(
+            _msgSender() == owner(),
+            //TODO: Add agent condition
+            UnauthorizedSender(_msgSender())
+        );
+        _;
+    }
+
+    /// @inheritdoc ERC165
+    function supportsInterface(bytes4 interfaceId) public view virtual override returns (bool) {
+        return interfaceId == type(IERCXXXXCRwa).interfaceId || super.supportsInterface(interfaceId);
+    }
+
+    /// @dev Pauses contract.
+    function pause() public virtual onlyOwnerOrAgent {
+        _pause();
+    }
+
+    /// @dev Unpauses contract.
+    function unpause() public virtual onlyOwnerOrAgent {
+        _unpause();
+    }
+
+    /// @dev Mints confidential amount of tokens to account with proof.
+    function mint(
+        address to,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) public virtual returns (euint64) {
+        return mint(to, FHE.fromExternal(encryptedAmount, inputProof));
+    }
+
+    /// @dev Mints confidential amount of tokens to account.
+    function mint(address to, euint64 encryptedAmount) public virtual onlyOwnerOrAgent returns (euint64) {
+        return _mint(to, encryptedAmount);
+    }
+
+    /// @dev Burns confidential amount of tokens from account with proof.
+    function burn(
+        address account,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) public virtual returns (euint64) {
+        return burn(account, FHE.fromExternal(encryptedAmount, inputProof));
+    }
+
+    /// @dev Burns confidential amount of tokens from account.
+    function burn(address account, euint64 encryptedAmount) public virtual onlyOwnerOrAgent returns (euint64) {
+        return _burn(account, encryptedAmount);
+    }
+
+    /// @dev Forces transfer of confidential amount of tokens from account to account with proof by skipping compliance checks.
+    function forceTransfer(
+        address from,
+        address to,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) public virtual returns (euint64) {
+        return forceTransfer(from, to, FHE.fromExternal(encryptedAmount, inputProof));
+    }
+
+    /// @dev Forces transfer of confidential amount of tokens from account to account by skipping compliance checks.
+    function forceTransfer(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) public virtual onlyOwnerOrAgent returns (euint64) {
+        //TODO: Add checks
+        return super._update(from, to, encryptedAmount);
+    }
+
+    function _update(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) internal override whenNotPaused returns (euint64) {
+        //TODO: Add checks
+        require(_isCompliantTransfer(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+        return super._update(from, to, encryptedAmount);
+    }
+
+    /// @dev Checks if a transfer follows token compliance.
+    function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+}

From 639e5a2c12fe1dd6d5ac060cc365dcdd64aa60bb Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 12 Aug 2025 12:41:42 +0200
Subject: [PATCH 02/52] Fix typos

---
 contracts/interfaces/draft-IERCXXXXCRwa.sol | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/contracts/interfaces/draft-IERCXXXXCRwa.sol b/contracts/interfaces/draft-IERCXXXXCRwa.sol
index 58348480..400dcf4a 100644
--- a/contracts/interfaces/draft-IERCXXXXCRwa.sol
+++ b/contracts/interfaces/draft-IERCXXXXCRwa.sol
@@ -7,11 +7,11 @@ import {IConfidentialFungibleToken} from "./IConfidentialFungibleToken.sol";
 
 /// @dev Interface for confidential RWA contracts.
 interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
-    /// @dev Emmited when the ownership of the contract changes.
+    /// @dev Emitted when the ownership of the contract changes.
     event OwnershipTransferred(address indexed previousOwner, address indexed newOwner);
-    /// @dev Emmited when the contract is paused.
+    /// @dev Emitted when the contract is paused.
     event Paused(address account);
-    /// @dev Emmited when the contract is unpaused.
+    /// @dev Emitted when the contract is unpaused.
     event Unpaused(address account);
 
     /// @dev The caller account is not authorized to perform an operation.
@@ -34,11 +34,15 @@ interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
     /// @dev Unpauses contract.
     function unpause() external;
     /// @dev Returns the confidential frozen balance of an account.
-    function confidentialFrozen(address acount) external view returns (euint64);
+    function confidentialFrozen(address account) external view returns (euint64);
     /// @dev Sets confidential amount of token for an account as frozen with proof.
-    function setConfidentialFrozen(address acount, externalEuint64 encryptedAmount, bytes calldata inputProof) external;
+    function setConfidentialFrozen(
+        address account,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) external;
     /// @dev Sets confidential amount of token for an account as frozen.
-    function setConfidentialFrozen(address acount, euint64 encryptedAmount) external;
+    function setConfidentialFrozen(address account, euint64 encryptedAmount) external;
     /// @dev Receives and executes a batch of function calls on this contract.
     function multicall(bytes[] calldata data) external returns (bytes[] memory results);
     /// @dev Mints confidential amount of tokens to account with proof.

From 85546dddd2cfe6be0bdc7d647531f828ea96581a Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 12 Aug 2025 18:08:21 +0200
Subject: [PATCH 03/52] Add agent role

---
 contracts/interfaces/draft-IERCXXXXCRwa.sol   |  17 +--
 .../token/extensions/ERC7984Freezable.sol     |  17 +++
 contracts/token/extensions/ERC7984Rwa.sol     | 101 ++++++++++++++----
 3 files changed, 109 insertions(+), 26 deletions(-)
 create mode 100644 contracts/token/extensions/ERC7984Freezable.sol

diff --git a/contracts/interfaces/draft-IERCXXXXCRwa.sol b/contracts/interfaces/draft-IERCXXXXCRwa.sol
index 400dcf4a..3586e90c 100644
--- a/contracts/interfaces/draft-IERCXXXXCRwa.sol
+++ b/contracts/interfaces/draft-IERCXXXXCRwa.sol
@@ -2,13 +2,12 @@
 pragma solidity ^0.8.24;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {IConfidentialFungibleToken} from "./IConfidentialFungibleToken.sol";
 
 /// @dev Interface for confidential RWA contracts.
-interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
-    /// @dev Emitted when the ownership of the contract changes.
-    event OwnershipTransferred(address indexed previousOwner, address indexed newOwner);
+interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165, IAccessControl {
     /// @dev Emitted when the contract is paused.
     event Paused(address account);
     /// @dev Emitted when the contract is unpaused.
@@ -23,10 +22,6 @@ interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
     /// @dev The operation failed because the contract is not paused.
     error ExpectedPause();
 
-    /// @dev Gets the address of the owner
-    function owner() external view returns (address);
-    /// @dev Transfers contract ownership to an account.
-    function transferOwnership(address _newOwner) external;
     /// @dev Returns true if the contract is paused, and false otherwise.
     function paused() external view returns (bool);
     /// @dev Pauses contract.
@@ -35,6 +30,8 @@ interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
     function unpause() external;
     /// @dev Returns the confidential frozen balance of an account.
     function confidentialFrozen(address account) external view returns (euint64);
+    /// @dev Returns the available (unfrozen) balance of an account. Up to {confidentialBalanceOf}.
+    function confidentialAvailable(address account) external returns (euint64);
     /// @dev Sets confidential amount of token for an account as frozen with proof.
     function setConfidentialFrozen(
         address account,
@@ -67,3 +64,9 @@ interface IERCXXXXCRwa is IConfidentialFungibleToken, IERC165 {
     /// @dev Forces transfer of confidential amount of tokens from account to account by skipping compliance checks.
     function forceTransfer(address from, address to, euint64 encryptedAmount) external returns (euint64);
 }
+
+/// @dev Interface for confidential RWA compliance.
+interface IERCXXXXCRWACompliance {
+    /// @dev Checks if a transfer follows token compliance.
+    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
+}
diff --git a/contracts/token/extensions/ERC7984Freezable.sol b/contracts/token/extensions/ERC7984Freezable.sol
new file mode 100644
index 00000000..47b1819f
--- /dev/null
+++ b/contracts/token/extensions/ERC7984Freezable.sol
@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
+
+abstract contract ERC7984Freezable {
+    function confidentialFrozen(address account) public view virtual returns (euint64);
+    function confidentialAvailable(address account) public virtual returns (euint64);
+    function setConfidentialFrozen(
+        address account,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) public virtual;
+    function setConfidentialFrozen(address account, euint64 encryptedAmount) public virtual;
+    function _checkFreezer() internal virtual;
+}
diff --git a/contracts/token/extensions/ERC7984Rwa.sol b/contracts/token/extensions/ERC7984Rwa.sol
index 2df8eb8b..11c4b314 100644
--- a/contracts/token/extensions/ERC7984Rwa.sol
+++ b/contracts/token/extensions/ERC7984Rwa.sol
@@ -3,49 +3,93 @@
 pragma solidity ^0.8.27;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
-import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
 import {ERC165, IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
 import {Pausable} from "@openzeppelin/contracts/utils/Pausable.sol";
 import {IERCXXXXCRwa} from "./../../interfaces/draft-IERCXXXXCRwa.sol";
 import {ConfidentialFungibleToken} from "./../ConfidentialFungibleToken.sol";
+import {ERC7984Freezable} from "./ERC7984Freezable.sol";
 
 /**
  * @dev Extension of {ConfidentialFungibleToken} supporting confidential Real World Assets.
  */
-abstract contract ERC7984Rwa is ConfidentialFungibleToken, Ownable, Pausable, Multicall, ERC165 {
+abstract contract ERC7984Rwa is
+    ConfidentialFungibleToken,
+    Pausable,
+    ERC7984Freezable,
+    Multicall,
+    ERC165,
+    AccessControl
+{
+    bytes32 public constant AGENT_ROLE = keccak256("AGENT_ROLE");
+
     /// @dev The caller account is not authorized to perform the operation.
     error UnauthorizedSender(address account);
     /// @dev The transfer does not follow token compliance.
     error UncompliantTransfer(address from, address to, euint64 encryptedAmount);
 
-    constructor() {}
+    constructor(
+        string memory name,
+        string memory symbol,
+        string memory tokenUri
+    ) ConfidentialFungibleToken(name, symbol, tokenUri) {
+        _grantRole(DEFAULT_ADMIN_ROLE, _msgSender());
+    }
 
-    /// @dev Checks the sender is the owner or an authorized agent.
-    modifier onlyOwnerOrAgent() {
-        require(
-            _msgSender() == owner(),
-            //TODO: Add agent condition
-            UnauthorizedSender(_msgSender())
-        );
+    /// @dev Checks if the sender is an admin.
+    modifier onlyAdmin() {
+        require(isAdmin(_msgSender()), UnauthorizedSender(_msgSender()));
+        _;
+    }
+
+    /// @dev Checks if the sender is an agent.
+    modifier onlyAgent() {
+        require(isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
+        _;
+    }
+
+    /// @dev Checks if the sender is an admin or an agent.
+    modifier onlyAdminOrAgent() {
+        require(isAdmin(_msgSender()) || isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
         _;
     }
 
     /// @inheritdoc ERC165
-    function supportsInterface(bytes4 interfaceId) public view virtual override returns (bool) {
+    function supportsInterface(bytes4 interfaceId) public view virtual override(ERC165, AccessControl) returns (bool) {
         return interfaceId == type(IERCXXXXCRwa).interfaceId || super.supportsInterface(interfaceId);
     }
 
     /// @dev Pauses contract.
-    function pause() public virtual onlyOwnerOrAgent {
+    function pause() public virtual onlyAdminOrAgent {
         _pause();
     }
 
     /// @dev Unpauses contract.
-    function unpause() public virtual onlyOwnerOrAgent {
+    function unpause() public virtual onlyAdminOrAgent {
         _unpause();
     }
 
+    /// @dev Returns true if has admin role, false otherwise.
+    function isAdmin(address account) public virtual returns (bool) {
+        return hasRole(DEFAULT_ADMIN_ROLE, account);
+    }
+
+    /// @dev Returns true if agent, false otherwise.
+    function isAgent(address account) public virtual returns (bool) {
+        return hasRole(AGENT_ROLE, account);
+    }
+
+    /// @dev Adds agent.
+    function addAgent(address account) public virtual {
+        _addAgent(account);
+    }
+
+    /// @dev Removes agent.
+    function removeAgent(address account) public virtual {
+        _removeAgent(account);
+    }
+
     /// @dev Mints confidential amount of tokens to account with proof.
     function mint(
         address to,
@@ -56,7 +100,7 @@ abstract contract ERC7984Rwa is ConfidentialFungibleToken, Ownable, Pausable, Mu
     }
 
     /// @dev Mints confidential amount of tokens to account.
-    function mint(address to, euint64 encryptedAmount) public virtual onlyOwnerOrAgent returns (euint64) {
+    function mint(address to, euint64 encryptedAmount) public virtual onlyAdminOrAgent returns (euint64) {
         return _mint(to, encryptedAmount);
     }
 
@@ -70,7 +114,7 @@ abstract contract ERC7984Rwa is ConfidentialFungibleToken, Ownable, Pausable, Mu
     }
 
     /// @dev Burns confidential amount of tokens from account.
-    function burn(address account, euint64 encryptedAmount) public virtual onlyOwnerOrAgent returns (euint64) {
+    function burn(address account, euint64 encryptedAmount) public virtual onlyAdminOrAgent returns (euint64) {
         return _burn(account, encryptedAmount);
     }
 
@@ -89,9 +133,26 @@ abstract contract ERC7984Rwa is ConfidentialFungibleToken, Ownable, Pausable, Mu
         address from,
         address to,
         euint64 encryptedAmount
-    ) public virtual onlyOwnerOrAgent returns (euint64) {
-        //TODO: Add checks
-        return super._update(from, to, encryptedAmount);
+    ) public virtual onlyAdminOrAgent returns (euint64 transferred) {
+        transferred = ConfidentialFungibleToken._update(from, to, encryptedAmount); // bypass frozen & compliance checks
+        setConfidentialFrozen(
+            from,
+            FHE.select(
+                FHE.gt(transferred, confidentialAvailable((from))),
+                confidentialBalanceOf(from),
+                confidentialFrozen(from)
+            )
+        );
+    }
+
+    /// @dev Adds an agent.
+    function _addAgent(address account) internal virtual {
+        _grantRole(AGENT_ROLE, account);
+    }
+
+    /// @dev Removes an agent.
+    function _removeAgent(address account) internal virtual {
+        _grantRole(AGENT_ROLE, account);
     }
 
     function _update(
@@ -99,11 +160,13 @@ abstract contract ERC7984Rwa is ConfidentialFungibleToken, Ownable, Pausable, Mu
         address to,
         euint64 encryptedAmount
     ) internal override whenNotPaused returns (euint64) {
-        //TODO: Add checks
         require(_isCompliantTransfer(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+        // frozen check perfomed through inheritance
         return super._update(from, to, encryptedAmount);
     }
 
+    function _checkFreezer() internal override onlyAdminOrAgent {}
+
     /// @dev Checks if a transfer follows token compliance.
     function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
 }

From 0029ad7583408d827accb680b55e80a95ef3d4fa Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 12 Aug 2025 18:12:29 +0200
Subject: [PATCH 04/52] Update spelling

---
 contracts/token/extensions/ERC7984Rwa.sol | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/contracts/token/extensions/ERC7984Rwa.sol b/contracts/token/extensions/ERC7984Rwa.sol
index 11c4b314..22ec226e 100644
--- a/contracts/token/extensions/ERC7984Rwa.sol
+++ b/contracts/token/extensions/ERC7984Rwa.sol
@@ -161,7 +161,7 @@ abstract contract ERC7984Rwa is
         euint64 encryptedAmount
     ) internal override whenNotPaused returns (euint64) {
         require(_isCompliantTransfer(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
-        // frozen check perfomed through inheritance
+        // frozen check performed through inheritance
         return super._update(from, to, encryptedAmount);
     }
 

From b2174ea6ef0bacfe8ae771e4efb14c6d2019abbe Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Mon, 18 Aug 2025 16:33:21 +0200
Subject: [PATCH 05/52] Add pausable & roles tests

---
 contracts/mocks/token/ERC7984RwaMock.sol  | 42 +++++++++++
 contracts/token/extensions/ERC7984Rwa.sol | 10 +--
 test/token/extensions/ERC7984Rwa.test.ts  | 87 +++++++++++++++++++++++
 3 files changed, 134 insertions(+), 5 deletions(-)
 create mode 100644 contracts/mocks/token/ERC7984RwaMock.sol
 create mode 100644 test/token/extensions/ERC7984Rwa.test.ts

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
new file mode 100644
index 00000000..c7064125
--- /dev/null
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+import {FHE, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
+import {Impl} from "@fhevm/solidity/lib/Impl.sol";
+import {ERC7984Rwa} from "../../token/extensions/ERC7984Rwa.sol";
+
+// solhint-disable func-name-mixedcase
+contract ERC7984RwaMock is ERC7984Rwa, SepoliaConfig {
+    mapping(address account => euint64 encryptedAmount) private _frozenBalances;
+    bool public compliantTransfer;
+
+    constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
+
+    function $_mint(address to, uint64 amount) public returns (euint64 transferred) {
+        return _mint(to, FHE.asEuint64(amount));
+    }
+
+    function _isCompliantTransfer(
+        address /*from*/,
+        address /*to*/,
+        euint64 /*encryptedAmount*/
+    ) internal override returns (bool) {
+        return compliantTransfer;
+    }
+
+    // TODO: Remove all below
+    function confidentialAvailable(address /*account*/) public override returns (euint64) {
+        return FHE.asEuint64(0);
+    }
+    function confidentialFrozen(address account) public view override returns (euint64) {
+        return _frozenBalances[account];
+    }
+    function setConfidentialFrozen(
+        address account,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) public override {}
+    function setConfidentialFrozen(address account, euint64 encryptedAmount) public override {}
+}
diff --git a/contracts/token/extensions/ERC7984Rwa.sol b/contracts/token/extensions/ERC7984Rwa.sol
index 22ec226e..3f84c291 100644
--- a/contracts/token/extensions/ERC7984Rwa.sol
+++ b/contracts/token/extensions/ERC7984Rwa.sol
@@ -71,12 +71,12 @@ abstract contract ERC7984Rwa is
     }
 
     /// @dev Returns true if has admin role, false otherwise.
-    function isAdmin(address account) public virtual returns (bool) {
+    function isAdmin(address account) public view virtual returns (bool) {
         return hasRole(DEFAULT_ADMIN_ROLE, account);
     }
 
     /// @dev Returns true if agent, false otherwise.
-    function isAgent(address account) public virtual returns (bool) {
+    function isAgent(address account) public view virtual returns (bool) {
         return hasRole(AGENT_ROLE, account);
     }
 
@@ -146,13 +146,13 @@ abstract contract ERC7984Rwa is
     }
 
     /// @dev Adds an agent.
-    function _addAgent(address account) internal virtual {
+    function _addAgent(address account) internal virtual onlyAdminOrAgent {
         _grantRole(AGENT_ROLE, account);
     }
 
     /// @dev Removes an agent.
-    function _removeAgent(address account) internal virtual {
-        _grantRole(AGENT_ROLE, account);
+    function _removeAgent(address account) internal virtual onlyAdminOrAgent {
+        _revokeRole(AGENT_ROLE, account);
     }
 
     function _update(
diff --git a/test/token/extensions/ERC7984Rwa.test.ts b/test/token/extensions/ERC7984Rwa.test.ts
new file mode 100644
index 00000000..027265da
--- /dev/null
+++ b/test/token/extensions/ERC7984Rwa.test.ts
@@ -0,0 +1,87 @@
+import { expect } from 'chai';
+import { ethers } from 'hardhat';
+
+/* eslint-disable no-unexpected-multiline */
+describe('ERC7984Rwa', function () {
+  async function deployFixture() {
+    const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
+    const token = await ethers.deployContract('ERC7984RwaMock', ['name', 'symbol', 'uri']);
+    token.connect(anyone);
+    return { token, admin, agent1, agent2, recipient, anyone };
+  }
+
+  describe('Pausable', async function () {
+    it('should pause & unpause', async function () {
+      const { token, admin, agent1 } = await deployFixture();
+      await token
+        .connect(admin)
+        .addAgent(agent1)
+        .then(tx => tx.wait());
+      for (const manager of [admin, agent1]) {
+        expect(await token.paused()).is.false;
+        await token
+          .connect(manager)
+          .pause()
+          .then(tx => tx.wait());
+        expect(await token.paused()).is.true;
+        await token
+          .connect(manager)
+          .unpause()
+          .then(tx => tx.wait());
+        expect(await token.paused()).is.false;
+      }
+    });
+
+    it('should not pause if neither admin nor agent', async function () {
+      const { token, anyone } = await deployFixture();
+      await expect(token.connect(anyone).pause())
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone);
+    });
+
+    it('should not unpause if neither admin nor agent', async function () {
+      const { token, anyone } = await deployFixture();
+      await expect(token.connect(anyone).unpause())
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone);
+    });
+  });
+
+  describe('Roles', async function () {
+    it('should check admin', async function () {
+      const { token, admin, anyone } = await deployFixture();
+      expect(await token.isAdmin(admin)).is.true;
+      expect(await token.isAdmin(anyone)).is.false;
+    });
+
+    it('should check/add/remove agent', async function () {
+      const { token, admin, agent1 } = await deployFixture();
+      expect(await token.isAgent(agent1)).is.false;
+      await token
+        .connect(admin)
+        .addAgent(agent1)
+        .then(tx => tx.wait());
+      expect(await token.isAgent(agent1)).is.true;
+      await token
+        .connect(admin)
+        .removeAgent(agent1)
+        .then(tx => tx.wait());
+      expect(await token.isAgent(agent1)).is.false;
+    });
+
+    it('should not add agent if neither admin nor agent', async function () {
+      const { token, agent1, anyone } = await deployFixture();
+      await expect(token.connect(anyone).addAgent(agent1))
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone);
+    });
+
+    it('should not remove agent if neither admin nor agent', async function () {
+      const { token, agent1, anyone } = await deployFixture();
+      await expect(token.connect(anyone).removeAgent(agent1))
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone);
+    });
+  });
+});
+/* eslint-disable no-unexpected-multiline */

From ce8286c207915dc903371c5273ae1b0d552b2974 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 19 Aug 2025 16:52:09 +0200
Subject: [PATCH 06/52] Add mint/burn/force/transfer tests

---
 contracts/mocks/token/ERC7984RwaMock.sol  |  34 +-
 contracts/token/extensions/ERC7984Rwa.sol |   8 +-
 test/token/extensions/ERC7984Rwa.test.ts  | 483 +++++++++++++++++++++-
 3 files changed, 497 insertions(+), 28 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index c7064125..fbaa06dc 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -3,17 +3,27 @@
 pragma solidity ^0.8.24;
 
 import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
-import {FHE, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, euint64, externalEuint64, ebool} from "@fhevm/solidity/lib/FHE.sol";
 import {Impl} from "@fhevm/solidity/lib/Impl.sol";
 import {ERC7984Rwa} from "../../token/extensions/ERC7984Rwa.sol";
+import {FHESafeMath} from "../../utils/FHESafeMath.sol";
+import {HandleAccessManager} from "../../utils/HandleAccessManager.sol";
 
 // solhint-disable func-name-mixedcase
-contract ERC7984RwaMock is ERC7984Rwa, SepoliaConfig {
+contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
     mapping(address account => euint64 encryptedAmount) private _frozenBalances;
     bool public compliantTransfer;
 
     constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
 
+    function $_setCompliantTransfer() public {
+        compliantTransfer = true;
+    }
+
+    function $_unsetCompliantTransfer() public {
+        compliantTransfer = false;
+    }
+
     function $_mint(address to, uint64 amount) public returns (euint64 transferred) {
         return _mint(to, FHE.asEuint64(amount));
     }
@@ -27,8 +37,14 @@ contract ERC7984RwaMock is ERC7984Rwa, SepoliaConfig {
     }
 
     // TODO: Remove all below
-    function confidentialAvailable(address /*account*/) public override returns (euint64) {
-        return FHE.asEuint64(0);
+    function confidentialAvailable(address account) public override returns (euint64) {
+        (ebool success, euint64 unfrozen) = FHESafeMath.tryDecrease(
+            confidentialBalanceOf(account),
+            confidentialFrozen(account)
+        );
+        unfrozen = FHE.select(success, unfrozen, FHE.asEuint64(0));
+        FHE.allowThis(unfrozen);
+        return unfrozen;
     }
     function confidentialFrozen(address account) public view override returns (euint64) {
         return _frozenBalances[account];
@@ -37,6 +53,12 @@ contract ERC7984RwaMock is ERC7984Rwa, SepoliaConfig {
         address account,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public override {}
-    function setConfidentialFrozen(address account, euint64 encryptedAmount) public override {}
+    ) public override {
+        return setConfidentialFrozen(account, FHE.fromExternal(encryptedAmount, inputProof));
+    }
+    function setConfidentialFrozen(address account, euint64 encryptedAmount) public override {
+        FHE.allowThis(_frozenBalances[account] = encryptedAmount);
+    }
+
+    function _validateHandleAllowance(bytes32 handle) internal view override onlyAdminOrAgent {}
 }
diff --git a/contracts/token/extensions/ERC7984Rwa.sol b/contracts/token/extensions/ERC7984Rwa.sol
index 3f84c291..98f54ae7 100644
--- a/contracts/token/extensions/ERC7984Rwa.sol
+++ b/contracts/token/extensions/ERC7984Rwa.sol
@@ -118,6 +118,7 @@ abstract contract ERC7984Rwa is
         return _burn(account, encryptedAmount);
     }
 
+    //TODO: Rename all to confidential
     /// @dev Forces transfer of confidential amount of tokens from account to account with proof by skipping compliance checks.
     function forceTransfer(
         address from,
@@ -134,14 +135,11 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) public virtual onlyAdminOrAgent returns (euint64 transferred) {
+        euint64 available = confidentialAvailable(from);
         transferred = ConfidentialFungibleToken._update(from, to, encryptedAmount); // bypass frozen & compliance checks
         setConfidentialFrozen(
             from,
-            FHE.select(
-                FHE.gt(transferred, confidentialAvailable((from))),
-                confidentialBalanceOf(from),
-                confidentialFrozen(from)
-            )
+            FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from))
         );
     }
 
diff --git a/test/token/extensions/ERC7984Rwa.test.ts b/test/token/extensions/ERC7984Rwa.test.ts
index 027265da..0f5ff399 100644
--- a/test/token/extensions/ERC7984Rwa.test.ts
+++ b/test/token/extensions/ERC7984Rwa.test.ts
@@ -1,11 +1,16 @@
+import { FhevmType } from '@fhevm/hardhat-plugin';
 import { expect } from 'chai';
-import { ethers } from 'hardhat';
+import { ethers, fhevm } from 'hardhat';
 
 /* eslint-disable no-unexpected-multiline */
 describe('ERC7984Rwa', function () {
   async function deployFixture() {
     const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
     const token = await ethers.deployContract('ERC7984RwaMock', ['name', 'symbol', 'uri']);
+    await token
+      .connect(admin)
+      .addAgent(agent1)
+      .then(tx => tx.wait());
     token.connect(anyone);
     return { token, admin, agent1, agent2, recipient, anyone };
   }
@@ -13,10 +18,6 @@ describe('ERC7984Rwa', function () {
   describe('Pausable', async function () {
     it('should pause & unpause', async function () {
       const { token, admin, agent1 } = await deployFixture();
-      await token
-        .connect(admin)
-        .addAgent(agent1)
-        .then(tx => tx.wait());
       for (const manager of [admin, agent1]) {
         expect(await token.paused()).is.false;
         await token
@@ -55,18 +56,20 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should check/add/remove agent', async function () {
-      const { token, admin, agent1 } = await deployFixture();
-      expect(await token.isAgent(agent1)).is.false;
-      await token
-        .connect(admin)
-        .addAgent(agent1)
-        .then(tx => tx.wait());
-      expect(await token.isAgent(agent1)).is.true;
-      await token
-        .connect(admin)
-        .removeAgent(agent1)
-        .then(tx => tx.wait());
-      expect(await token.isAgent(agent1)).is.false;
+      const { token, admin, agent1, agent2 } = await deployFixture();
+      for (const manager of [admin, agent1]) {
+        expect(await token.isAgent(agent2)).is.false;
+        await token
+          .connect(manager)
+          .addAgent(agent2)
+          .then(tx => tx.wait());
+        expect(await token.isAgent(agent2)).is.true;
+        await token
+          .connect(manager)
+          .removeAgent(agent2)
+          .then(tx => tx.wait());
+        expect(await token.isAgent(agent2)).is.false;
+      }
     });
 
     it('should not add agent if neither admin nor agent', async function () {
@@ -83,5 +86,451 @@ describe('ERC7984Rwa', function () {
         .withArgs(anyone);
     });
   });
+
+  describe('Mintable', async function () {
+    it('should mint by admin or agent', async function () {
+      const { admin, agent1, recipient } = await deployFixture();
+      for (const manager of [admin, agent1]) {
+        const { token } = await deployFixture();
+        const encryptedInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token
+          .connect(manager)
+          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof)
+          .then(tx => tx.wait());
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(balanceHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+        ).to.eventually.equal(100);
+      }
+    });
+
+    it('should not mint if neither admin nor agent', async function () {
+      const { token, recipient, anyone } = await deployFixture();
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), anyone.address)
+        .add64(100)
+        .encrypt();
+      await expect(
+        token
+          .connect(anyone)
+          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+      )
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone);
+    });
+
+    it('should not mint if transfer not compliant', async function () {
+      const { token, admin, recipient } = await deployFixture();
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(100)
+        .encrypt();
+      await expect(
+        token
+          .connect(admin)
+          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+      )
+        .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
+        .withArgs(ethers.ZeroAddress, recipient, encryptedInput.handles[0]);
+    });
+
+    it('should not mint if paused', async function () {
+      const { token, admin, recipient } = await deployFixture();
+      await token
+        .connect(admin)
+        .pause()
+        .then(tx => tx.wait());
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(100)
+        .encrypt();
+      await expect(
+        token
+          .connect(admin)
+          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+      ).to.be.revertedWithCustomError(token, 'EnforcedPause');
+    });
+  });
+
+  describe('Burnable', async function () {
+    it('should burn by admin or agent', async function () {
+      const { admin, agent1, recipient } = await deployFixture();
+      for (const manager of [admin, agent1]) {
+        const { token } = await deployFixture();
+        const encryptedInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token
+          .connect(manager)
+          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof)
+          .then(tx => tx.wait());
+        const balanceBeforeHandle = await token.confidentialBalanceOf(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(balanceBeforeHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceBeforeHandle, await token.getAddress(), manager),
+        ).to.eventually.greaterThan(0);
+        await token
+          .connect(manager)
+          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof)
+          .then(tx => tx.wait());
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(balanceHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+        ).to.eventually.equal(0);
+      }
+    });
+
+    it('should not burn if neither admin nor agent', async function () {
+      const { token, recipient, anyone } = await deployFixture();
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), anyone.address)
+        .add64(100)
+        .encrypt();
+      await expect(
+        token
+          .connect(anyone)
+          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+      )
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone);
+    });
+
+    it('should not mint if transfer not compliant', async function () {
+      const { token, admin, recipient } = await deployFixture();
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(100)
+        .encrypt();
+      await expect(
+        token
+          .connect(admin)
+          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+      )
+        .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
+        .withArgs(recipient, ethers.ZeroAddress, encryptedInput.handles[0]);
+    });
+
+    it('should not burn if paused', async function () {
+      const { token, admin, recipient } = await deployFixture();
+      await token
+        .connect(admin)
+        .pause()
+        .then(tx => tx.wait());
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(100)
+        .encrypt();
+      await expect(
+        token
+          .connect(admin)
+          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+      ).to.be.revertedWithCustomError(token, 'EnforcedPause');
+    });
+  });
+
+  describe('Force transfer', async function () {
+    it('should force transfer by admin or agent', async function () {
+      const { admin, agent1, recipient, anyone } = await deployFixture();
+      for (const manager of [admin, agent1]) {
+        const { token } = await deployFixture();
+        const encryptedMintValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token
+          .connect(manager)
+          ['mint(address,bytes32,bytes)'](
+            recipient,
+            encryptedMintValueInput.handles[0],
+            encryptedMintValueInput.inputProof,
+          )
+          .then(tx => tx.wait());
+        // set frozen (50 available and about to force transfer 25)
+        const encryptedFrozenValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(50)
+          .encrypt();
+        await token
+          .connect(manager)
+          ['setConfidentialFrozen(address,bytes32,bytes)'](
+            recipient,
+            encryptedFrozenValueInput.handles[0],
+            encryptedFrozenValueInput.inputProof,
+          )
+          .then(tx => tx.wait());
+        const encryptedTransferValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(25)
+          .encrypt();
+        await token.$_unsetCompliantTransfer().then(tx => tx.wait());
+        expect(await token.compliantTransfer()).to.be.false;
+        await token
+          .connect(manager)
+          ['forceTransfer(address,address,bytes32,bytes)'](
+            recipient,
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          )
+          .then(tx => tx.wait());
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(balanceHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+        ).to.eventually.equal(75);
+        const frozenHandle = await token.confidentialFrozen(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(frozenHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
+        ).to.eventually.equal(50); // frozen is left unchanged
+      }
+    });
+
+    it('should force transfer even if frozen', async function () {
+      const { admin, agent1, recipient, anyone } = await deployFixture();
+      for (const manager of [admin, agent1]) {
+        const { token } = await deployFixture();
+        const encryptedMintValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token
+          .connect(manager)
+          ['mint(address,bytes32,bytes)'](
+            recipient,
+            encryptedMintValueInput.handles[0],
+            encryptedMintValueInput.inputProof,
+          )
+          .then(tx => tx.wait());
+        // set frozen (only 20 available but about to force transfer 25)
+        const encryptedFrozenValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(80)
+          .encrypt();
+        await token
+          .connect(manager)
+          ['setConfidentialFrozen(address,bytes32,bytes)'](
+            recipient,
+            encryptedFrozenValueInput.handles[0],
+            encryptedFrozenValueInput.inputProof,
+          )
+          .then(tx => tx.wait());
+        const encryptedTransferValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), manager.address)
+          .add64(25)
+          .encrypt();
+        await token.$_unsetCompliantTransfer().then(tx => tx.wait());
+        expect(await token.compliantTransfer()).to.be.false;
+        // should force transfer even if paused
+        await token
+          .connect(manager)
+          .pause()
+          .then(tx => tx.wait());
+        expect(await token.paused()).to.be.true;
+        await token
+          .connect(manager)
+          ['forceTransfer(address,address,bytes32,bytes)'](
+            recipient,
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          )
+          .then(tx => tx.wait());
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(balanceHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+        ).to.eventually.equal(75);
+        const frozenHandle = await token.confidentialFrozen(recipient);
+        await token
+          .connect(manager)
+          .getHandleAllowance(frozenHandle, manager, true)
+          .then(tx => tx.wait());
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
+        ).to.eventually.equal(75); // frozen got reset to available balance
+      }
+    });
+  });
+
+  describe('Transfer', async function () {
+    it('should transfer', async function () {
+      const { token, admin: manager, recipient, anyone } = await deployFixture();
+      const encryptedMintValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), manager.address)
+        .add64(100)
+        .encrypt();
+      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      await token
+        .connect(manager)
+        ['mint(address,bytes32,bytes)'](
+          recipient,
+          encryptedMintValueInput.handles[0],
+          encryptedMintValueInput.inputProof,
+        )
+        .then(tx => tx.wait());
+      // set frozen (50 available and about to transfer 25)
+      const encryptedFrozenValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), manager.address)
+        .add64(50)
+        .encrypt();
+      await token
+        .connect(manager)
+        ['setConfidentialFrozen(address,bytes32,bytes)'](
+          recipient,
+          encryptedFrozenValueInput.handles[0],
+          encryptedFrozenValueInput.inputProof,
+        )
+        .then(tx => tx.wait());
+      const encryptedTransferValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), recipient.address)
+        .add64(25)
+        .encrypt();
+      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      expect(await token.compliantTransfer()).to.be.true;
+      await expect(
+        token
+          .connect(recipient)
+          ['confidentialTransfer(address,bytes32,bytes)'](
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          ),
+      ).to.emit(token, 'ConfidentialTransfer');
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await token.confidentialBalanceOf(recipient),
+          await token.getAddress(),
+          recipient,
+        ),
+      ).to.eventually.equal(75);
+    });
+
+    it('should not transfer if paused', async function () {
+      const { token, admin: manager, recipient, anyone } = await deployFixture();
+      const encryptedTransferValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), recipient.address)
+        .add64(25)
+        .encrypt();
+      await token
+        .connect(manager)
+        .pause()
+        .then(tx => tx.wait());
+      await expect(
+        token
+          .connect(recipient)
+          ['confidentialTransfer(address,bytes32,bytes)'](
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          ),
+      ).to.be.revertedWithCustomError(token, 'EnforcedPause');
+    });
+
+    it('should not transfer if transfer not compliant', async function () {
+      const { token, recipient, anyone } = await deployFixture();
+      const encryptedTransferValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), recipient.address)
+        .add64(25)
+        .encrypt();
+      expect(await token.compliantTransfer()).to.be.false;
+      await expect(
+        token
+          .connect(recipient)
+          ['confidentialTransfer(address,bytes32,bytes)'](
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          ),
+      )
+        .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
+        .withArgs(recipient, anyone, encryptedTransferValueInput.handles[0]);
+    });
+
+    it('should not transfer if frozen', async function () {
+      const { token, admin: manager, recipient, anyone } = await deployFixture();
+      const encryptedMintValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), manager.address)
+        .add64(100)
+        .encrypt();
+      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      await token
+        .connect(manager)
+        ['mint(address,bytes32,bytes)'](
+          recipient,
+          encryptedMintValueInput.handles[0],
+          encryptedMintValueInput.inputProof,
+        )
+        .then(tx => tx.wait());
+      // set frozen (20 available but about to transfer 25)
+      const encryptedFrozenValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), manager.address)
+        .add64(80)
+        .encrypt();
+      await token
+        .connect(manager)
+        ['setConfidentialFrozen(address,bytes32,bytes)'](
+          recipient,
+          encryptedFrozenValueInput.handles[0],
+          encryptedFrozenValueInput.inputProof,
+        )
+        .then(tx => tx.wait());
+      const encryptedTransferValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), recipient.address)
+        .add64(25)
+        .encrypt();
+      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      expect(await token.compliantTransfer()).to.be.true;
+      await expect(
+        token
+          .connect(recipient)
+          ['confidentialTransfer(address,bytes32,bytes)'](
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          ),
+      ).to.emit(token, 'ConfidentialTransfer');
+      /* TODO: Enable when freezable ready
+      // Balance is unchanged
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await token.confidentialBalanceOf(recipient),
+          await token.getAddress(),
+          recipient,
+        ),
+      ).to.eventually.equal(100);
+      */
+    });
+  });
 });
 /* eslint-disable no-unexpected-multiline */

From 6cb98a5d9df91ad7e8daecf81f8deef4d066ae14 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Mon, 25 Aug 2025 17:01:58 +0200
Subject: [PATCH 07/52] Remove tmp freezable

---
 contracts/token/extensions/ERC7984Freezable.sol | 17 -----------------
 1 file changed, 17 deletions(-)
 delete mode 100644 contracts/token/extensions/ERC7984Freezable.sol

diff --git a/contracts/token/extensions/ERC7984Freezable.sol b/contracts/token/extensions/ERC7984Freezable.sol
deleted file mode 100644
index 47b1819f..00000000
--- a/contracts/token/extensions/ERC7984Freezable.sol
+++ /dev/null
@@ -1,17 +0,0 @@
-// SPDX-License-Identifier: MIT
-
-pragma solidity ^0.8.27;
-
-import {euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
-
-abstract contract ERC7984Freezable {
-    function confidentialFrozen(address account) public view virtual returns (euint64);
-    function confidentialAvailable(address account) public virtual returns (euint64);
-    function setConfidentialFrozen(
-        address account,
-        externalEuint64 encryptedAmount,
-        bytes calldata inputProof
-    ) public virtual;
-    function setConfidentialFrozen(address account, euint64 encryptedAmount) public virtual;
-    function _checkFreezer() internal virtual;
-}

From d2562aa38430818af9297f20a0e0605c7e5015d3 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 09:30:20 +0200
Subject: [PATCH 08/52] Name ERC7984Rwa

---
 .../interfaces/{draft-IERCXXXXCRwa.sol => IERC7984Rwa.sol}  | 4 ++--
 contracts/token/ERC7984/extensions/ERC7984Rwa.sol           | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)
 rename contracts/interfaces/{draft-IERCXXXXCRwa.sol => IERC7984Rwa.sol} (97%)

diff --git a/contracts/interfaces/draft-IERCXXXXCRwa.sol b/contracts/interfaces/IERC7984Rwa.sol
similarity index 97%
rename from contracts/interfaces/draft-IERCXXXXCRwa.sol
rename to contracts/interfaces/IERC7984Rwa.sol
index 274536f6..c0889397 100644
--- a/contracts/interfaces/draft-IERCXXXXCRwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -7,7 +7,7 @@ import {IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 
 /// @dev Interface for confidential RWA contracts.
-interface IERCXXXXCRwa is IERC7984, IERC165, IAccessControl {
+interface IERC7984Rwa is IERC7984, IERC165, IAccessControl {
     /// @dev Emitted when the contract is paused.
     event Paused(address account);
     /// @dev Emitted when the contract is unpaused.
@@ -66,7 +66,7 @@ interface IERCXXXXCRwa is IERC7984, IERC165, IAccessControl {
 }
 
 /// @dev Interface for confidential RWA compliance.
-interface IERCXXXXCRWACompliance {
+interface IERC7984RwaCompliance {
     /// @dev Checks if a transfer follows token compliance.
     function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 9389df86..5cdccca2 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -7,14 +7,14 @@ import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
 import {ERC165, IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
 import {Pausable} from "@openzeppelin/contracts/utils/Pausable.sol";
-import {IERCXXXXCRwa} from "./../../../interfaces/draft-IERCXXXXCRwa.sol";
+import {IERC7984Rwa} from "./../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984} from "./../ERC7984.sol";
 import {ERC7984Freezable} from "./ERC7984Freezable.sol";
 
 /**
  * @dev Extension of {ERC7984} supporting confidential Real World Assets.
  */
-abstract contract ERC7984Rwa is ERC7984, Pausable, ERC7984Freezable, Multicall, ERC165, AccessControl {
+abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall, ERC165, AccessControl {
     bytes32 public constant AGENT_ROLE = keccak256("AGENT_ROLE");
 
     /// @dev The caller account is not authorized to perform the operation.
@@ -46,7 +46,7 @@ abstract contract ERC7984Rwa is ERC7984, Pausable, ERC7984Freezable, Multicall,
 
     /// @inheritdoc ERC165
     function supportsInterface(bytes4 interfaceId) public view virtual override(ERC165, AccessControl) returns (bool) {
-        return interfaceId == type(IERCXXXXCRwa).interfaceId || super.supportsInterface(interfaceId);
+        return interfaceId == type(IERC7984Rwa).interfaceId || super.supportsInterface(interfaceId);
     }
 
     /// @dev Pauses contract.

From f58c1f3d05ac3b4493265ce8250d3fbf36431c20 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 09:56:37 +0200
Subject: [PATCH 09/52] Name confidential

---
 contracts/interfaces/IERC7984Rwa.sol          |  24 ++-
 .../token/ERC7984/extensions/ERC7984Rwa.sol   |  71 +++++---
 test/token/extensions/ERC7984Rwa.test.ts      | 165 ++++++------------
 3 files changed, 119 insertions(+), 141 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index c0889397..81106c8a 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -40,29 +40,37 @@ interface IERC7984Rwa is IERC7984, IERC165, IAccessControl {
     ) external;
     /// @dev Sets confidential amount of token for an account as frozen.
     function setConfidentialFrozen(address account, euint64 encryptedAmount) external;
-    /// @dev Receives and executes a batch of function calls on this contract.
-    function multicall(bytes[] calldata data) external returns (bytes[] memory results);
     /// @dev Mints confidential amount of tokens to account with proof.
-    function mint(address to, externalEuint64 encryptedAmount, bytes calldata inputProof) external returns (euint64);
+    function confidentialMint(
+        address to,
+        externalEuint64 encryptedAmount,
+        bytes calldata inputProof
+    ) external returns (euint64);
     /// @dev Mints confidential amount of tokens to account.
-    function mint(address to, euint64 encryptedAmount) external;
+    function confidentialMint(address to, euint64 encryptedAmount) external;
     /// @dev Burns confidential amount of tokens from account with proof.
-    function burn(
+    function confidentialBurn(
         address account,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
     ) external returns (euint64);
     /// @dev Burns confidential amount of tokens from account.
-    function burn(address account, euint64 encryptedAmount) external returns (euint64);
+    function confidentialBurn(address account, euint64 encryptedAmount) external returns (euint64);
     /// @dev Forces transfer of confidential amount of tokens from account to account with proof by skipping compliance checks.
-    function forceTransfer(
+    function forceConfidentialTransferFrom(
         address from,
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
     ) external returns (euint64);
     /// @dev Forces transfer of confidential amount of tokens from account to account by skipping compliance checks.
-    function forceTransfer(address from, address to, euint64 encryptedAmount) external returns (euint64);
+    function forceConfidentialTransferFrom(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) external returns (euint64);
+    /// @dev Receives and executes a batch of function calls on this contract.
+    function multicall(bytes[] calldata data) external returns (bytes[] memory results);
 }
 
 /// @dev Interface for confidential RWA compliance.
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 5cdccca2..070d0e9b 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -80,68 +80,93 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
     }
 
     /// @dev Mints confidential amount of tokens to account with proof.
-    function mint(
+    function confidentialMint(
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
     ) public virtual returns (euint64) {
-        return mint(to, FHE.fromExternal(encryptedAmount, inputProof));
+        return _confidentialMint(to, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Mints confidential amount of tokens to account.
-    function mint(address to, euint64 encryptedAmount) public virtual onlyAdminOrAgent returns (euint64) {
-        return _mint(to, encryptedAmount);
+    function confidentialMint(address to, euint64 encryptedAmount) public virtual returns (euint64) {
+        return _confidentialMint(to, encryptedAmount);
     }
 
     /// @dev Burns confidential amount of tokens from account with proof.
-    function burn(
+    function confidentialBurn(
         address account,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
     ) public virtual returns (euint64) {
-        return burn(account, FHE.fromExternal(encryptedAmount, inputProof));
+        return _confidentialBurn(account, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Burns confidential amount of tokens from account.
-    function burn(address account, euint64 encryptedAmount) public virtual onlyAdminOrAgent returns (euint64) {
-        return _burn(account, encryptedAmount);
+    function confidentialBurn(address account, euint64 encryptedAmount) public virtual returns (euint64) {
+        return _confidentialBurn(account, encryptedAmount);
     }
 
-    //TODO: Rename all to confidential
     /// @dev Forces transfer of confidential amount of tokens from account to account with proof by skipping compliance checks.
-    function forceTransfer(
+    function forceConfidentialTransferFrom(
         address from,
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
     ) public virtual returns (euint64) {
-        return forceTransfer(from, to, FHE.fromExternal(encryptedAmount, inputProof));
+        return _forceConfidentialTransferFrom(from, to, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Forces transfer of confidential amount of tokens from account to account by skipping compliance checks.
-    function forceTransfer(
+    function forceConfidentialTransferFrom(
         address from,
         address to,
         euint64 encryptedAmount
-    ) public virtual onlyAdminOrAgent returns (euint64 transferred) {
-        euint64 available = confidentialAvailable(from);
-        transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen & compliance checks
-        setConfidentialFrozen(
-            from,
-            FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from))
-        );
+    ) public virtual returns (euint64 transferred) {
+        return _forceConfidentialTransferFrom(from, to, encryptedAmount);
     }
 
-    /// @dev Adds an agent.
+    /// @dev Internal function which adds an agent.
     function _addAgent(address account) internal virtual onlyAdminOrAgent {
         _grantRole(AGENT_ROLE, account);
     }
 
-    /// @dev Removes an agent.
+    /// @dev Internal function which removes an agent.
     function _removeAgent(address account) internal virtual onlyAdminOrAgent {
         _revokeRole(AGENT_ROLE, account);
     }
 
+    /// @dev Internal function which mints confidential amount of tokens to account.
+    function _confidentialMint(
+        address to,
+        euint64 encryptedAmount
+    ) internal virtual onlyAdminOrAgent returns (euint64) {
+        return _mint(to, encryptedAmount);
+    }
+
+    /// @dev Internal function which burns confidential amount of tokens from account.
+    function _confidentialBurn(
+        address account,
+        euint64 encryptedAmount
+    ) internal virtual onlyAdminOrAgent returns (euint64) {
+        return _burn(account, encryptedAmount);
+    }
+
+    /// @dev Internal function which forces transfer of confidential amount of tokens from account to account by skipping compliance checks.
+    function _forceConfidentialTransferFrom(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) internal virtual onlyAdminOrAgent returns (euint64 transferred) {
+        euint64 available = confidentialAvailable(from);
+        transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen & compliance checks
+        setConfidentialFrozen(
+            from,
+            FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from))
+        );
+    }
+
+    /// @dev Internal function which updates confidential balances while performing frozen and compliance checks.
     function _update(
         address from,
         address to,
@@ -152,6 +177,10 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
         return super._update(from, to, encryptedAmount);
     }
 
+    /**
+     * @dev Internal function which reverts if `msg.sender` is not authorized as a freezer.
+     * This freezer role is only granted to admin or agent.
+     */
     function _checkFreezer() internal override onlyAdminOrAgent {}
 
     /// @dev Checks if a transfer follows token compliance.
diff --git a/test/token/extensions/ERC7984Rwa.test.ts b/test/token/extensions/ERC7984Rwa.test.ts
index 0f5ff399..40ac094d 100644
--- a/test/token/extensions/ERC7984Rwa.test.ts
+++ b/test/token/extensions/ERC7984Rwa.test.ts
@@ -7,10 +7,7 @@ describe('ERC7984Rwa', function () {
   async function deployFixture() {
     const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
     const token = await ethers.deployContract('ERC7984RwaMock', ['name', 'symbol', 'uri']);
-    await token
-      .connect(admin)
-      .addAgent(agent1)
-      .then(tx => tx.wait());
+    await token.connect(admin).addAgent(agent1);
     token.connect(anyone);
     return { token, admin, agent1, agent2, recipient, anyone };
   }
@@ -20,15 +17,9 @@ describe('ERC7984Rwa', function () {
       const { token, admin, agent1 } = await deployFixture();
       for (const manager of [admin, agent1]) {
         expect(await token.paused()).is.false;
-        await token
-          .connect(manager)
-          .pause()
-          .then(tx => tx.wait());
+        await token.connect(manager).pause();
         expect(await token.paused()).is.true;
-        await token
-          .connect(manager)
-          .unpause()
-          .then(tx => tx.wait());
+        await token.connect(manager).unpause();
         expect(await token.paused()).is.false;
       }
     });
@@ -59,15 +50,9 @@ describe('ERC7984Rwa', function () {
       const { token, admin, agent1, agent2 } = await deployFixture();
       for (const manager of [admin, agent1]) {
         expect(await token.isAgent(agent2)).is.false;
-        await token
-          .connect(manager)
-          .addAgent(agent2)
-          .then(tx => tx.wait());
+        await token.connect(manager).addAgent(agent2);
         expect(await token.isAgent(agent2)).is.true;
-        await token
-          .connect(manager)
-          .removeAgent(agent2)
-          .then(tx => tx.wait());
+        await token.connect(manager).removeAgent(agent2);
         expect(await token.isAgent(agent2)).is.false;
       }
     });
@@ -96,16 +81,12 @@ describe('ERC7984Rwa', function () {
           .createEncryptedInput(await token.getAddress(), manager.address)
           .add64(100)
           .encrypt();
-        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token.$_setCompliantTransfer();
         await token
           .connect(manager)
-          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof)
-          .then(tx => tx.wait());
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
         const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(balanceHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
         ).to.eventually.equal(100);
@@ -118,10 +99,11 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
         .encrypt();
+      await token.$_setCompliantTransfer();
       await expect(
         token
           .connect(anyone)
-          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
         .withArgs(anyone);
@@ -136,7 +118,7 @@ describe('ERC7984Rwa', function () {
       await expect(
         token
           .connect(admin)
-          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
         .withArgs(ethers.ZeroAddress, recipient, encryptedInput.handles[0]);
@@ -144,10 +126,7 @@ describe('ERC7984Rwa', function () {
 
     it('should not mint if paused', async function () {
       const { token, admin, recipient } = await deployFixture();
-      await token
-        .connect(admin)
-        .pause()
-        .then(tx => tx.wait());
+      await token.connect(admin).pause();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
@@ -155,7 +134,7 @@ describe('ERC7984Rwa', function () {
       await expect(
         token
           .connect(admin)
-          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       ).to.be.revertedWithCustomError(token, 'EnforcedPause');
     });
   });
@@ -169,28 +148,20 @@ describe('ERC7984Rwa', function () {
           .createEncryptedInput(await token.getAddress(), manager.address)
           .add64(100)
           .encrypt();
-        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token.$_setCompliantTransfer();
         await token
           .connect(manager)
-          ['mint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof)
-          .then(tx => tx.wait());
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
         const balanceBeforeHandle = await token.confidentialBalanceOf(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(balanceBeforeHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(balanceBeforeHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, balanceBeforeHandle, await token.getAddress(), manager),
         ).to.eventually.greaterThan(0);
         await token
           .connect(manager)
-          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof)
-          .then(tx => tx.wait());
+          ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
         const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(balanceHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
         ).to.eventually.equal(0);
@@ -203,10 +174,11 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
         .encrypt();
+      await token.$_setCompliantTransfer();
       await expect(
         token
           .connect(anyone)
-          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+          ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
         .withArgs(anyone);
@@ -221,7 +193,7 @@ describe('ERC7984Rwa', function () {
       await expect(
         token
           .connect(admin)
-          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+          ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
         .withArgs(recipient, ethers.ZeroAddress, encryptedInput.handles[0]);
@@ -229,10 +201,7 @@ describe('ERC7984Rwa', function () {
 
     it('should not burn if paused', async function () {
       const { token, admin, recipient } = await deployFixture();
-      await token
-        .connect(admin)
-        .pause()
-        .then(tx => tx.wait());
+      await token.connect(admin).pause();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
@@ -240,7 +209,7 @@ describe('ERC7984Rwa', function () {
       await expect(
         token
           .connect(admin)
-          ['burn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
+          ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       ).to.be.revertedWithCustomError(token, 'EnforcedPause');
     });
   });
@@ -254,15 +223,14 @@ describe('ERC7984Rwa', function () {
           .createEncryptedInput(await token.getAddress(), manager.address)
           .add64(100)
           .encrypt();
-        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token.$_setCompliantTransfer();
         await token
           .connect(manager)
-          ['mint(address,bytes32,bytes)'](
+          ['confidentialMint(address,bytes32,bytes)'](
             recipient,
             encryptedMintValueInput.handles[0],
             encryptedMintValueInput.inputProof,
-          )
-          .then(tx => tx.wait());
+          );
         // set frozen (50 available and about to force transfer 25)
         const encryptedFrozenValueInput = await fhevm
           .createEncryptedInput(await token.getAddress(), manager.address)
@@ -274,36 +242,28 @@ describe('ERC7984Rwa', function () {
             recipient,
             encryptedFrozenValueInput.handles[0],
             encryptedFrozenValueInput.inputProof,
-          )
-          .then(tx => tx.wait());
+          );
         const encryptedTransferValueInput = await fhevm
           .createEncryptedInput(await token.getAddress(), manager.address)
           .add64(25)
           .encrypt();
-        await token.$_unsetCompliantTransfer().then(tx => tx.wait());
+        await token.$_unsetCompliantTransfer();
         expect(await token.compliantTransfer()).to.be.false;
         await token
           .connect(manager)
-          ['forceTransfer(address,address,bytes32,bytes)'](
+          ['forceConfidentialTransferFrom(address,address,bytes32,bytes)'](
             recipient,
             anyone,
             encryptedTransferValueInput.handles[0],
             encryptedTransferValueInput.inputProof,
-          )
-          .then(tx => tx.wait());
+          );
         const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(balanceHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
         ).to.eventually.equal(75);
         const frozenHandle = await token.confidentialFrozen(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(frozenHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
         ).to.eventually.equal(50); // frozen is left unchanged
@@ -318,15 +278,14 @@ describe('ERC7984Rwa', function () {
           .createEncryptedInput(await token.getAddress(), manager.address)
           .add64(100)
           .encrypt();
-        await token.$_setCompliantTransfer().then(tx => tx.wait());
+        await token.$_setCompliantTransfer();
         await token
           .connect(manager)
-          ['mint(address,bytes32,bytes)'](
+          ['confidentialMint(address,bytes32,bytes)'](
             recipient,
             encryptedMintValueInput.handles[0],
             encryptedMintValueInput.inputProof,
-          )
-          .then(tx => tx.wait());
+          );
         // set frozen (only 20 available but about to force transfer 25)
         const encryptedFrozenValueInput = await fhevm
           .createEncryptedInput(await token.getAddress(), manager.address)
@@ -338,42 +297,31 @@ describe('ERC7984Rwa', function () {
             recipient,
             encryptedFrozenValueInput.handles[0],
             encryptedFrozenValueInput.inputProof,
-          )
-          .then(tx => tx.wait());
+          );
         const encryptedTransferValueInput = await fhevm
           .createEncryptedInput(await token.getAddress(), manager.address)
           .add64(25)
           .encrypt();
-        await token.$_unsetCompliantTransfer().then(tx => tx.wait());
+        await token.$_unsetCompliantTransfer();
         expect(await token.compliantTransfer()).to.be.false;
         // should force transfer even if paused
-        await token
-          .connect(manager)
-          .pause()
-          .then(tx => tx.wait());
+        await token.connect(manager).pause();
         expect(await token.paused()).to.be.true;
         await token
           .connect(manager)
-          ['forceTransfer(address,address,bytes32,bytes)'](
+          ['forceConfidentialTransferFrom(address,address,bytes32,bytes)'](
             recipient,
             anyone,
             encryptedTransferValueInput.handles[0],
             encryptedTransferValueInput.inputProof,
-          )
-          .then(tx => tx.wait());
+          );
         const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(balanceHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
         ).to.eventually.equal(75);
         const frozenHandle = await token.confidentialFrozen(recipient);
-        await token
-          .connect(manager)
-          .getHandleAllowance(frozenHandle, manager, true)
-          .then(tx => tx.wait());
+        await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
         await expect(
           fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
         ).to.eventually.equal(75); // frozen got reset to available balance
@@ -388,15 +336,14 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      await token.$_setCompliantTransfer();
       await token
         .connect(manager)
-        ['mint(address,bytes32,bytes)'](
+        ['confidentialMint(address,bytes32,bytes)'](
           recipient,
           encryptedMintValueInput.handles[0],
           encryptedMintValueInput.inputProof,
-        )
-        .then(tx => tx.wait());
+        );
       // set frozen (50 available and about to transfer 25)
       const encryptedFrozenValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), manager.address)
@@ -408,13 +355,12 @@ describe('ERC7984Rwa', function () {
           recipient,
           encryptedFrozenValueInput.handles[0],
           encryptedFrozenValueInput.inputProof,
-        )
-        .then(tx => tx.wait());
+        );
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      await token.$_setCompliantTransfer();
       expect(await token.compliantTransfer()).to.be.true;
       await expect(
         token
@@ -441,10 +387,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      await token
-        .connect(manager)
-        .pause()
-        .then(tx => tx.wait());
+      await token.connect(manager).pause();
       await expect(
         token
           .connect(recipient)
@@ -482,15 +425,14 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      await token.$_setCompliantTransfer();
       await token
         .connect(manager)
-        ['mint(address,bytes32,bytes)'](
+        ['confidentialMint(address,bytes32,bytes)'](
           recipient,
           encryptedMintValueInput.handles[0],
           encryptedMintValueInput.inputProof,
-        )
-        .then(tx => tx.wait());
+        );
       // set frozen (20 available but about to transfer 25)
       const encryptedFrozenValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), manager.address)
@@ -502,13 +444,12 @@ describe('ERC7984Rwa', function () {
           recipient,
           encryptedFrozenValueInput.handles[0],
           encryptedFrozenValueInput.inputProof,
-        )
-        .then(tx => tx.wait());
+        );
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      await token.$_setCompliantTransfer().then(tx => tx.wait());
+      await token.$_setCompliantTransfer();
       expect(await token.compliantTransfer()).to.be.true;
       await expect(
         token

From 76b21ae0624dc97b63ce7e4b07ec9c9c07106de2 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 10:15:42 +0200
Subject: [PATCH 10/52] Move RWA test

---
 test/token/{ => ERC7984}/extensions/ERC7984Rwa.test.ts | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename test/token/{ => ERC7984}/extensions/ERC7984Rwa.test.ts (100%)

diff --git a/test/token/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
similarity index 100%
rename from test/token/extensions/ERC7984Rwa.test.ts
rename to test/token/ERC7984/extensions/ERC7984Rwa.test.ts

From 127aff5da6b78dab7543e1f8f4b730c9d07dfa04 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 10:52:41 +0200
Subject: [PATCH 11/52] Test with & without proof

---
 contracts/mocks/token/ERC7984RwaMock.sol      |   5 +
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 252 +++++++++++-------
 2 files changed, 157 insertions(+), 100 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index d55d8397..cb093b4a 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -16,6 +16,11 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
 
     constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
 
+    function createEncryptedAmount(uint64 amount) public returns (euint64 encryptedAmount) {
+        FHE.allowThis(encryptedAmount = FHE.asEuint64(amount));
+        FHE.allow(encryptedAmount, msg.sender);
+    }
+
     function $_setCompliantTransfer() public {
         compliantTransfer = true;
     }
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 40ac094d..e84151ae 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -1,5 +1,6 @@
 import { FhevmType } from '@fhevm/hardhat-plugin';
 import { expect } from 'chai';
+import { AddressLike, BytesLike } from 'ethers';
 import { ethers, fhevm } from 'hardhat';
 
 /* eslint-disable no-unexpected-multiline */
@@ -73,25 +74,39 @@ describe('ERC7984Rwa', function () {
   });
 
   describe('Mintable', async function () {
-    it('should mint by admin or agent', async function () {
-      const { admin, agent1, recipient } = await deployFixture();
-      for (const manager of [admin, agent1]) {
-        const { token } = await deployFixture();
-        const encryptedInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(100)
-          .encrypt();
-        await token.$_setCompliantTransfer();
-        await token
-          .connect(manager)
-          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
-        const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-        ).to.eventually.equal(100);
-      }
-    });
+    for (const withProof of [true, false]) {
+      it(`should mint by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
+        const { admin, agent1, recipient } = await deployFixture();
+        for (const manager of [admin, agent1]) {
+          const { token } = await deployFixture();
+          await token.$_setCompliantTransfer();
+          const amount = 100;
+          let params = [recipient.address] as unknown as [
+            account: AddressLike,
+            encryptedAmount: BytesLike,
+            inputProof: BytesLike,
+          ];
+          if (withProof) {
+            const { handles, inputProof } = await fhevm
+              .createEncryptedInput(await token.getAddress(), manager.address)
+              .add64(amount)
+              .encrypt();
+            params.push(handles[0], inputProof);
+          } else {
+            await token.connect(manager).createEncryptedAmount(amount);
+            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
+          }
+          await token
+            .connect(manager)
+            [withProof ? 'confidentialMint(address,bytes32,bytes)' : 'confidentialMint(address,bytes32)'](...params);
+          const balanceHandle = await token.confidentialBalanceOf(recipient);
+          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+          ).to.eventually.equal(100);
+        }
+      });
+    }
 
     it('should not mint if neither admin nor agent', async function () {
       const { token, recipient, anyone } = await deployFixture();
@@ -140,33 +155,55 @@ describe('ERC7984Rwa', function () {
   });
 
   describe('Burnable', async function () {
-    it('should burn by admin or agent', async function () {
-      const { admin, agent1, recipient } = await deployFixture();
-      for (const manager of [admin, agent1]) {
-        const { token } = await deployFixture();
-        const encryptedInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(100)
-          .encrypt();
-        await token.$_setCompliantTransfer();
-        await token
-          .connect(manager)
-          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
-        const balanceBeforeHandle = await token.confidentialBalanceOf(recipient);
-        await token.connect(manager).getHandleAllowance(balanceBeforeHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, balanceBeforeHandle, await token.getAddress(), manager),
-        ).to.eventually.greaterThan(0);
-        await token
-          .connect(manager)
-          ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
-        const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-        ).to.eventually.equal(0);
-      }
-    });
+    for (const withProof of [true, false]) {
+      it(`should burn by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
+        const { admin, agent1, recipient } = await deployFixture();
+        for (const manager of [admin, agent1]) {
+          const { token } = await deployFixture();
+          const encryptedInput = await fhevm
+            .createEncryptedInput(await token.getAddress(), manager.address)
+            .add64(100)
+            .encrypt();
+          await token.$_setCompliantTransfer();
+          await token
+            .connect(manager)
+            ['confidentialMint(address,bytes32,bytes)'](
+              recipient,
+              encryptedInput.handles[0],
+              encryptedInput.inputProof,
+            );
+          const balanceBeforeHandle = await token.confidentialBalanceOf(recipient);
+          await token.connect(manager).getHandleAllowance(balanceBeforeHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, balanceBeforeHandle, await token.getAddress(), manager),
+          ).to.eventually.greaterThan(0);
+          const amount = 100;
+          let params = [recipient.address] as unknown as [
+            account: AddressLike,
+            encryptedAmount: BytesLike,
+            inputProof: BytesLike,
+          ];
+          if (withProof) {
+            const { handles, inputProof } = await fhevm
+              .createEncryptedInput(await token.getAddress(), manager.address)
+              .add64(amount)
+              .encrypt();
+            params.push(handles[0], inputProof);
+          } else {
+            await token.connect(manager).createEncryptedAmount(amount);
+            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
+          }
+          await token
+            .connect(manager)
+            [withProof ? 'confidentialBurn(address,bytes32,bytes)' : 'confidentialBurn(address,bytes32)'](...params);
+          const balanceHandle = await token.confidentialBalanceOf(recipient);
+          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+          ).to.eventually.equal(0);
+        }
+      });
+    }
 
     it('should not burn if neither admin nor agent', async function () {
       const { token, recipient, anyone } = await deployFixture();
@@ -215,60 +252,75 @@ describe('ERC7984Rwa', function () {
   });
 
   describe('Force transfer', async function () {
-    it('should force transfer by admin or agent', async function () {
-      const { admin, agent1, recipient, anyone } = await deployFixture();
-      for (const manager of [admin, agent1]) {
-        const { token } = await deployFixture();
-        const encryptedMintValueInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(100)
-          .encrypt();
-        await token.$_setCompliantTransfer();
-        await token
-          .connect(manager)
-          ['confidentialMint(address,bytes32,bytes)'](
-            recipient,
-            encryptedMintValueInput.handles[0],
-            encryptedMintValueInput.inputProof,
-          );
-        // set frozen (50 available and about to force transfer 25)
-        const encryptedFrozenValueInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(50)
-          .encrypt();
-        await token
-          .connect(manager)
-          ['setConfidentialFrozen(address,bytes32,bytes)'](
-            recipient,
-            encryptedFrozenValueInput.handles[0],
-            encryptedFrozenValueInput.inputProof,
-          );
-        const encryptedTransferValueInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(25)
-          .encrypt();
-        await token.$_unsetCompliantTransfer();
-        expect(await token.compliantTransfer()).to.be.false;
-        await token
-          .connect(manager)
-          ['forceConfidentialTransferFrom(address,address,bytes32,bytes)'](
-            recipient,
-            anyone,
-            encryptedTransferValueInput.handles[0],
-            encryptedTransferValueInput.inputProof,
-          );
-        const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-        ).to.eventually.equal(75);
-        const frozenHandle = await token.confidentialFrozen(recipient);
-        await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
-        ).to.eventually.equal(50); // frozen is left unchanged
-      }
-    });
+    for (const withProof of [true, false]) {
+      it(`should force transfer by admin or agent
+         ${withProof ? 'with proof' : ''}`, async function () {
+        const { admin, agent1, recipient, anyone } = await deployFixture();
+        for (const manager of [admin, agent1]) {
+          const { token } = await deployFixture();
+          const encryptedMintValueInput = await fhevm
+            .createEncryptedInput(await token.getAddress(), manager.address)
+            .add64(100)
+            .encrypt();
+          await token.$_setCompliantTransfer();
+          await token
+            .connect(manager)
+            ['confidentialMint(address,bytes32,bytes)'](
+              recipient,
+              encryptedMintValueInput.handles[0],
+              encryptedMintValueInput.inputProof,
+            );
+          // set frozen (50 available and about to force transfer 25)
+          const encryptedFrozenValueInput = await fhevm
+            .createEncryptedInput(await token.getAddress(), manager.address)
+            .add64(50)
+            .encrypt();
+          await token
+            .connect(manager)
+            ['setConfidentialFrozen(address,bytes32,bytes)'](
+              recipient,
+              encryptedFrozenValueInput.handles[0],
+              encryptedFrozenValueInput.inputProof,
+            );
+          await token.$_unsetCompliantTransfer();
+          expect(await token.compliantTransfer()).to.be.false;
+          const amount = 25;
+          let params = [recipient.address, anyone.address] as unknown as [
+            from: AddressLike,
+            to: AddressLike,
+            encryptedAmount: BytesLike,
+            inputProof: BytesLike,
+          ];
+          if (withProof) {
+            const { handles, inputProof } = await fhevm
+              .createEncryptedInput(await token.getAddress(), manager.address)
+              .add64(amount)
+              .encrypt();
+            params.push(handles[0], inputProof);
+          } else {
+            await token.connect(manager).createEncryptedAmount(amount);
+            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
+          }
+          await token
+            .connect(manager)
+            [
+              withProof
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'forceConfidentialTransferFrom(address,address,bytes32)'
+            ](...params);
+          const balanceHandle = await token.confidentialBalanceOf(recipient);
+          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+          ).to.eventually.equal(75);
+          const frozenHandle = await token.confidentialFrozen(recipient);
+          await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
+          ).to.eventually.equal(50); // frozen is left unchanged
+        }
+      });
+    }
 
     it('should force transfer even if frozen', async function () {
       const { admin, agent1, recipient, anyone } = await deployFixture();

From 84af6873252c7a8d9e1f555d8e3974317b1fc5d4 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 12:03:13 +0200
Subject: [PATCH 12/52] Rwa mock uses freezable

---
 contracts/mocks/token/ERC7984RwaMock.sol      |  24 ----
 .../token/ERC7984/extensions/ERC7984Rwa.sol   |   2 +-
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 133 ++++++++++--------
 3 files changed, 74 insertions(+), 85 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index cb093b4a..97533af5 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -41,29 +41,5 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return compliantTransfer;
     }
 
-    // TODO: Remove all below
-    function confidentialAvailable(address account) public override returns (euint64) {
-        (ebool success, euint64 unfrozen) = FHESafeMath.tryDecrease(
-            confidentialBalanceOf(account),
-            confidentialFrozen(account)
-        );
-        unfrozen = FHE.select(success, unfrozen, FHE.asEuint64(0));
-        FHE.allowThis(unfrozen);
-        return unfrozen;
-    }
-    function confidentialFrozen(address account) public view override returns (euint64) {
-        return _frozenBalances[account];
-    }
-    function setConfidentialFrozen(
-        address account,
-        externalEuint64 encryptedAmount,
-        bytes calldata inputProof
-    ) public override {
-        return setConfidentialFrozen(account, FHE.fromExternal(encryptedAmount, inputProof));
-    }
-    function setConfidentialFrozen(address account, euint64 encryptedAmount) public override {
-        FHE.allowThis(_frozenBalances[account] = encryptedAmount);
-    }
-
     function _validateHandleAllowance(bytes32 handle) internal view override onlyAdminOrAgent {}
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 070d0e9b..5ea2e855 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -160,7 +160,7 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
     ) internal virtual onlyAdminOrAgent returns (euint64 transferred) {
         euint64 available = confidentialAvailable(from);
         transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen & compliance checks
-        setConfidentialFrozen(
+        _setConfidentialFrozen(
             from,
             FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from))
         );
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index e84151ae..24886bf5 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -253,8 +253,7 @@ describe('ERC7984Rwa', function () {
 
   describe('Force transfer', async function () {
     for (const withProof of [true, false]) {
-      it(`should force transfer by admin or agent
-         ${withProof ? 'with proof' : ''}`, async function () {
+      it(`should force transfer by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
         const { admin, agent1, recipient, anyone } = await deployFixture();
         for (const manager of [admin, agent1]) {
           const { token } = await deployFixture();
@@ -321,64 +320,78 @@ describe('ERC7984Rwa', function () {
         }
       });
     }
-
-    it('should force transfer even if frozen', async function () {
-      const { admin, agent1, recipient, anyone } = await deployFixture();
-      for (const manager of [admin, agent1]) {
-        const { token } = await deployFixture();
-        const encryptedMintValueInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(100)
-          .encrypt();
-        await token.$_setCompliantTransfer();
-        await token
-          .connect(manager)
-          ['confidentialMint(address,bytes32,bytes)'](
-            recipient,
-            encryptedMintValueInput.handles[0],
-            encryptedMintValueInput.inputProof,
-          );
-        // set frozen (only 20 available but about to force transfer 25)
-        const encryptedFrozenValueInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(80)
-          .encrypt();
-        await token
-          .connect(manager)
-          ['setConfidentialFrozen(address,bytes32,bytes)'](
-            recipient,
-            encryptedFrozenValueInput.handles[0],
-            encryptedFrozenValueInput.inputProof,
-          );
-        const encryptedTransferValueInput = await fhevm
-          .createEncryptedInput(await token.getAddress(), manager.address)
-          .add64(25)
-          .encrypt();
-        await token.$_unsetCompliantTransfer();
-        expect(await token.compliantTransfer()).to.be.false;
-        // should force transfer even if paused
-        await token.connect(manager).pause();
-        expect(await token.paused()).to.be.true;
-        await token
-          .connect(manager)
-          ['forceConfidentialTransferFrom(address,address,bytes32,bytes)'](
-            recipient,
-            anyone,
-            encryptedTransferValueInput.handles[0],
-            encryptedTransferValueInput.inputProof,
-          );
-        const balanceHandle = await token.confidentialBalanceOf(recipient);
-        await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-        ).to.eventually.equal(75);
-        const frozenHandle = await token.confidentialFrozen(recipient);
-        await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
-        await expect(
-          fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
-        ).to.eventually.equal(75); // frozen got reset to available balance
-      }
-    });
+    for (const withProof of [true, false]) {
+      it(`should force transfer even if frozen ${withProof ? 'with proof' : ''}`, async function () {
+        const { admin, agent1, recipient, anyone } = await deployFixture();
+        for (const manager of [admin, agent1]) {
+          const { token } = await deployFixture();
+          const encryptedMintValueInput = await fhevm
+            .createEncryptedInput(await token.getAddress(), manager.address)
+            .add64(100)
+            .encrypt();
+          await token.$_setCompliantTransfer();
+          await token
+            .connect(manager)
+            ['confidentialMint(address,bytes32,bytes)'](
+              recipient,
+              encryptedMintValueInput.handles[0],
+              encryptedMintValueInput.inputProof,
+            );
+          // set frozen (only 20 available but about to force transfer 25)
+          const encryptedFrozenValueInput = await fhevm
+            .createEncryptedInput(await token.getAddress(), manager.address)
+            .add64(80)
+            .encrypt();
+          await token
+            .connect(manager)
+            ['setConfidentialFrozen(address,bytes32,bytes)'](
+              recipient,
+              encryptedFrozenValueInput.handles[0],
+              encryptedFrozenValueInput.inputProof,
+            );
+          // should force transfer even if not compliant
+          await token.$_unsetCompliantTransfer();
+          expect(await token.compliantTransfer()).to.be.false;
+          // should force transfer even if paused
+          await token.connect(manager).pause();
+          expect(await token.paused()).to.be.true;
+          const amount = 25;
+          let params = [recipient.address, anyone.address] as unknown as [
+            from: AddressLike,
+            to: AddressLike,
+            encryptedAmount: BytesLike,
+            inputProof: BytesLike,
+          ];
+          if (withProof) {
+            const { handles, inputProof } = await fhevm
+              .createEncryptedInput(await token.getAddress(), manager.address)
+              .add64(amount)
+              .encrypt();
+            params.push(handles[0], inputProof);
+          } else {
+            await token.connect(manager).createEncryptedAmount(amount);
+            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
+          }
+          await token
+            .connect(manager)
+            [
+              withProof
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'forceConfidentialTransferFrom(address,address,bytes32)'
+            ](...params);
+          const balanceHandle = await token.confidentialBalanceOf(recipient);
+          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
+          ).to.eventually.equal(75);
+          const frozenHandle = await token.confidentialFrozen(recipient);
+          await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
+          ).to.eventually.equal(75); // frozen got reset to available balance
+        }
+      });
+    }
   });
 
   describe('Transfer', async function () {

From b0d5ffa5947ffa0623c413afbc3db5c12f72aff0 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 12:44:34 +0200
Subject: [PATCH 13/52] Check transferred amounts in tests

---
 test/helpers/event.ts                         |   9 ++
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 104 +++++++++++++-----
 2 files changed, 85 insertions(+), 28 deletions(-)
 create mode 100644 test/helpers/event.ts

diff --git a/test/helpers/event.ts b/test/helpers/event.ts
new file mode 100644
index 00000000..99e13ed4
--- /dev/null
+++ b/test/helpers/event.ts
@@ -0,0 +1,9 @@
+import { EventLog } from 'ethers';
+import { ContractTransactionResponse } from 'ethers';
+import { ethers } from 'ethers';
+
+export async function callAndGetResult(txPromise: Promise<ContractTransactionResponse>, eventName: string) {
+  const receipt = await txPromise.then(tx => tx.wait());
+  const logs = receipt?.logs.filter(log => log.address == receipt.to && log.topics[0] == ethers.id(eventName));
+  return (logs![0] as EventLog).args;
+}
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 24886bf5..7feb30bc 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -1,8 +1,12 @@
+import { callAndGetResult } from '../../../helpers/event';
 import { FhevmType } from '@fhevm/hardhat-plugin';
 import { expect } from 'chai';
 import { AddressLike, BytesLike } from 'ethers';
 import { ethers, fhevm } from 'hardhat';
 
+const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
+const frozenEventSignature = 'TokensFrozen(address,bytes32)';
+
 /* eslint-disable no-unexpected-multiline */
 describe('ERC7984Rwa', function () {
   async function deployFixture() {
@@ -96,14 +100,20 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
-          await token
-            .connect(manager)
-            [withProof ? 'confidentialMint(address,bytes32,bytes)' : 'confidentialMint(address,bytes32)'](...params);
+          const [, , transferredHandle] = await callAndGetResult(
+            token
+              .connect(manager)
+              [withProof ? 'confidentialMint(address,bytes32,bytes)' : 'confidentialMint(address,bytes32)'](...params),
+            transferEventSignature,
+          );
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+          ).to.eventually.equal(amount);
           const balanceHandle = await token.confidentialBalanceOf(recipient);
           await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
           await expect(
             fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(100);
+          ).to.eventually.equal(amount);
         }
       });
     }
@@ -193,9 +203,15 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
-          await token
-            .connect(manager)
-            [withProof ? 'confidentialBurn(address,bytes32,bytes)' : 'confidentialBurn(address,bytes32)'](...params);
+          const [, , transferredHandle] = await callAndGetResult(
+            token
+              .connect(manager)
+              [withProof ? 'confidentialBurn(address,bytes32,bytes)' : 'confidentialBurn(address,bytes32)'](...params),
+            transferEventSignature,
+          );
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+          ).to.eventually.equal(amount);
           const balanceHandle = await token.confidentialBalanceOf(recipient);
           await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
           await expect(
@@ -300,13 +316,21 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
-          await token
-            .connect(manager)
-            [
-              withProof
-                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
-                : 'forceConfidentialTransferFrom(address,address,bytes32)'
-            ](...params);
+          const [from, to, transferredHandle] = await callAndGetResult(
+            token
+              .connect(manager)
+              [
+                withProof
+                  ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                  : 'forceConfidentialTransferFrom(address,address,bytes32)'
+              ](...params),
+            transferEventSignature,
+          );
+          expect(from).equal(recipient.address);
+          expect(to).equal(anyone.address);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), anyone),
+          ).to.eventually.equal(amount);
           const balanceHandle = await token.confidentialBalanceOf(recipient);
           await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
           await expect(
@@ -372,13 +396,20 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
-          await token
-            .connect(manager)
-            [
-              withProof
-                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
-                : 'forceConfidentialTransferFrom(address,address,bytes32)'
-            ](...params);
+          const [account, frozenAmountHandle] = await callAndGetResult(
+            token
+              .connect(manager)
+              [
+                withProof
+                  ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                  : 'forceConfidentialTransferFrom(address,address,bytes32)'
+              ](...params),
+            frozenEventSignature,
+          );
+          expect(account).equal(recipient.address);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, frozenAmountHandle, await token.getAddress(), recipient),
+          ).to.eventually.equal(75);
           const balanceHandle = await token.confidentialBalanceOf(recipient);
           await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
           await expect(
@@ -421,13 +452,14 @@ describe('ERC7984Rwa', function () {
           encryptedFrozenValueInput.handles[0],
           encryptedFrozenValueInput.inputProof,
         );
+      const amount = 25;
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
-        .add64(25)
+        .add64(amount)
         .encrypt();
       await token.$_setCompliantTransfer();
       expect(await token.compliantTransfer()).to.be.true;
-      await expect(
+      const [from, to, transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)
           ['confidentialTransfer(address,bytes32,bytes)'](
@@ -435,7 +467,21 @@ describe('ERC7984Rwa', function () {
             encryptedTransferValueInput.handles[0],
             encryptedTransferValueInput.inputProof,
           ),
-      ).to.emit(token, 'ConfidentialTransfer');
+        transferEventSignature,
+      );
+      expect(from).equal(recipient.address);
+      expect(to).equal(anyone.address);
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), anyone),
+      ).to.eventually.equal(amount);
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await token.confidentialBalanceOf(anyone),
+          await token.getAddress(),
+          anyone,
+        ),
+      ).to.eventually.equal(amount);
       await expect(
         fhevm.userDecryptEuint(
           FhevmType.euint64,
@@ -516,7 +562,7 @@ describe('ERC7984Rwa', function () {
         .encrypt();
       await token.$_setCompliantTransfer();
       expect(await token.compliantTransfer()).to.be.true;
-      await expect(
+      const [, , transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)
           ['confidentialTransfer(address,bytes32,bytes)'](
@@ -524,8 +570,11 @@ describe('ERC7984Rwa', function () {
             encryptedTransferValueInput.handles[0],
             encryptedTransferValueInput.inputProof,
           ),
-      ).to.emit(token, 'ConfidentialTransfer');
-      /* TODO: Enable when freezable ready
+        transferEventSignature,
+      );
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), anyone),
+      ).to.eventually.equal(0);
       // Balance is unchanged
       await expect(
         fhevm.userDecryptEuint(
@@ -535,7 +584,6 @@ describe('ERC7984Rwa', function () {
           recipient,
         ),
       ).to.eventually.equal(100);
-      */
     });
   });
 });

From 6fb7f97dc4f60443bff7206d3d1623486e5b1879 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 14:50:32 +0200
Subject: [PATCH 14/52] Bypass hardhat fhevm behaviour

---
 test/utils/HandleAccessManager.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/utils/HandleAccessManager.test.ts b/test/utils/HandleAccessManager.test.ts
index 086065a3..db896f56 100644
--- a/test/utils/HandleAccessManager.test.ts
+++ b/test/utils/HandleAccessManager.test.ts
@@ -12,7 +12,7 @@ describe('HandleAccessManager', function () {
   });
 
   it('should not be allowed to reencrypt unallowed handle', async function () {
-    const handle = await createHandle(this.mock, 100);
+    const handle = await createHandle(this.mock, 101);
 
     await expect(fhevm.userDecryptEuint(FhevmType.euint64, handle, this.mock.target, this.holder)).to.be.rejectedWith(
       `User ${this.holder.address} is not authorized to user decrypt handle ${handle}`,

From 0cb02084b5663c7219fd2a1274f4ae9db53b0970 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 17:37:50 +0200
Subject: [PATCH 15/52] Add support interface test

---
 contracts/interfaces/IERC7984Rwa.sol          |  7 ++++--
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 10 +++++---
 test/helpers/interface.ts                     | 15 ++++++++++++
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 23 +++++++++++++++++++
 4 files changed, 50 insertions(+), 5 deletions(-)
 create mode 100644 test/helpers/interface.ts

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 81106c8a..c5ea9e48 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -6,8 +6,8 @@ import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol"
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 
-/// @dev Interface for confidential RWA contracts.
-interface IERC7984Rwa is IERC7984, IERC165, IAccessControl {
+/// @dev Base interface for confidential RWA contracts.
+interface IERC7984RwaBase {
     /// @dev Emitted when the contract is paused.
     event Paused(address account);
     /// @dev Emitted when the contract is unpaused.
@@ -73,6 +73,9 @@ interface IERC7984Rwa is IERC7984, IERC165, IAccessControl {
     function multicall(bytes[] calldata data) external returns (bytes[] memory results);
 }
 
+/// @dev Full interface for confidential RWA contracts.
+interface IERC7984Rwa is IERC7984, IERC7984RwaBase, IERC165, IAccessControl {}
+
 /// @dev Interface for confidential RWA compliance.
 interface IERC7984RwaCompliance {
     /// @dev Checks if a transfer follows token compliance.
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 5ea2e855..b8122e49 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -4,10 +4,11 @@ pragma solidity ^0.8.27;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
-import {ERC165, IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import {ERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
 import {Pausable} from "@openzeppelin/contracts/utils/Pausable.sol";
-import {IERC7984Rwa} from "./../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984} from "./../../../interfaces/IERC7984.sol";
+import {IERC7984RwaBase} from "./../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984} from "./../ERC7984.sol";
 import {ERC7984Freezable} from "./ERC7984Freezable.sol";
 
@@ -46,7 +47,10 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
 
     /// @inheritdoc ERC165
     function supportsInterface(bytes4 interfaceId) public view virtual override(ERC165, AccessControl) returns (bool) {
-        return interfaceId == type(IERC7984Rwa).interfaceId || super.supportsInterface(interfaceId);
+        return
+            interfaceId == type(IERC7984RwaBase).interfaceId ||
+            interfaceId == type(IERC7984).interfaceId ||
+            super.supportsInterface(interfaceId);
     }
 
     /// @dev Pauses contract.
diff --git a/test/helpers/interface.ts b/test/helpers/interface.ts
new file mode 100644
index 00000000..4b4bf675
--- /dev/null
+++ b/test/helpers/interface.ts
@@ -0,0 +1,15 @@
+import { Interface } from 'ethers';
+import { ethers } from 'hardhat';
+
+export function getFunctions(interfaceFactory: any) {
+  return (interfaceFactory.createInterface() as Interface).fragments
+    .filter(f => f.type == 'function')
+    .map(f => f.format());
+}
+
+export function getInterfaceId(signatures: string[]) {
+  return ethers.toBeHex(
+    signatures.reduce((acc, signature) => acc ^ ethers.toBigInt(ethers.FunctionFragment.from(signature).selector), 0n),
+    4,
+  );
+}
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 7feb30bc..6d39a988 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -1,4 +1,11 @@
+import {
+  IAccessControl__factory,
+  IERC165__factory,
+  IERC7984__factory,
+  IERC7984RwaBase__factory,
+} from '../../../../types';
 import { callAndGetResult } from '../../../helpers/event';
+import { getFunctions, getInterfaceId } from '../../../helpers/interface';
 import { FhevmType } from '@fhevm/hardhat-plugin';
 import { expect } from 'chai';
 import { AddressLike, BytesLike } from 'ethers';
@@ -17,6 +24,22 @@ describe('ERC7984Rwa', function () {
     return { token, admin, agent1, agent2, recipient, anyone };
   }
 
+  describe('ERC165', async function () {
+    it('should support interfaces', async function () {
+      const { token } = await deployFixture();
+      const interfaceFactories = [
+        IERC7984RwaBase__factory,
+        IERC7984__factory,
+        IERC165__factory,
+        IAccessControl__factory,
+      ];
+      for (const interfaceFactory of interfaceFactories) {
+        const functions = getFunctions(interfaceFactory);
+        expect(await token.supportsInterface(getInterfaceId(functions))).is.true;
+      }
+    });
+  });
+
   describe('Pausable', async function () {
     it('should pause & unpause', async function () {
       const { token, admin, agent1 } = await deployFixture();

From 90ebfa09654acfee38772da306c5876d27f17470 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 17:49:58 +0200
Subject: [PATCH 16/52] Add should not force transfer if anyone

---
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 6d39a988..fc5aff6e 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -367,6 +367,7 @@ describe('ERC7984Rwa', function () {
         }
       });
     }
+
     for (const withProof of [true, false]) {
       it(`should force transfer even if frozen ${withProof ? 'with proof' : ''}`, async function () {
         const { admin, agent1, recipient, anyone } = await deployFixture();
@@ -446,6 +447,40 @@ describe('ERC7984Rwa', function () {
         }
       });
     }
+
+    for (const withProof of [true, false]) {
+      it(`should not force transfer if neither admin nor agent ${withProof ? 'with proof' : ''}`, async function () {
+        const { token, recipient, anyone } = await deployFixture();
+        let params = [recipient.address, anyone.address] as unknown as [
+          from: AddressLike,
+          to: AddressLike,
+          encryptedAmount: BytesLike,
+          inputProof: BytesLike,
+        ];
+        const amount = 100;
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await token.getAddress(), anyone.address)
+            .add64(amount)
+            .encrypt();
+          params.push(handles[0], inputProof);
+        } else {
+          await token.connect(anyone).createEncryptedAmount(amount);
+          params.push(await token.connect(anyone).createEncryptedAmount.staticCall(amount));
+        }
+        await expect(
+          token
+            .connect(anyone)
+            [
+              withProof
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'forceConfidentialTransferFrom(address,address,bytes32)'
+            ](...params),
+        )
+          .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+          .withArgs(anyone.address);
+      });
+    }
   });
 
   describe('Transfer', async function () {

From c5d07fe9fa22860338d2c07ef045caf3eb81d5d1 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 18:03:03 +0200
Subject: [PATCH 17/52] Move some modifiers to mock

---
 contracts/mocks/token/ERC7984RwaMock.sol      | 12 +++++++++++
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 20 ++++---------------
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index 97533af5..0ae0200e 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -14,6 +14,18 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
     mapping(address account => euint64 encryptedAmount) private _frozenBalances;
     bool public compliantTransfer;
 
+    // TODO: Move modifiers to `ERC7984Rwa` or remove from mock if useless
+    /// @dev Checks if the sender is an admin.
+    modifier onlyAdmin() {
+        require(isAdmin(_msgSender()), UnauthorizedSender(_msgSender()));
+        _;
+    }
+    /// @dev Checks if the sender is an agent.
+    modifier onlyAgent() {
+        require(isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
+        _;
+    }
+
     constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
 
     function createEncryptedAmount(uint64 amount) public returns (euint64 encryptedAmount) {
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index b8122e49..7f420904 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -23,28 +23,16 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
     /// @dev The transfer does not follow token compliance.
     error UncompliantTransfer(address from, address to, euint64 encryptedAmount);
 
-    constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984(name, symbol, tokenUri) {
-        _grantRole(DEFAULT_ADMIN_ROLE, _msgSender());
-    }
-
-    /// @dev Checks if the sender is an admin.
-    modifier onlyAdmin() {
-        require(isAdmin(_msgSender()), UnauthorizedSender(_msgSender()));
-        _;
-    }
-
-    /// @dev Checks if the sender is an agent.
-    modifier onlyAgent() {
-        require(isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
-        _;
-    }
-
     /// @dev Checks if the sender is an admin or an agent.
     modifier onlyAdminOrAgent() {
         require(isAdmin(_msgSender()) || isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
         _;
     }
 
+    constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984(name, symbol, tokenUri) {
+        _grantRole(DEFAULT_ADMIN_ROLE, _msgSender());
+    }
+
     /// @inheritdoc ERC165
     function supportsInterface(bytes4 interfaceId) public view virtual override(ERC165, AccessControl) returns (bool) {
         return

From e484066f5d7704bcf2b45da6463d708429168fd3 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 26 Aug 2025 18:32:23 +0200
Subject: [PATCH 18/52] Update doc

---
 .changeset/new-crews-boil.md                      | 2 +-
 contracts/token/ERC7984/extensions/ERC7984Rwa.sol | 3 ++-
 contracts/token/README.adoc                       | 2 ++
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/.changeset/new-crews-boil.md b/.changeset/new-crews-boil.md
index e8c18888..f4702693 100644
--- a/.changeset/new-crews-boil.md
+++ b/.changeset/new-crews-boil.md
@@ -2,4 +2,4 @@
 'openzeppelin-confidential-contracts': minor
 ---
 
-Add `ERC7984Rwa` extension.
+`ERC7984Rwa`: An extension of `ERC7984`, that supports confidential Real World Assets (RWAs).
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 7f420904..7067f5cf 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -13,7 +13,8 @@ import {ERC7984} from "./../ERC7984.sol";
 import {ERC7984Freezable} from "./ERC7984Freezable.sol";
 
 /**
- * @dev Extension of {ERC7984} supporting confidential Real World Assets.
+ * @dev Extension of {ERC7984} that supports confidential Real World Assets (RWAs).
+ * This interface provides compliance checks, transfer controls and enforcement actions.
  */
 abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall, ERC165, AccessControl {
     bytes32 public constant AGENT_ROLE = keccak256("AGENT_ROLE");
diff --git a/contracts/token/README.adoc b/contracts/token/README.adoc
index 25777d53..cbe08e31 100644
--- a/contracts/token/README.adoc
+++ b/contracts/token/README.adoc
@@ -9,6 +9,7 @@ This set of interfaces, contracts, and utilities are all related to `ERC7984`, a
 - {ERC7984ERC20Wrapper}: Extension of {ERC7984} which wraps an `ERC20` into a confidential token. The wrapper allows for free conversion in both directions at a fixed rate.
 - {ERC7984Freezable}: An extension for {ERC7984}, which allows accounts granted the "freezer" role to freeze and unfreeze tokens.
 - {ERC7984ObserverAccess}: An extension for {ERC7984}, which allows each account to add an observer who is given access to their transfer and balance amounts.
+- {ERC7984Rwa}: Extension of {ERC7984} that supports confidential Real World Assets (RWAs) by providing compliance checks, transfer controls and enforcement actions.
 - {ERC7984Utils}: A library that provides the on-transfer callback check used by {ERC7984}.
 
 == Core
@@ -18,6 +19,7 @@ This set of interfaces, contracts, and utilities are all related to `ERC7984`, a
 {{ERC7984ERC20Wrapper}}
 {{ERC7984Freezable}}
 {{ERC7984ObserverAccess}}
+{{ERC7984Rwa}}
 
 == Utilities
 {{ERC7984Utils}}
\ No newline at end of file

From 64c6c9ba16cb932903b76b5bceb4828bc22736a8 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 29 Aug 2025 10:46:50 +0200
Subject: [PATCH 19/52] Swap items in doc

---
 contracts/token/README.adoc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/contracts/token/README.adoc b/contracts/token/README.adoc
index bb4b8ec8..724d9acc 100644
--- a/contracts/token/README.adoc
+++ b/contracts/token/README.adoc
@@ -20,8 +20,8 @@ This set of interfaces, contracts, and utilities are all related to `ERC7984`, a
 {{ERC7984ERC20Wrapper}}
 {{ERC7984Freezable}}
 {{ERC7984ObserverAccess}}
-{{ERC7984Rwa}}
 {{ERC7984Restricted}}
+{{ERC7984Rwa}}
 
 == Utilities
 {{ERC7984Utils}}
\ No newline at end of file

From 1ad9ccde6219f728ebf67a79bf832df5cd271760 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 29 Aug 2025 14:34:34 +0200
Subject: [PATCH 20/52] Add suggestions

---
 contracts/interfaces/IERC7984Rwa.sol          |  4 +--
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 28 +++++++++++--------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index c5ea9e48..63498c79 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -3,7 +3,7 @@ pragma solidity ^0.8.24;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
-import {IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 
 /// @dev Base interface for confidential RWA contracts.
@@ -47,7 +47,7 @@ interface IERC7984RwaBase {
         bytes calldata inputProof
     ) external returns (euint64);
     /// @dev Mints confidential amount of tokens to account.
-    function confidentialMint(address to, euint64 encryptedAmount) external;
+    function confidentialMint(address to, euint64 encryptedAmount) external returns (euint64);
     /// @dev Burns confidential amount of tokens from account with proof.
     function confidentialBurn(
         address account,
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index fc5aff6e..f3b47342 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -16,16 +16,16 @@ const frozenEventSignature = 'TokensFrozen(address,bytes32)';
 
 /* eslint-disable no-unexpected-multiline */
 describe('ERC7984Rwa', function () {
-  async function deployFixture() {
+  const deployFixture = async () => {
     const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
     const token = await ethers.deployContract('ERC7984RwaMock', ['name', 'symbol', 'uri']);
     await token.connect(admin).addAgent(agent1);
     token.connect(anyone);
     return { token, admin, agent1, agent2, recipient, anyone };
-  }
+  };
 
   describe('ERC165', async function () {
-    it('should support interfaces', async function () {
+    it('should support interface', async function () {
       const { token } = await deployFixture();
       const interfaceFactories = [
         IERC7984RwaBase__factory,
@@ -38,6 +38,10 @@ describe('ERC7984Rwa', function () {
         expect(await token.supportsInterface(getInterfaceId(functions))).is.true;
       }
     });
+    it('should not support interface', async function () {
+      const { token } = await deployFixture();
+      expect(await token.supportsInterface('0xbadbadba')).is.false;
+    });
   });
 
   describe('Pausable', async function () {
@@ -56,14 +60,14 @@ describe('ERC7984Rwa', function () {
       const { token, anyone } = await deployFixture();
       await expect(token.connect(anyone).pause())
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone);
+        .withArgs(anyone.address);
     });
 
     it('should not unpause if neither admin nor agent', async function () {
       const { token, anyone } = await deployFixture();
       await expect(token.connect(anyone).unpause())
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone);
+        .withArgs(anyone.address);
     });
   });
 
@@ -89,14 +93,14 @@ describe('ERC7984Rwa', function () {
       const { token, agent1, anyone } = await deployFixture();
       await expect(token.connect(anyone).addAgent(agent1))
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone);
+        .withArgs(anyone.address);
     });
 
     it('should not remove agent if neither admin nor agent', async function () {
       const { token, agent1, anyone } = await deployFixture();
       await expect(token.connect(anyone).removeAgent(agent1))
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone);
+        .withArgs(anyone.address);
     });
   });
 
@@ -154,7 +158,7 @@ describe('ERC7984Rwa', function () {
           ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone);
+        .withArgs(anyone.address);
     });
 
     it('should not mint if transfer not compliant', async function () {
@@ -169,7 +173,7 @@ describe('ERC7984Rwa', function () {
           ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
-        .withArgs(ethers.ZeroAddress, recipient, encryptedInput.handles[0]);
+        .withArgs(ethers.ZeroAddress, recipient.address, encryptedInput.handles[0]);
     });
 
     it('should not mint if paused', async function () {
@@ -257,7 +261,7 @@ describe('ERC7984Rwa', function () {
           ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone);
+        .withArgs(anyone.address);
     });
 
     it('should not mint if transfer not compliant', async function () {
@@ -272,7 +276,7 @@ describe('ERC7984Rwa', function () {
           ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
-        .withArgs(recipient, ethers.ZeroAddress, encryptedInput.handles[0]);
+        .withArgs(recipient.address, ethers.ZeroAddress, encryptedInput.handles[0]);
     });
 
     it('should not burn if paused', async function () {
@@ -585,7 +589,7 @@ describe('ERC7984Rwa', function () {
           ),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
-        .withArgs(recipient, anyone, encryptedTransferValueInput.handles[0]);
+        .withArgs(recipient.address, anyone.address, encryptedTransferValueInput.handles[0]);
     });
 
     it('should not transfer if frozen', async function () {

From 628d1434730726cf065e4f3777011a022a63c087 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 29 Aug 2025 15:01:39 +0200
Subject: [PATCH 21/52] Remove lint annotation

---
 test/token/ERC7984/extensions/ERC7984Rwa.test.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index f3b47342..917514d6 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -14,7 +14,6 @@ import { ethers, fhevm } from 'hardhat';
 const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
 const frozenEventSignature = 'TokensFrozen(address,bytes32)';
 
-/* eslint-disable no-unexpected-multiline */
 describe('ERC7984Rwa', function () {
   const deployFixture = async () => {
     const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
@@ -649,4 +648,3 @@ describe('ERC7984Rwa', function () {
     });
   });
 });
-/* eslint-disable no-unexpected-multiline */

From 0b1d87cbcad5123584a342540158da7da3cede26 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 29 Aug 2025 15:11:27 +0200
Subject: [PATCH 22/52] Update test name

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
---
 test/token/ERC7984/extensions/ERC7984Rwa.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 917514d6..a7115ba2 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -263,7 +263,7 @@ describe('ERC7984Rwa', function () {
         .withArgs(anyone.address);
     });
 
-    it('should not mint if transfer not compliant', async function () {
+    it('should not burn if transfer not compliant', async function () {
       const { token, admin, recipient } = await deployFixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)

From b6b68270ce37335647b340e7e1d8b72ea155e397 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 3 Sep 2025 16:03:58 +0200
Subject: [PATCH 23/52] Add restriction to ERC7984Rwa

---
 contracts/interfaces/IERC7984Restricted.sol   | 23 +++++++++
 contracts/interfaces/IERC7984Rwa.sol          |  9 ++++
 .../ERC7984/extensions/ERC7984Restricted.sol  | 15 +-----
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 29 ++++++++++--
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 47 +++++++++++++++++++
 5 files changed, 105 insertions(+), 18 deletions(-)
 create mode 100644 contracts/interfaces/IERC7984Restricted.sol

diff --git a/contracts/interfaces/IERC7984Restricted.sol b/contracts/interfaces/IERC7984Restricted.sol
new file mode 100644
index 00000000..c129460f
--- /dev/null
+++ b/contracts/interfaces/IERC7984Restricted.sol
@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+/// @dev Interface for contracts that implements user account transfer restrictions.
+interface IERC7984Restricted {
+    enum Restriction {
+        DEFAULT, // User has no explicit restriction
+        BLOCKED, // User is explicitly blocked
+        ALLOWED // User is explicitly allowed
+    }
+
+    /// @dev Emitted when a user account's restriction is updated.
+    event UserRestrictionUpdated(address indexed account, Restriction restriction);
+
+    /// @dev The operation failed because the user account is restricted.
+    error UserRestricted(address account);
+
+    /// @dev Returns the restriction of a user account.
+    function getRestriction(address account) external view returns (Restriction);
+    /// @dev Returns whether a user account is allowed to interact with the token.
+    function isUserAllowed(address account) external view returns (bool);
+}
diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 63498c79..af72e460 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -5,6 +5,7 @@ import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
 import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
+import {IERC7984Restricted} from "./IERC7984Restricted.sol";
 
 /// @dev Base interface for confidential RWA contracts.
 interface IERC7984RwaBase {
@@ -28,6 +29,14 @@ interface IERC7984RwaBase {
     function pause() external;
     /// @dev Unpauses contract.
     function unpause() external;
+    /// @dev Returns the restriction of a user account.
+    function getRestriction(address account) external view returns (IERC7984Restricted.Restriction);
+    /// @dev Blocks an account.
+    function block(address account) external;
+    /// @dev Unblocks an account.
+    function unblock(address account) external;
+    /// @dev Returns whether an account is allowed to interact with the token.
+    function isUserAllowed(address account) external view returns (bool);
     /// @dev Returns the confidential frozen balance of an account.
     function confidentialFrozen(address account) external view returns (euint64);
     /// @dev Returns the available (unfrozen) balance of an account. Up to {confidentialBalanceOf}.
diff --git a/contracts/token/ERC7984/extensions/ERC7984Restricted.sol b/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
index de66f2e7..1f3dbae9 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
@@ -2,6 +2,7 @@
 
 pragma solidity ^0.8.27;
 
+import {IERC7984Restricted} from "../../../interfaces/IERC7984Restricted.sol";
 import {ERC7984, euint64} from "../ERC7984.sol";
 
 /**
@@ -13,21 +14,9 @@ import {ERC7984, euint64} from "../ERC7984.sol";
  * a blocklist. Developers can override {isUserAllowed} to check that `restriction == ALLOWED`
  * to implement an allowlist.
  */
-abstract contract ERC7984Restricted is ERC7984 {
-    enum Restriction {
-        DEFAULT, // User has no explicit restriction
-        BLOCKED, // User is explicitly blocked
-        ALLOWED // User is explicitly allowed
-    }
-
+abstract contract ERC7984Restricted is ERC7984, IERC7984Restricted {
     mapping(address account => Restriction) private _restrictions;
 
-    /// @dev Emitted when a user account's restriction is updated.
-    event UserRestrictionUpdated(address indexed account, Restriction restriction);
-
-    /// @dev The operation failed because the user account is restricted.
-    error UserRestricted(address account);
-
     /// @dev Returns the restriction of a user account.
     function getRestriction(address account) public view virtual returns (Restriction) {
         return _restrictions[account];
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 7067f5cf..ffc29341 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -11,12 +11,21 @@ import {IERC7984} from "./../../../interfaces/IERC7984.sol";
 import {IERC7984RwaBase} from "./../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984} from "./../ERC7984.sol";
 import {ERC7984Freezable} from "./ERC7984Freezable.sol";
+import {ERC7984Restricted} from "./ERC7984Restricted.sol";
 
 /**
  * @dev Extension of {ERC7984} that supports confidential Real World Assets (RWAs).
  * This interface provides compliance checks, transfer controls and enforcement actions.
  */
-abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall, ERC165, AccessControl {
+abstract contract ERC7984Rwa is
+    ERC7984,
+    ERC7984Freezable,
+    ERC7984Restricted,
+    Pausable,
+    Multicall,
+    ERC165,
+    AccessControl
+{
     bytes32 public constant AGENT_ROLE = keccak256("AGENT_ROLE");
 
     /// @dev The caller account is not authorized to perform the operation.
@@ -72,6 +81,16 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
         _removeAgent(account);
     }
 
+    /// @dev Blocks an account.
+    function block(address account) public virtual onlyAdminOrAgent {
+        _blockUser(account);
+    }
+
+    /// @dev Unblocks an account.
+    function unblock(address account) public virtual onlyAdminOrAgent {
+        _allowUser(account);
+    }
+
     /// @dev Mints confidential amount of tokens to account with proof.
     function confidentialMint(
         address to,
@@ -152,21 +171,21 @@ abstract contract ERC7984Rwa is ERC7984, ERC7984Freezable, Pausable, Multicall,
         euint64 encryptedAmount
     ) internal virtual onlyAdminOrAgent returns (euint64 transferred) {
         euint64 available = confidentialAvailable(from);
-        transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen & compliance checks
+        transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen, restrictions & compliance checks
         _setConfidentialFrozen(
             from,
             FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from))
         );
     }
 
-    /// @dev Internal function which updates confidential balances while performing frozen and compliance checks.
+    /// @dev Internal function which updates confidential balances while performing frozen, restrictions and compliance checks.
     function _update(
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal override(ERC7984, ERC7984Freezable) whenNotPaused returns (euint64) {
+    ) internal override(ERC7984, ERC7984Freezable, ERC7984Restricted) whenNotPaused returns (euint64) {
         require(_isCompliantTransfer(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
-        // frozen check performed through inheritance
+        // frozen and restrictions checks performed through inheritance
         return super._update(from, to, encryptedAmount);
     }
 
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index a7115ba2..523d8794 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -103,6 +103,28 @@ describe('ERC7984Rwa', function () {
     });
   });
 
+  describe('ERC7984Restricted', async function () {
+    it('should block & unblock', async function () {
+      const { token, admin, agent1, recipient } = await deployFixture();
+      for (const manager of [admin, agent1]) {
+        await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
+        await token.connect(manager).block(recipient);
+        await expect(token.isUserAllowed(recipient)).to.eventually.be.false;
+        await token.connect(manager).unblock(recipient);
+        await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
+      }
+    });
+
+    for (const arg of [true, false]) {
+      it(`should not ${arg ? 'block' : 'unblock'} if neither admin nor agent`, async function () {
+        const { token, anyone } = await deployFixture();
+        await expect(token.connect(anyone)[arg ? 'block' : 'unblock'](anyone))
+          .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+          .withArgs(anyone.address);
+      });
+    }
+  });
+
   describe('Mintable', async function () {
     for (const withProof of [true, false]) {
       it(`should mint by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
@@ -646,5 +668,30 @@ describe('ERC7984Rwa', function () {
         ),
       ).to.eventually.equal(100);
     });
+
+    for (const arg of [true, false]) {
+      it(`should not transfer if ${arg ? 'sender' : 'receiver'} blocked `, async function () {
+        const { token, admin: manager, recipient, anyone } = await deployFixture();
+        const account = arg ? recipient : anyone;
+        await token.$_setCompliantTransfer();
+        const encryptedInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), recipient.address)
+          .add64(25)
+          .encrypt();
+        await token.connect(manager).block(account);
+
+        await expect(
+          token
+            .connect(recipient)
+            ['confidentialTransfer(address,bytes32,bytes)'](
+              anyone,
+              encryptedInput.handles[0],
+              encryptedInput.inputProof,
+            ),
+        )
+          .to.be.revertedWithCustomError(token, 'UserRestricted')
+          .withArgs(account);
+      });
+    }
   });
 });

From 3185336e44f27f78acc114863baa9cb9cc034cd8 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 3 Sep 2025 16:21:23 +0200
Subject: [PATCH 24/52] Move gates

---
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 35 +++++++++----------
 1 file changed, 16 insertions(+), 19 deletions(-)

diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index ffc29341..9a3f4c4b 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -72,12 +72,12 @@ abstract contract ERC7984Rwa is
     }
 
     /// @dev Adds agent.
-    function addAgent(address account) public virtual {
+    function addAgent(address account) public virtual onlyAdminOrAgent {
         _addAgent(account);
     }
 
     /// @dev Removes agent.
-    function removeAgent(address account) public virtual {
+    function removeAgent(address account) public virtual onlyAdminOrAgent {
         _removeAgent(account);
     }
 
@@ -96,12 +96,12 @@ abstract contract ERC7984Rwa is
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual returns (euint64) {
+    ) public virtual onlyAdminOrAgent returns (euint64) {
         return _confidentialMint(to, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Mints confidential amount of tokens to account.
-    function confidentialMint(address to, euint64 encryptedAmount) public virtual returns (euint64) {
+    function confidentialMint(address to, euint64 encryptedAmount) public virtual onlyAdminOrAgent returns (euint64) {
         return _confidentialMint(to, encryptedAmount);
     }
 
@@ -110,12 +110,15 @@ abstract contract ERC7984Rwa is
         address account,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual returns (euint64) {
+    ) public virtual onlyAdminOrAgent returns (euint64) {
         return _confidentialBurn(account, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Burns confidential amount of tokens from account.
-    function confidentialBurn(address account, euint64 encryptedAmount) public virtual returns (euint64) {
+    function confidentialBurn(
+        address account,
+        euint64 encryptedAmount
+    ) public virtual onlyAdminOrAgent returns (euint64) {
         return _confidentialBurn(account, encryptedAmount);
     }
 
@@ -125,7 +128,7 @@ abstract contract ERC7984Rwa is
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual returns (euint64) {
+    ) public virtual onlyAdminOrAgent returns (euint64) {
         return _forceConfidentialTransferFrom(from, to, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
@@ -134,33 +137,27 @@ abstract contract ERC7984Rwa is
         address from,
         address to,
         euint64 encryptedAmount
-    ) public virtual returns (euint64 transferred) {
+    ) public virtual onlyAdminOrAgent returns (euint64 transferred) {
         return _forceConfidentialTransferFrom(from, to, encryptedAmount);
     }
 
     /// @dev Internal function which adds an agent.
-    function _addAgent(address account) internal virtual onlyAdminOrAgent {
+    function _addAgent(address account) internal virtual {
         _grantRole(AGENT_ROLE, account);
     }
 
     /// @dev Internal function which removes an agent.
-    function _removeAgent(address account) internal virtual onlyAdminOrAgent {
+    function _removeAgent(address account) internal virtual {
         _revokeRole(AGENT_ROLE, account);
     }
 
     /// @dev Internal function which mints confidential amount of tokens to account.
-    function _confidentialMint(
-        address to,
-        euint64 encryptedAmount
-    ) internal virtual onlyAdminOrAgent returns (euint64) {
+    function _confidentialMint(address to, euint64 encryptedAmount) internal virtual returns (euint64) {
         return _mint(to, encryptedAmount);
     }
 
     /// @dev Internal function which burns confidential amount of tokens from account.
-    function _confidentialBurn(
-        address account,
-        euint64 encryptedAmount
-    ) internal virtual onlyAdminOrAgent returns (euint64) {
+    function _confidentialBurn(address account, euint64 encryptedAmount) internal virtual returns (euint64) {
         return _burn(account, encryptedAmount);
     }
 
@@ -169,7 +166,7 @@ abstract contract ERC7984Rwa is
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal virtual onlyAdminOrAgent returns (euint64 transferred) {
+    ) internal virtual returns (euint64 transferred) {
         euint64 available = confidentialAvailable(from);
         transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen, restrictions & compliance checks
         _setConfidentialFrozen(

From b4f8c03ff217d4972b3fb3cf4b8c5fea07bf87f3 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 3 Sep 2025 16:45:12 +0200
Subject: [PATCH 25/52] Remove ExpectedPause error

---
 contracts/interfaces/IERC7984Rwa.sol | 2 --
 1 file changed, 2 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index af72e460..8c5e1d7c 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -20,8 +20,6 @@ interface IERC7984RwaBase {
     error OwnableInvalidOwner(address owner);
     /// @dev The operation failed because the contract is paused.
     error EnforcedPause();
-    /// @dev The operation failed because the contract is not paused.
-    error ExpectedPause();
 
     /// @dev Returns true if the contract is paused, and false otherwise.
     function paused() external view returns (bool);

From 28973a2c65796627de7dd0cb92fe33679871bbad Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 3 Sep 2025 17:04:58 +0200
Subject: [PATCH 26/52] Rename block functions

---
 contracts/interfaces/IERC7984Rwa.sol              | 8 ++++----
 contracts/token/ERC7984/extensions/ERC7984Rwa.sol | 8 ++++----
 test/token/ERC7984/extensions/ERC7984Rwa.test.ts  | 8 ++++----
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 8c5e1d7c..aad61239 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -29,10 +29,10 @@ interface IERC7984RwaBase {
     function unpause() external;
     /// @dev Returns the restriction of a user account.
     function getRestriction(address account) external view returns (IERC7984Restricted.Restriction);
-    /// @dev Blocks an account.
-    function block(address account) external;
-    /// @dev Unblocks an account.
-    function unblock(address account) external;
+    /// @dev Blocks a user account.
+    function blockUser(address account) external;
+    /// @dev Unblocks a user account.
+    function unblockUser(address account) external;
     /// @dev Returns whether an account is allowed to interact with the token.
     function isUserAllowed(address account) external view returns (bool);
     /// @dev Returns the confidential frozen balance of an account.
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 9a3f4c4b..02d3f2fb 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -81,13 +81,13 @@ abstract contract ERC7984Rwa is
         _removeAgent(account);
     }
 
-    /// @dev Blocks an account.
-    function block(address account) public virtual onlyAdminOrAgent {
+    /// @dev Blocks a user account.
+    function blockUser(address account) public virtual onlyAdminOrAgent {
         _blockUser(account);
     }
 
-    /// @dev Unblocks an account.
-    function unblock(address account) public virtual onlyAdminOrAgent {
+    /// @dev Unblocks a user account.
+    function unblockUser(address account) public virtual onlyAdminOrAgent {
         _allowUser(account);
     }
 
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 523d8794..ef50e729 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -108,9 +108,9 @@ describe('ERC7984Rwa', function () {
       const { token, admin, agent1, recipient } = await deployFixture();
       for (const manager of [admin, agent1]) {
         await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
-        await token.connect(manager).block(recipient);
+        await token.connect(manager).blockUser(recipient);
         await expect(token.isUserAllowed(recipient)).to.eventually.be.false;
-        await token.connect(manager).unblock(recipient);
+        await token.connect(manager).unblockUser(recipient);
         await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
       }
     });
@@ -118,7 +118,7 @@ describe('ERC7984Rwa', function () {
     for (const arg of [true, false]) {
       it(`should not ${arg ? 'block' : 'unblock'} if neither admin nor agent`, async function () {
         const { token, anyone } = await deployFixture();
-        await expect(token.connect(anyone)[arg ? 'block' : 'unblock'](anyone))
+        await expect(token.connect(anyone)[arg ? 'blockUser' : 'unblockUser'](anyone))
           .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
           .withArgs(anyone.address);
       });
@@ -678,7 +678,7 @@ describe('ERC7984Rwa', function () {
           .createEncryptedInput(await token.getAddress(), recipient.address)
           .add64(25)
           .encrypt();
-        await token.connect(manager).block(account);
+        await token.connect(manager).blockUser(account);
 
         await expect(
           token

From 0facae57fbbe3bb3849503371f3718b52e4864b5 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 3 Sep 2025 17:07:44 +0200
Subject: [PATCH 27/52] Rename fixture

---
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 78 +++++++++----------
 1 file changed, 39 insertions(+), 39 deletions(-)

diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index ef50e729..cc2d0884 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -14,18 +14,18 @@ import { ethers, fhevm } from 'hardhat';
 const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
 const frozenEventSignature = 'TokensFrozen(address,bytes32)';
 
-describe('ERC7984Rwa', function () {
-  const deployFixture = async () => {
-    const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
-    const token = await ethers.deployContract('ERC7984RwaMock', ['name', 'symbol', 'uri']);
-    await token.connect(admin).addAgent(agent1);
-    token.connect(anyone);
-    return { token, admin, agent1, agent2, recipient, anyone };
-  };
+const fixture = async () => {
+  const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
+  const token = await ethers.deployContract('ERC7984RwaMock', ['name', 'symbol', 'uri']);
+  await token.connect(admin).addAgent(agent1);
+  token.connect(anyone);
+  return { token, admin, agent1, agent2, recipient, anyone };
+};
 
+describe('ERC7984Rwa', function () {
   describe('ERC165', async function () {
     it('should support interface', async function () {
-      const { token } = await deployFixture();
+      const { token } = await fixture();
       const interfaceFactories = [
         IERC7984RwaBase__factory,
         IERC7984__factory,
@@ -38,14 +38,14 @@ describe('ERC7984Rwa', function () {
       }
     });
     it('should not support interface', async function () {
-      const { token } = await deployFixture();
+      const { token } = await fixture();
       expect(await token.supportsInterface('0xbadbadba')).is.false;
     });
   });
 
   describe('Pausable', async function () {
     it('should pause & unpause', async function () {
-      const { token, admin, agent1 } = await deployFixture();
+      const { token, admin, agent1 } = await fixture();
       for (const manager of [admin, agent1]) {
         expect(await token.paused()).is.false;
         await token.connect(manager).pause();
@@ -56,14 +56,14 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not pause if neither admin nor agent', async function () {
-      const { token, anyone } = await deployFixture();
+      const { token, anyone } = await fixture();
       await expect(token.connect(anyone).pause())
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
         .withArgs(anyone.address);
     });
 
     it('should not unpause if neither admin nor agent', async function () {
-      const { token, anyone } = await deployFixture();
+      const { token, anyone } = await fixture();
       await expect(token.connect(anyone).unpause())
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
         .withArgs(anyone.address);
@@ -72,13 +72,13 @@ describe('ERC7984Rwa', function () {
 
   describe('Roles', async function () {
     it('should check admin', async function () {
-      const { token, admin, anyone } = await deployFixture();
+      const { token, admin, anyone } = await fixture();
       expect(await token.isAdmin(admin)).is.true;
       expect(await token.isAdmin(anyone)).is.false;
     });
 
     it('should check/add/remove agent', async function () {
-      const { token, admin, agent1, agent2 } = await deployFixture();
+      const { token, admin, agent1, agent2 } = await fixture();
       for (const manager of [admin, agent1]) {
         expect(await token.isAgent(agent2)).is.false;
         await token.connect(manager).addAgent(agent2);
@@ -89,14 +89,14 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not add agent if neither admin nor agent', async function () {
-      const { token, agent1, anyone } = await deployFixture();
+      const { token, agent1, anyone } = await fixture();
       await expect(token.connect(anyone).addAgent(agent1))
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
         .withArgs(anyone.address);
     });
 
     it('should not remove agent if neither admin nor agent', async function () {
-      const { token, agent1, anyone } = await deployFixture();
+      const { token, agent1, anyone } = await fixture();
       await expect(token.connect(anyone).removeAgent(agent1))
         .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
         .withArgs(anyone.address);
@@ -105,7 +105,7 @@ describe('ERC7984Rwa', function () {
 
   describe('ERC7984Restricted', async function () {
     it('should block & unblock', async function () {
-      const { token, admin, agent1, recipient } = await deployFixture();
+      const { token, admin, agent1, recipient } = await fixture();
       for (const manager of [admin, agent1]) {
         await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
         await token.connect(manager).blockUser(recipient);
@@ -117,7 +117,7 @@ describe('ERC7984Rwa', function () {
 
     for (const arg of [true, false]) {
       it(`should not ${arg ? 'block' : 'unblock'} if neither admin nor agent`, async function () {
-        const { token, anyone } = await deployFixture();
+        const { token, anyone } = await fixture();
         await expect(token.connect(anyone)[arg ? 'blockUser' : 'unblockUser'](anyone))
           .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
           .withArgs(anyone.address);
@@ -128,9 +128,9 @@ describe('ERC7984Rwa', function () {
   describe('Mintable', async function () {
     for (const withProof of [true, false]) {
       it(`should mint by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient } = await deployFixture();
+        const { admin, agent1, recipient } = await fixture();
         for (const manager of [admin, agent1]) {
-          const { token } = await deployFixture();
+          const { token } = await fixture();
           await token.$_setCompliantTransfer();
           const amount = 100;
           let params = [recipient.address] as unknown as [
@@ -167,7 +167,7 @@ describe('ERC7984Rwa', function () {
     }
 
     it('should not mint if neither admin nor agent', async function () {
-      const { token, recipient, anyone } = await deployFixture();
+      const { token, recipient, anyone } = await fixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
@@ -183,7 +183,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not mint if transfer not compliant', async function () {
-      const { token, admin, recipient } = await deployFixture();
+      const { token, admin, recipient } = await fixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
@@ -198,7 +198,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not mint if paused', async function () {
-      const { token, admin, recipient } = await deployFixture();
+      const { token, admin, recipient } = await fixture();
       await token.connect(admin).pause();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
@@ -215,9 +215,9 @@ describe('ERC7984Rwa', function () {
   describe('Burnable', async function () {
     for (const withProof of [true, false]) {
       it(`should burn by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient } = await deployFixture();
+        const { admin, agent1, recipient } = await fixture();
         for (const manager of [admin, agent1]) {
-          const { token } = await deployFixture();
+          const { token } = await fixture();
           const encryptedInput = await fhevm
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
@@ -270,7 +270,7 @@ describe('ERC7984Rwa', function () {
     }
 
     it('should not burn if neither admin nor agent', async function () {
-      const { token, recipient, anyone } = await deployFixture();
+      const { token, recipient, anyone } = await fixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
@@ -286,7 +286,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not burn if transfer not compliant', async function () {
-      const { token, admin, recipient } = await deployFixture();
+      const { token, admin, recipient } = await fixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
@@ -301,7 +301,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not burn if paused', async function () {
-      const { token, admin, recipient } = await deployFixture();
+      const { token, admin, recipient } = await fixture();
       await token.connect(admin).pause();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
@@ -318,9 +318,9 @@ describe('ERC7984Rwa', function () {
   describe('Force transfer', async function () {
     for (const withProof of [true, false]) {
       it(`should force transfer by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient, anyone } = await deployFixture();
+        const { admin, agent1, recipient, anyone } = await fixture();
         for (const manager of [admin, agent1]) {
-          const { token } = await deployFixture();
+          const { token } = await fixture();
           const encryptedMintValueInput = await fhevm
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
@@ -395,9 +395,9 @@ describe('ERC7984Rwa', function () {
 
     for (const withProof of [true, false]) {
       it(`should force transfer even if frozen ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient, anyone } = await deployFixture();
+        const { admin, agent1, recipient, anyone } = await fixture();
         for (const manager of [admin, agent1]) {
-          const { token } = await deployFixture();
+          const { token } = await fixture();
           const encryptedMintValueInput = await fhevm
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
@@ -475,7 +475,7 @@ describe('ERC7984Rwa', function () {
 
     for (const withProof of [true, false]) {
       it(`should not force transfer if neither admin nor agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { token, recipient, anyone } = await deployFixture();
+        const { token, recipient, anyone } = await fixture();
         let params = [recipient.address, anyone.address] as unknown as [
           from: AddressLike,
           to: AddressLike,
@@ -510,7 +510,7 @@ describe('ERC7984Rwa', function () {
 
   describe('Transfer', async function () {
     it('should transfer', async function () {
-      const { token, admin: manager, recipient, anyone } = await deployFixture();
+      const { token, admin: manager, recipient, anyone } = await fixture();
       const encryptedMintValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
@@ -576,7 +576,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not transfer if paused', async function () {
-      const { token, admin: manager, recipient, anyone } = await deployFixture();
+      const { token, admin: manager, recipient, anyone } = await fixture();
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
@@ -594,7 +594,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not transfer if transfer not compliant', async function () {
-      const { token, recipient, anyone } = await deployFixture();
+      const { token, recipient, anyone } = await fixture();
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
@@ -614,7 +614,7 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not transfer if frozen', async function () {
-      const { token, admin: manager, recipient, anyone } = await deployFixture();
+      const { token, admin: manager, recipient, anyone } = await fixture();
       const encryptedMintValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
@@ -671,7 +671,7 @@ describe('ERC7984Rwa', function () {
 
     for (const arg of [true, false]) {
       it(`should not transfer if ${arg ? 'sender' : 'receiver'} blocked `, async function () {
-        const { token, admin: manager, recipient, anyone } = await deployFixture();
+        const { token, admin: manager, recipient, anyone } = await fixture();
         const account = arg ? recipient : anyone;
         await token.$_setCompliantTransfer();
         const encryptedInput = await fhevm

From 2d0ef0ea9c9ba36f030d3364e2e8cbd3b103ab47 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Thu, 4 Sep 2025 13:13:38 +0200
Subject: [PATCH 28/52] Force transfer with all update effects

---
 .../ERC7984/extensions/ERC7984Freezable.sol   | 43 +++++++++-
 .../ERC7984/extensions/ERC7984Restricted.sol  | 22 ++++-
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 18 ++--
 .../extensions/ERC7984Freezable.test.ts       | 84 +++++++++++++++++++
 .../extensions/ERC7984Restricted.test.ts      | 11 +++
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 37 +++++++-
 6 files changed, 199 insertions(+), 16 deletions(-)

diff --git a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
index 50fe89d5..66360036 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
@@ -21,6 +21,8 @@ import {ERC7984} from "../ERC7984.sol";
 abstract contract ERC7984Freezable is ERC7984 {
     /// @dev Confidential frozen amount of tokens per address.
     mapping(address account => euint64 encryptedAmount) private _frozenBalances;
+    /// @dev Skips frozen checks in {_update}.
+    bool private _skipUpdateCheck;
 
     /// @dev Emitted when a confidential amount of token is frozen for an account
     event TokensFrozen(address indexed account, euint64 encryptedAmount);
@@ -59,7 +61,14 @@ abstract contract ERC7984Freezable is ERC7984 {
 
     /// @dev Internal function to freeze a confidential amount of tokens for an account.
     function _setConfidentialFrozen(address account, euint64 encryptedAmount) internal virtual {
-        _checkFreezer();
+        _setConfidentialFrozen(account, encryptedAmount, true);
+    }
+
+    /// @dev Private function to freeze a confidential amount of tokens for an account
+    function _setConfidentialFrozen(address account, euint64 encryptedAmount, bool checkFreezer) internal virtual {
+        if (checkFreezer) {
+            _checkFreezer();
+        }
         FHE.allowThis(encryptedAmount);
         FHE.allow(encryptedAmount, account);
         _frozenBalances[account] = encryptedAmount;
@@ -69,15 +78,41 @@ abstract contract ERC7984Freezable is ERC7984 {
     /// @dev Unimplemented function that must revert if `msg.sender` is not authorized as a freezer.
     function _checkFreezer() internal virtual;
 
+    /// @dev Internal function to skip update check. Check can be restored with {_restoreERC7984FreezableUpdateCheck}.
+    function _disableERC7984FreezableUpdateCheck() internal virtual {
+        if (!_skipUpdateCheck) {
+            _skipUpdateCheck = true;
+        }
+    }
+
+    /// @dev Internal function to restore update check previously disabled by {_disableERC7984FreezableUpdateCheck}.
+    function _restoreERC7984FreezableUpdateCheck() internal virtual {
+        if (_skipUpdateCheck) {
+            _skipUpdateCheck = false;
+        }
+    }
+
     /**
      * @dev See {ERC7984-_update}. The `from` account must have sufficient unfrozen balance,
      * otherwise 0 tokens are transferred.
      */
     function _update(address from, address to, euint64 encryptedAmount) internal virtual override returns (euint64) {
+        euint64 available;
         if (from != address(0)) {
-            euint64 unfrozen = confidentialAvailable(from);
-            encryptedAmount = FHE.select(FHE.le(encryptedAmount, unfrozen), encryptedAmount, FHE.asEuint64(0));
+            available = confidentialAvailable(from);
+            if (!_skipUpdateCheck) {
+                encryptedAmount = FHE.select(FHE.le(encryptedAmount, available), encryptedAmount, FHE.asEuint64(0));
+            }
+        }
+        euint64 transferred = super._update(from, to, encryptedAmount);
+        if (from != address(0) && _skipUpdateCheck) {
+            // Reset frozen to balance if transferred more than available
+            _setConfidentialFrozen(
+                from,
+                FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from)),
+                false
+            );
         }
-        return super._update(from, to, encryptedAmount);
+        return transferred;
     }
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984Restricted.sol b/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
index 1f3dbae9..37ca2d93 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
@@ -16,6 +16,8 @@ import {ERC7984, euint64} from "../ERC7984.sol";
  */
 abstract contract ERC7984Restricted is ERC7984, IERC7984Restricted {
     mapping(address account => Restriction) private _restrictions;
+    /// @dev Skips restriction checks in {_update}.
+    bool private _skipUpdateCheck;
 
     /// @dev Returns the restriction of a user account.
     function getRestriction(address account) public view virtual returns (Restriction) {
@@ -39,6 +41,20 @@ abstract contract ERC7984Restricted is ERC7984, IERC7984Restricted {
         return getRestriction(account) != Restriction.BLOCKED; // i.e. DEFAULT && ALLOWED
     }
 
+    /// @dev Internal function to skip update check. Check can be restored with {_restoreERC7984RestrictedUpdateCheck}.
+    function _disableERC7984RestrictedUpdateCheck() internal virtual {
+        if (!_skipUpdateCheck) {
+            _skipUpdateCheck = true;
+        }
+    }
+
+    /// @dev Internal function to restore update check previously disabled by {_disableERC7984RestrictedUpdateCheck}.
+    function _restoreERC7984RestrictedUpdateCheck() internal virtual {
+        if (_skipUpdateCheck) {
+            _skipUpdateCheck = false;
+        }
+    }
+
     /**
      * @dev See {ERC7984-_update}. Enforces transfer restrictions (excluding minting and burning).
      *
@@ -48,8 +64,10 @@ abstract contract ERC7984Restricted is ERC7984, IERC7984Restricted {
      * * `to` must be allowed to receive tokens (see {isUserAllowed}).
      */
     function _update(address from, address to, euint64 value) internal virtual override returns (euint64) {
-        if (from != address(0)) _checkRestriction(from); // Not minting
-        if (to != address(0)) _checkRestriction(to); // Not burning
+        if (!_skipUpdateCheck) {
+            if (from != address(0)) _checkRestriction(from); // Not minting
+            if (to != address(0)) _checkRestriction(to); // Not burning
+        }
         return super._update(from, to, value);
     }
 
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 02d3f2fb..9d6a1821 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -167,22 +167,22 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) internal virtual returns (euint64 transferred) {
-        euint64 available = confidentialAvailable(from);
-        transferred = ERC7984._update(from, to, encryptedAmount); // bypass frozen, restrictions & compliance checks
-        _setConfidentialFrozen(
-            from,
-            FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from))
-        );
+        _disableERC7984FreezableUpdateCheck(); // bypass frozen check
+        _disableERC7984RestrictedUpdateCheck(); // bypass default restriction check
+        if (to != address(0)) _checkRestriction(to); // only perform restriction check on `to`
+        transferred = super._update(from, to, encryptedAmount); // bypass compliance check
+        _restoreERC7984FreezableUpdateCheck();
+        _restoreERC7984RestrictedUpdateCheck();
     }
 
-    /// @dev Internal function which updates confidential balances while performing frozen, restrictions and compliance checks.
+    /// @dev Internal function which updates confidential balances while performing frozen, restriction and compliance checks.
     function _update(
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal override(ERC7984, ERC7984Freezable, ERC7984Restricted) whenNotPaused returns (euint64) {
+    ) internal override(ERC7984Freezable, ERC7984Restricted, ERC7984) whenNotPaused returns (euint64) {
         require(_isCompliantTransfer(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
-        // frozen and restrictions checks performed through inheritance
+        // frozen and restriction checks performed through inheritance
         return super._update(from, to, encryptedAmount);
     }
 
diff --git a/test/token/ERC7984/extensions/ERC7984Freezable.test.ts b/test/token/ERC7984/extensions/ERC7984Freezable.test.ts
index 6bef4f42..c9e7cb68 100644
--- a/test/token/ERC7984/extensions/ERC7984Freezable.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Freezable.test.ts
@@ -1,6 +1,7 @@
 import { IACL__factory } from '../../../../types';
 import { $ERC7984FreezableMock } from '../../../../types/contracts-exposed/mocks/token/ERC7984FreezableMock.sol/$ERC7984FreezableMock';
 import { ACL_ADDRESS } from '../../../helpers/accounts';
+import { callAndGetResult } from '../../../helpers/event';
 import { FhevmType } from '@fhevm/hardhat-plugin';
 import { expect } from 'chai';
 import { AddressLike, BytesLike, EventLog } from 'ethers';
@@ -217,6 +218,89 @@ describe('ERC7984Freezable', function () {
     ).to.eventually.equal(1000);
   });
 
+  it('should transfer all if transferring more than available but check disabled', async function () {
+    const { token, holder, recipient, freezer, anyone } = await deployFixture();
+    const encryptedRecipientMintInput = await fhevm
+      .createEncryptedInput(await token.getAddress(), holder.address)
+      .add64(1000)
+      .encrypt();
+    await token
+      .connect(holder)
+      ['$_mint(address,bytes32,bytes)'](
+        recipient.address,
+        encryptedRecipientMintInput.handles[0],
+        encryptedRecipientMintInput.inputProof,
+      );
+    const encryptedInput = await fhevm
+      .createEncryptedInput(await token.getAddress(), freezer.address)
+      .add64(500)
+      .encrypt();
+    await token
+      .connect(freezer)
+      ['setConfidentialFrozen(address,bytes32,bytes)'](
+        recipient.address,
+        encryptedInput.handles[0],
+        encryptedInput.inputProof,
+      );
+    await token.$_disableERC7984FreezableUpdateCheck();
+    const encryptedInput1 = await fhevm
+      .createEncryptedInput(await token.getAddress(), recipient.address)
+      .add64(501)
+      .encrypt();
+    const tx = await token
+      .connect(recipient)
+      ['confidentialTransfer(address,bytes32,bytes)'](
+        anyone.address,
+        encryptedInput1.handles[0],
+        encryptedInput1.inputProof,
+      );
+    await expect(tx).to.emit(token, 'ConfidentialTransfer');
+    const transferEvent = (await tx
+      .wait()
+      .then(receipt => receipt!.logs.filter((log: any) => log.address === token.target)[0])) as EventLog;
+    expect(transferEvent.args[0]).to.equal(recipient.address);
+    expect(transferEvent.args[1]).to.equal(anyone.address);
+    await expect(
+      fhevm.userDecryptEuint(FhevmType.euint64, transferEvent.args[2], await token.getAddress(), recipient),
+    ).to.eventually.equal(501);
+    await expect(
+      fhevm.userDecryptEuint(
+        FhevmType.euint64,
+        await token.confidentialBalanceOf(recipient.address),
+        await token.getAddress(),
+        recipient,
+      ),
+    ).to.eventually.equal(499);
+    await expect(
+      fhevm.userDecryptEuint(
+        FhevmType.euint64,
+        await token.confidentialFrozen(recipient),
+        await token.getAddress(),
+        recipient,
+      ),
+    ).to.eventually.equal(499); // frozen got reset to balance
+
+    // should transfer zero if frozen check is restored
+    await token.$_restoreERC7984FreezableUpdateCheck();
+    const encryptedInput2 = await fhevm
+      .createEncryptedInput(await token.getAddress(), recipient.address)
+      .add64(499)
+      .encrypt();
+    const [, , transferred] = await callAndGetResult(
+      token
+        .connect(recipient)
+        ['confidentialTransfer(address,bytes32,bytes)'](
+          anyone.address,
+          encryptedInput2.handles[0],
+          encryptedInput2.inputProof,
+        ),
+      'ConfidentialTransfer(address,address,bytes32)',
+    );
+    await expect(
+      fhevm.userDecryptEuint(FhevmType.euint64, transferred, await token.getAddress(), recipient),
+    ).to.eventually.equal(0);
+  });
+
   it('should not set confidential frozen if unauthorized', async function () {
     const { token, recipient, freezer, anyone } = await deployFixture();
     const encryptedInput = await fhevm
diff --git a/test/token/ERC7984/extensions/ERC7984Restricted.test.ts b/test/token/ERC7984/extensions/ERC7984Restricted.test.ts
index 302c9245..7f3e37af 100644
--- a/test/token/ERC7984/extensions/ERC7984Restricted.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Restricted.test.ts
@@ -79,6 +79,17 @@ describe('ERC7984Restricted', function () {
         await this.token.connect(this.holder).transfer(this.recipient, initialSupply);
       });
 
+      it('allows when sender and recipient are BLOCKED but restriction checks are disabled', async function () {
+        await this.token.$_blockUser(this.holder); // Sets to BLOCKED
+        await this.token.$_blockUser(this.recipient); // Sets to BLOCKED
+        await this.token.$_disableERC7984RestrictedUpdateCheck();
+        await this.token.connect(this.holder).transfer(this.recipient, initialSupply);
+        await this.token.$_restoreERC7984RestrictedUpdateCheck();
+        await expect(this.token.connect(this.holder).transfer(this.recipient, initialSupply))
+          .to.be.revertedWithCustomError(this.token, 'UserRestricted')
+          .withArgs(this.holder);
+      });
+
       it('reverts when sender is BLOCKED', async function () {
         await this.token.$_blockUser(this.holder); // Sets to BLOCKED
 
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index cc2d0884..ae9de6fb 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -468,7 +468,7 @@ describe('ERC7984Rwa', function () {
           await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
           await expect(
             fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(75); // frozen got reset to available balance
+          ).to.eventually.equal(75); // frozen got reset to balance
         }
       });
     }
@@ -506,6 +506,41 @@ describe('ERC7984Rwa', function () {
           .withArgs(anyone.address);
       });
     }
+
+    for (const withProof of [true, false]) {
+      it(`should not force transfer if receiver blocked ${withProof ? 'with proof' : ''}`, async function () {
+        const { token, recipient, anyone } = await fixture();
+        let params = [recipient.address, anyone.address] as unknown as [
+          from: AddressLike,
+          to: AddressLike,
+          encryptedAmount: BytesLike,
+          inputProof: BytesLike,
+        ];
+        const amount = 100;
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await token.getAddress(), anyone.address)
+            .add64(amount)
+            .encrypt();
+          params.push(handles[0], inputProof);
+        } else {
+          await token.connect(anyone).createEncryptedAmount(amount);
+          params.push(await token.connect(anyone).createEncryptedAmount.staticCall(amount));
+        }
+        await token.blockUser(anyone);
+        await expect(
+          token
+            .connect(anyone)
+            [
+              withProof
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'forceConfidentialTransferFrom(address,address,bytes32)'
+            ](...params),
+        )
+          .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+          .withArgs(anyone.address);
+      });
+    }
   });
 
   describe('Transfer', async function () {

From 54517770b6d0dc0fd0717c40acb8af820de853a7 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Thu, 4 Sep 2025 13:23:23 +0200
Subject: [PATCH 29/52] Update set frozen doc

---
 contracts/token/ERC7984/extensions/ERC7984Freezable.sol | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
index 66360036..025be19a 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
@@ -64,7 +64,7 @@ abstract contract ERC7984Freezable is ERC7984 {
         _setConfidentialFrozen(account, encryptedAmount, true);
     }
 
-    /// @dev Private function to freeze a confidential amount of tokens for an account
+    /// @dev Private function to freeze a confidential amount of tokens for an account with optional freezer check.
     function _setConfidentialFrozen(address account, euint64 encryptedAmount, bool checkFreezer) internal virtual {
         if (checkFreezer) {
             _checkFreezer();

From 30650021a9d6e5b3a59d7150ac395d2ac28c1a45 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Thu, 4 Sep 2025 16:23:09 +0200
Subject: [PATCH 30/52] Refactor event checks in freezable tests

---
 contracts/interfaces/IERC7984Rwa.sol          |  2 +-
 .../extensions/ERC7984Freezable.test.ts       | 21 +++++++++----------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index aad61239..bbf51585 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
-import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
 import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
diff --git a/test/token/ERC7984/extensions/ERC7984Freezable.test.ts b/test/token/ERC7984/extensions/ERC7984Freezable.test.ts
index c9e7cb68..72496ed5 100644
--- a/test/token/ERC7984/extensions/ERC7984Freezable.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Freezable.test.ts
@@ -10,6 +10,7 @@ import { ethers, fhevm } from 'hardhat';
 const name = 'ConfidentialFungibleToken';
 const symbol = 'CFT';
 const uri = 'https://example.com/metadata';
+const confidentialTransferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
 
 describe('ERC7984Freezable', function () {
   async function deployFixture() {
@@ -247,7 +248,7 @@ describe('ERC7984Freezable', function () {
       .createEncryptedInput(await token.getAddress(), recipient.address)
       .add64(501)
       .encrypt();
-    const tx = await token
+    const tx = token
       .connect(recipient)
       ['confidentialTransfer(address,bytes32,bytes)'](
         anyone.address,
@@ -255,13 +256,11 @@ describe('ERC7984Freezable', function () {
         encryptedInput1.inputProof,
       );
     await expect(tx).to.emit(token, 'ConfidentialTransfer');
-    const transferEvent = (await tx
-      .wait()
-      .then(receipt => receipt!.logs.filter((log: any) => log.address === token.target)[0])) as EventLog;
-    expect(transferEvent.args[0]).to.equal(recipient.address);
-    expect(transferEvent.args[1]).to.equal(anyone.address);
+    const [from, to, transferred1] = await callAndGetResult(tx, confidentialTransferEventSignature);
+    expect(from).to.equal(recipient.address);
+    expect(to).to.equal(anyone.address);
     await expect(
-      fhevm.userDecryptEuint(FhevmType.euint64, transferEvent.args[2], await token.getAddress(), recipient),
+      fhevm.userDecryptEuint(FhevmType.euint64, transferred1, await token.getAddress(), recipient),
     ).to.eventually.equal(501);
     await expect(
       fhevm.userDecryptEuint(
@@ -274,7 +273,7 @@ describe('ERC7984Freezable', function () {
     await expect(
       fhevm.userDecryptEuint(
         FhevmType.euint64,
-        await token.confidentialFrozen(recipient),
+        await token.confidentialFrozen(recipient.address),
         await token.getAddress(),
         recipient,
       ),
@@ -286,7 +285,7 @@ describe('ERC7984Freezable', function () {
       .createEncryptedInput(await token.getAddress(), recipient.address)
       .add64(499)
       .encrypt();
-    const [, , transferred] = await callAndGetResult(
+    const [, , transferred2] = await callAndGetResult(
       token
         .connect(recipient)
         ['confidentialTransfer(address,bytes32,bytes)'](
@@ -294,10 +293,10 @@ describe('ERC7984Freezable', function () {
           encryptedInput2.handles[0],
           encryptedInput2.inputProof,
         ),
-      'ConfidentialTransfer(address,address,bytes32)',
+      confidentialTransferEventSignature,
     );
     await expect(
-      fhevm.userDecryptEuint(FhevmType.euint64, transferred, await token.getAddress(), recipient),
+      fhevm.userDecryptEuint(FhevmType.euint64, transferred2, await token.getAddress(), recipient),
     ).to.eventually.equal(0);
   });
 

From 4a2e41ea08f70fe7d24467923d177b0fdafd675c Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 5 Sep 2025 15:40:25 +0200
Subject: [PATCH 31/52] Init compliance modules for confidential RWAs

---
 .changeset/wet-results-doubt.md               |  5 +
 contracts/interfaces/IERC7984Rwa.sol          | 20 +++-
 contracts/mocks/token/ERC7984RwaMock.sol      |  2 +-
 .../token/ERC7984/extensions/ERC7984Rwa.sol   |  4 +-
 .../extensions/ERC7984RwaCompliance.sol       | 98 +++++++++++++++++++
 .../extensions/ERC7984RwaComplianceModule.sol | 10 ++
 6 files changed, 135 insertions(+), 4 deletions(-)
 create mode 100644 .changeset/wet-results-doubt.md
 create mode 100644 contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
 create mode 100644 contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol

diff --git a/.changeset/wet-results-doubt.md b/.changeset/wet-results-doubt.md
new file mode 100644
index 00000000..e4b5af97
--- /dev/null
+++ b/.changeset/wet-results-doubt.md
@@ -0,0 +1,5 @@
+---
+'openzeppelin-confidential-contracts': minor
+---
+
+`ERC7984RwaCompliance`: Support compliance modules for confidential RWAs.
diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index bbf51585..48918955 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -86,5 +86,23 @@ interface IERC7984Rwa is IERC7984, IERC7984RwaBase, IERC165, IAccessControl {}
 /// @dev Interface for confidential RWA compliance.
 interface IERC7984RwaCompliance {
     /// @dev Checks if a transfer follows token compliance.
-    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
+    function isCompliant(address from, address to, euint64 encryptedAmount) external returns (bool);
+}
+
+/// @dev Interface for confidential RWA compliance module.
+interface IERC7984RwaComplianceModule {
+    /// @dev Returns true if module is a certain type, false otherwise.
+    function isModuleType(uint256 moduleTypeId) external returns (bool);
+}
+
+/// @dev Interface for confidential RWA identity compliance module.
+interface IERC7984RwaIdentityComplianceModule is IERC7984RwaComplianceModule {
+    /// @dev Checks if an identity is authorized.
+    function isAuthorizedIdentity(address identity) external returns (bool);
+}
+
+/// @dev Interface for confidential RWA transfer compliance module.
+interface IERC7984RwaTransferComplianceModule is IERC7984RwaComplianceModule {
+    /// @dev Checks if an identity is authorized.
+    function isAuthorizedTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
 }
diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index 0ae0200e..22f3ea98 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -45,7 +45,7 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return _mint(to, FHE.asEuint64(amount));
     }
 
-    function _isCompliantTransfer(
+    function _isCompliant(
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 9d6a1821..6b41175d 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -181,7 +181,7 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) internal override(ERC7984Freezable, ERC7984Restricted, ERC7984) whenNotPaused returns (euint64) {
-        require(_isCompliantTransfer(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+        require(_isCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
         // frozen and restriction checks performed through inheritance
         return super._update(from, to, encryptedAmount);
     }
@@ -193,5 +193,5 @@ abstract contract ERC7984Rwa is
     function _checkFreezer() internal override onlyAdminOrAgent {}
 
     /// @dev Checks if a transfer follows token compliance.
-    function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+    function _isCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
new file mode 100644
index 00000000..84109a15
--- /dev/null
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
@@ -0,0 +1,98 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+import {IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IERC7984RwaTransferComplianceModule} from "./../../../interfaces/IERC7984Rwa.sol";
+import {ERC7984Rwa} from "./ERC7984Rwa.sol";
+
+/**
+ * @dev Extension of {ERC7984Rwa} that supports compliance modules for confidential Real World Assets (RWAs).
+ * Inspired by ERC-7579 modules.
+ */
+abstract contract ERC7984RwaCompliance is ERC7984Rwa {
+    using EnumerableSet for *;
+
+    uint256 constant IDENTITY_COMPLIANCE_MODULE_TYPE = 1;
+    uint256 constant TRANSFER_COMPLIANCE_MODULE_TYPE = 2;
+    EnumerableSet.AddressSet private _identityComplianceModules;
+    EnumerableSet.AddressSet private _transferComplianceModules;
+
+    /// @dev Emitted when a module is installed.
+    event ModuleInstalled(uint256 moduleTypeId, address module);
+    /// @dev Emitted when a module is uninstalled.
+    event ModuleUninstalled(uint256 moduleTypeId, address module);
+
+    /// @dev The module type is not supported.
+    error ERC7984RwaUnsupportedModuleType(uint256 moduleTypeId);
+    /// @dev The provided module doesn't match the provided module type.
+    error ERC7984RwaMismatchedModuleTypeId(uint256 moduleTypeId, address module);
+    /// @dev The module is already installed.
+    error ERC7984RwaAlreadyInstalledModule(uint256 moduleTypeId, address module);
+
+    /**
+     * @dev Check if a certain module typeId is supported.
+     *
+     * Supported module types:
+     *
+     * * Identity compliance moduleÒ
+     * * Transfer compliance module
+     */
+    function supportsModule(uint256 moduleTypeId) public view virtual returns (bool) {
+        return moduleTypeId == IDENTITY_COMPLIANCE_MODULE_TYPE || moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE;
+    }
+
+    function installModule(uint256 moduleTypeId, address module) public virtual onlyAdminOrAgent {
+        _installModule(moduleTypeId, module);
+    }
+
+    function _installModule(uint256 moduleTypeId, address module) internal virtual {
+        require(supportsModule(moduleTypeId), ERC7984RwaUnsupportedModuleType(moduleTypeId));
+        require(
+            IERC7984RwaComplianceModule(module).isModuleType(moduleTypeId),
+            ERC7984RwaMismatchedModuleTypeId(moduleTypeId, module)
+        );
+
+        if (moduleTypeId == IDENTITY_COMPLIANCE_MODULE_TYPE) {
+            require(_identityComplianceModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
+        } else if (moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE) {
+            require(_transferComplianceModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
+        }
+        emit ModuleInstalled(moduleTypeId, module);
+    }
+
+    /// @dev Checks if an identity is authorized.
+    function _isAuthorizedIdentity(address identity) internal virtual returns (bool) {
+        address[] memory modules = _identityComplianceModules.values();
+        uint256 modulesLength = modules.length;
+        for (uint256 index = 0; index < modulesLength; index++) {
+            address module = modules[index];
+            if (!IERC7984RwaIdentityComplianceModule(module).isAuthorizedIdentity(identity)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /// @dev Checks if a transfer is authorized.
+    function _isAuthorizedTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool) {
+        address[] memory modules = _transferComplianceModules.values();
+        uint256 modulesLength = modules.length;
+        for (uint256 index = 0; index < modulesLength; index++) {
+            address module = modules[index];
+            if (!IERC7984RwaTransferComplianceModule(module).isAuthorizedTransfer(from, to, encryptedAmount)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /// @dev Checks if a transfer follows token compliance.
+    function _isCompliant(address from, address to, euint64 encryptedAmount) internal override returns (bool) {
+        return
+            _isAuthorizedIdentity(from) &&
+            _isAuthorizedIdentity(to) &&
+            _isAuthorizedTransfer(from, to, encryptedAmount);
+    }
+}
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol
new file mode 100644
index 00000000..0ee19f30
--- /dev/null
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {IERC7984RwaComplianceModule} from "./../../../interfaces/IERC7984Rwa.sol";
+
+/**
+ * @dev A contract which allows to build a compliance module for confidential Real World Assets (RWAs).
+ */
+abstract contract ERC7984RwaComplianceModule is IERC7984RwaComplianceModule {}

From 3946b30c89958e0c43ce0ac32c7f20638fcc0d04 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 5 Sep 2025 18:09:44 +0200
Subject: [PATCH 32/52] Add abstract compliance modules

---
 contracts/interfaces/IERC7984Rwa.sol          | 15 ++++++---
 contracts/mocks/token/ERC7984RwaMock.sol      | 21 ++++++++++--
 .../token/ERC7984/extensions/ERC7984Rwa.sol   |  6 +++-
 .../extensions/ERC7984RwaCompliance.sol       | 33 +++++++++++--------
 .../extensions/ERC7984RwaComplianceModule.sol | 10 ------
 .../ERC7984RwaIdentityComplianceModule.sol    | 25 ++++++++++++++
 .../ERC7984RwaTransferComplianceModule.sol    | 26 +++++++++++++++
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 28 ++++++++--------
 8 files changed, 119 insertions(+), 45 deletions(-)
 delete mode 100644 contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol
 create mode 100644 contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol
 create mode 100644 contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 48918955..1d31f153 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -7,6 +7,9 @@ import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 import {IERC7984Restricted} from "./IERC7984Restricted.sol";
 
+uint256 constant IDENTITY_COMPLIANCE_MODULE_TYPE = 1;
+uint256 constant TRANSFER_COMPLIANCE_MODULE_TYPE = 2;
+
 /// @dev Base interface for confidential RWA contracts.
 interface IERC7984RwaBase {
     /// @dev Emitted when the contract is paused.
@@ -85,8 +88,10 @@ interface IERC7984Rwa is IERC7984, IERC7984RwaBase, IERC165, IAccessControl {}
 
 /// @dev Interface for confidential RWA compliance.
 interface IERC7984RwaCompliance {
-    /// @dev Checks if a transfer follows token compliance.
+    /// @dev Checks if a transfer follows compliance.
     function isCompliant(address from, address to, euint64 encryptedAmount) external returns (bool);
+    /// @dev Checks if a force transfer follows compliance.
+    function isForceCompliantForce(address from, address to, euint64 encryptedAmount) external returns (bool);
 }
 
 /// @dev Interface for confidential RWA compliance module.
@@ -97,12 +102,12 @@ interface IERC7984RwaComplianceModule {
 
 /// @dev Interface for confidential RWA identity compliance module.
 interface IERC7984RwaIdentityComplianceModule is IERC7984RwaComplianceModule {
-    /// @dev Checks if an identity is authorized.
-    function isAuthorizedIdentity(address identity) external returns (bool);
+    /// @dev Checks if an identity is compliant.
+    function isCompliantIdentity(address identity) external returns (bool);
 }
 
 /// @dev Interface for confidential RWA transfer compliance module.
 interface IERC7984RwaTransferComplianceModule is IERC7984RwaComplianceModule {
-    /// @dev Checks if an identity is authorized.
-    function isAuthorizedTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
+    /// @dev Checks if an transfer is compliant.
+    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
 }
diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index 22f3ea98..c385f78e 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -13,6 +13,7 @@ import {HandleAccessManager} from "../../utils/HandleAccessManager.sol";
 contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
     mapping(address account => euint64 encryptedAmount) private _frozenBalances;
     bool public compliantTransfer;
+    bool public compliantForceTransfer;
 
     // TODO: Move modifiers to `ERC7984Rwa` or remove from mock if useless
     /// @dev Checks if the sender is an admin.
@@ -33,14 +34,22 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         FHE.allow(encryptedAmount, msg.sender);
     }
 
-    function $_setCompliantTransfer() public {
+    function $_setCompliant() public {
         compliantTransfer = true;
     }
 
-    function $_unsetCompliantTransfer() public {
+    function $_unsetCompliant() public {
         compliantTransfer = false;
     }
 
+    function $_setForceCompliant() public {
+        compliantForceTransfer = true;
+    }
+
+    function $_unsetForceCompliant() public {
+        compliantForceTransfer = false;
+    }
+
     function $_mint(address to, uint64 amount) public returns (euint64 transferred) {
         return _mint(to, FHE.asEuint64(amount));
     }
@@ -53,5 +62,13 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return compliantTransfer;
     }
 
+    function _isForceCompliant(
+        address /*from*/,
+        address /*to*/,
+        euint64 /*encryptedAmount*/
+    ) internal override returns (bool) {
+        return compliantForceTransfer;
+    }
+
     function _validateHandleAllowance(bytes32 handle) internal view override onlyAdminOrAgent {}
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 6b41175d..a8a28a31 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -167,6 +167,7 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) internal virtual returns (euint64 transferred) {
+        require(_isForceCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
         _disableERC7984FreezableUpdateCheck(); // bypass frozen check
         _disableERC7984RestrictedUpdateCheck(); // bypass default restriction check
         if (to != address(0)) _checkRestriction(to); // only perform restriction check on `to`
@@ -192,6 +193,9 @@ abstract contract ERC7984Rwa is
      */
     function _checkFreezer() internal override onlyAdminOrAgent {}
 
-    /// @dev Checks if a transfer follows token compliance.
+    /// @dev Checks if a transfer follows compliance.
     function _isCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+
+    /// @dev Checks if a force transfer follows compliance.
+    function _isForceCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
index 84109a15..dadfd14c 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IERC7984RwaTransferComplianceModule} from "./../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IERC7984RwaTransferComplianceModule, IDENTITY_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984Rwa} from "./ERC7984Rwa.sol";
 
 /**
@@ -14,8 +14,6 @@ import {ERC7984Rwa} from "./ERC7984Rwa.sol";
 abstract contract ERC7984RwaCompliance is ERC7984Rwa {
     using EnumerableSet for *;
 
-    uint256 constant IDENTITY_COMPLIANCE_MODULE_TYPE = 1;
-    uint256 constant TRANSFER_COMPLIANCE_MODULE_TYPE = 2;
     EnumerableSet.AddressSet private _identityComplianceModules;
     EnumerableSet.AddressSet private _transferComplianceModules;
 
@@ -36,7 +34,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa {
      *
      * Supported module types:
      *
-     * * Identity compliance moduleÒ
+     * * Identity compliance module
      * * Transfer compliance module
      */
     function supportsModule(uint256 moduleTypeId) public view virtual returns (bool) {
@@ -62,37 +60,44 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa {
         emit ModuleInstalled(moduleTypeId, module);
     }
 
-    /// @dev Checks if an identity is authorized.
-    function _isAuthorizedIdentity(address identity) internal virtual returns (bool) {
+    /// @dev Checks if an identity is compliant.
+    function _isCompliantIdentity(address identity) internal virtual returns (bool) {
         address[] memory modules = _identityComplianceModules.values();
         uint256 modulesLength = modules.length;
         for (uint256 index = 0; index < modulesLength; index++) {
             address module = modules[index];
-            if (!IERC7984RwaIdentityComplianceModule(module).isAuthorizedIdentity(identity)) {
+            if (!IERC7984RwaIdentityComplianceModule(module).isCompliantIdentity(identity)) {
                 return false;
             }
         }
         return true;
     }
 
-    /// @dev Checks if a transfer is authorized.
-    function _isAuthorizedTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool) {
+    /// @dev Checks if a transfer is compliant.
+    function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool) {
         address[] memory modules = _transferComplianceModules.values();
         uint256 modulesLength = modules.length;
         for (uint256 index = 0; index < modulesLength; index++) {
             address module = modules[index];
-            if (!IERC7984RwaTransferComplianceModule(module).isAuthorizedTransfer(from, to, encryptedAmount)) {
+            if (!IERC7984RwaTransferComplianceModule(module).isCompliantTransfer(from, to, encryptedAmount)) {
                 return false;
             }
         }
         return true;
     }
 
-    /// @dev Checks if a transfer follows token compliance.
+    /// @dev Checks if a transfer follows compliance.
     function _isCompliant(address from, address to, euint64 encryptedAmount) internal override returns (bool) {
         return
-            _isAuthorizedIdentity(from) &&
-            _isAuthorizedIdentity(to) &&
-            _isAuthorizedTransfer(from, to, encryptedAmount);
+            _isCompliantIdentity(from) && _isCompliantIdentity(to) && _isCompliantTransfer(from, to, encryptedAmount);
+    }
+
+    /// @dev Checks if a force transfer follows compliance.
+    function _isForceCompliant(
+        address /*from*/,
+        address to,
+        euint64 /*encryptedAmount*/
+    ) internal override returns (bool) {
+        return _isCompliantIdentity(to);
     }
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol
deleted file mode 100644
index 0ee19f30..00000000
--- a/contracts/token/ERC7984/extensions/ERC7984RwaComplianceModule.sol
+++ /dev/null
@@ -1,10 +0,0 @@
-// SPDX-License-Identifier: MIT
-
-pragma solidity ^0.8.27;
-
-import {IERC7984RwaComplianceModule} from "./../../../interfaces/IERC7984Rwa.sol";
-
-/**
- * @dev A contract which allows to build a compliance module for confidential Real World Assets (RWAs).
- */
-abstract contract ERC7984RwaComplianceModule is IERC7984RwaComplianceModule {}
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol
new file mode 100644
index 00000000..7c1e9abc
--- /dev/null
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IDENTITY_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
+
+/**
+ * @dev A contract which allows to build an identity compliance module for confidential Real World Assets (RWAs).
+ */
+abstract contract ERC7984RwaIdentityComplianceModule is
+    IERC7984RwaComplianceModule,
+    IERC7984RwaIdentityComplianceModule
+{
+    /// @inheritdoc IERC7984RwaComplianceModule
+    function isModuleType(uint256 moduleTypeId) public pure override returns (bool) {
+        return moduleTypeId == IDENTITY_COMPLIANCE_MODULE_TYPE;
+    }
+
+    /// @inheritdoc IERC7984RwaIdentityComplianceModule
+    function isCompliantIdentity(address identity) public virtual returns (bool) {
+        return _isCompliantIdentity(identity);
+    }
+
+    function _isCompliantIdentity(address identity) internal virtual returns (bool);
+}
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
new file mode 100644
index 00000000..ba8b0283
--- /dev/null
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {IERC7984RwaComplianceModule, IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
+
+/**
+ * @dev A contract which allows to build an transfer compliance module for confidential Real World Assets (RWAs).
+ */
+abstract contract ERC7984RwaTransferComplianceModule is
+    IERC7984RwaComplianceModule,
+    IERC7984RwaTransferComplianceModule
+{
+    /// @inheritdoc IERC7984RwaComplianceModule
+    function isModuleType(uint256 moduleTypeId) public pure override returns (bool) {
+        return moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE;
+    }
+
+    /// @inheritdoc IERC7984RwaTransferComplianceModule
+    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) public virtual returns (bool) {
+        return _isCompliantTransfer(from, to, encryptedAmount);
+    }
+
+    function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+}
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index ae9de6fb..3f3ef38d 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -131,7 +131,7 @@ describe('ERC7984Rwa', function () {
         const { admin, agent1, recipient } = await fixture();
         for (const manager of [admin, agent1]) {
           const { token } = await fixture();
-          await token.$_setCompliantTransfer();
+          await token.$_setCompliant();
           const amount = 100;
           let params = [recipient.address] as unknown as [
             account: AddressLike,
@@ -172,7 +172,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliantTransfer();
+      await token.$_setCompliant();
       await expect(
         token
           .connect(anyone)
@@ -222,7 +222,7 @@ describe('ERC7984Rwa', function () {
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
             .encrypt();
-          await token.$_setCompliantTransfer();
+          await token.$_setCompliant();
           await token
             .connect(manager)
             ['confidentialMint(address,bytes32,bytes)'](
@@ -275,7 +275,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliantTransfer();
+      await token.$_setCompliant();
       await expect(
         token
           .connect(anyone)
@@ -325,7 +325,7 @@ describe('ERC7984Rwa', function () {
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
             .encrypt();
-          await token.$_setCompliantTransfer();
+          await token.$_setCompliant();
           await token
             .connect(manager)
             ['confidentialMint(address,bytes32,bytes)'](
@@ -345,7 +345,7 @@ describe('ERC7984Rwa', function () {
               encryptedFrozenValueInput.handles[0],
               encryptedFrozenValueInput.inputProof,
             );
-          await token.$_unsetCompliantTransfer();
+          await token.$_unsetCompliant();
           expect(await token.compliantTransfer()).to.be.false;
           const amount = 25;
           let params = [recipient.address, anyone.address] as unknown as [
@@ -364,6 +364,7 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
+          await token.$_setForceCompliant();
           const [from, to, transferredHandle] = await callAndGetResult(
             token
               .connect(manager)
@@ -402,7 +403,7 @@ describe('ERC7984Rwa', function () {
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
             .encrypt();
-          await token.$_setCompliantTransfer();
+          await token.$_setCompliant();
           await token
             .connect(manager)
             ['confidentialMint(address,bytes32,bytes)'](
@@ -423,7 +424,7 @@ describe('ERC7984Rwa', function () {
               encryptedFrozenValueInput.inputProof,
             );
           // should force transfer even if not compliant
-          await token.$_unsetCompliantTransfer();
+          await token.$_unsetCompliant();
           expect(await token.compliantTransfer()).to.be.false;
           // should force transfer even if paused
           await token.connect(manager).pause();
@@ -445,6 +446,7 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
+          await token.$_setForceCompliant();
           const [account, frozenAmountHandle] = await callAndGetResult(
             token
               .connect(manager)
@@ -550,7 +552,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliantTransfer();
+      await token.$_setCompliant();
       await token
         .connect(manager)
         ['confidentialMint(address,bytes32,bytes)'](
@@ -575,7 +577,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(amount)
         .encrypt();
-      await token.$_setCompliantTransfer();
+      await token.$_setCompliant();
       expect(await token.compliantTransfer()).to.be.true;
       const [from, to, transferredHandle] = await callAndGetResult(
         token
@@ -654,7 +656,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliantTransfer();
+      await token.$_setCompliant();
       await token
         .connect(manager)
         ['confidentialMint(address,bytes32,bytes)'](
@@ -678,7 +680,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      await token.$_setCompliantTransfer();
+      await token.$_setCompliant();
       expect(await token.compliantTransfer()).to.be.true;
       const [, , transferredHandle] = await callAndGetResult(
         token
@@ -708,7 +710,7 @@ describe('ERC7984Rwa', function () {
       it(`should not transfer if ${arg ? 'sender' : 'receiver'} blocked `, async function () {
         const { token, admin: manager, recipient, anyone } = await fixture();
         const account = arg ? recipient : anyone;
-        await token.$_setCompliantTransfer();
+        await token.$_setCompliant();
         const encryptedInput = await fhevm
           .createEncryptedInput(await token.getAddress(), recipient.address)
           .add64(25)

From 4debb472ba28db8d5145f3f87d13b1a123fc173b Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Fri, 5 Sep 2025 18:19:34 +0200
Subject: [PATCH 33/52] Compliance implements interface

---
 contracts/interfaces/IERC7984Rwa.sol                        | 2 +-
 contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 1d31f153..f0363a62 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -91,7 +91,7 @@ interface IERC7984RwaCompliance {
     /// @dev Checks if a transfer follows compliance.
     function isCompliant(address from, address to, euint64 encryptedAmount) external returns (bool);
     /// @dev Checks if a force transfer follows compliance.
-    function isForceCompliantForce(address from, address to, euint64 encryptedAmount) external returns (bool);
+    function isForceCompliant(address from, address to, euint64 encryptedAmount) external returns (bool);
 }
 
 /// @dev Interface for confidential RWA compliance module.
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
index dadfd14c..24f9ab46 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
@@ -4,14 +4,14 @@ pragma solidity ^0.8.27;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IERC7984RwaTransferComplianceModule, IDENTITY_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaCompliance, IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IERC7984RwaTransferComplianceModule, IDENTITY_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984Rwa} from "./ERC7984Rwa.sol";
 
 /**
  * @dev Extension of {ERC7984Rwa} that supports compliance modules for confidential Real World Assets (RWAs).
  * Inspired by ERC-7579 modules.
  */
-abstract contract ERC7984RwaCompliance is ERC7984Rwa {
+abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
     using EnumerableSet for *;
 
     EnumerableSet.AddressSet private _identityComplianceModules;

From 86d5250216bd1617376026f80a5e869eb4562b7a Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Mon, 8 Sep 2025 16:15:17 +0200
Subject: [PATCH 34/52] Add post transfer hook

---
 contracts/interfaces/IERC7984Rwa.sol          |  32 ++---
 contracts/mocks/token/ERC7984RwaMock.sol      |   4 +-
 .../token/ERC7984/extensions/ERC7984Rwa.sol   |  24 +++-
 .../extensions/ERC7984RwaCompliance.sol       | 111 ++++++++++++------
 .../ERC7984RwaIdentityComplianceModule.sol    |  25 ----
 .../ERC7984RwaTransferComplianceModule.sol    |  35 ++++--
 6 files changed, 134 insertions(+), 97 deletions(-)
 delete mode 100644 contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index f0363a62..f0c06628 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -7,8 +7,8 @@ import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 import {IERC7984Restricted} from "./IERC7984Restricted.sol";
 
-uint256 constant IDENTITY_COMPLIANCE_MODULE_TYPE = 1;
-uint256 constant TRANSFER_COMPLIANCE_MODULE_TYPE = 2;
+uint256 constant TRANSFER_COMPLIANCE_MODULE_TYPE = 1;
+uint256 constant FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE = 2;
 
 /// @dev Base interface for confidential RWA contracts.
 interface IERC7984RwaBase {
@@ -88,26 +88,18 @@ interface IERC7984Rwa is IERC7984, IERC7984RwaBase, IERC165, IAccessControl {}
 
 /// @dev Interface for confidential RWA compliance.
 interface IERC7984RwaCompliance {
-    /// @dev Checks if a transfer follows compliance.
-    function isCompliant(address from, address to, euint64 encryptedAmount) external returns (bool);
-    /// @dev Checks if a force transfer follows compliance.
-    function isForceCompliant(address from, address to, euint64 encryptedAmount) external returns (bool);
-}
-
-/// @dev Interface for confidential RWA compliance module.
-interface IERC7984RwaComplianceModule {
-    /// @dev Returns true if module is a certain type, false otherwise.
-    function isModuleType(uint256 moduleTypeId) external returns (bool);
-}
-
-/// @dev Interface for confidential RWA identity compliance module.
-interface IERC7984RwaIdentityComplianceModule is IERC7984RwaComplianceModule {
-    /// @dev Checks if an identity is compliant.
-    function isCompliantIdentity(address identity) external returns (bool);
+    /// @dev Installs a transfer compliance module.
+    function installModule(uint256 moduleTypeId, address module) external;
+    /// @dev Uninstalls a transfer compliance module.
+    function uninstallModule(uint256 moduleTypeId, address module) external;
 }
 
 /// @dev Interface for confidential RWA transfer compliance module.
-interface IERC7984RwaTransferComplianceModule is IERC7984RwaComplianceModule {
-    /// @dev Checks if an transfer is compliant.
+interface IERC7984RwaTransferComplianceModule {
+    /// @dev Returns magic number if it is a module.
+    function isModule() external returns (bytes4);
+    /// @dev Checks if a transfer is compliant. Should be non-mutating.
     function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
+    /// @dev Peforms operation after transfer.
+    function postTransferHook(address from, address to, euint64 encryptedAmount) external;
 }
diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index c385f78e..f3c93708 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -54,7 +54,7 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return _mint(to, FHE.asEuint64(amount));
     }
 
-    function _isCompliant(
+    function _isTransferCompliant(
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
@@ -62,7 +62,7 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return compliantTransfer;
     }
 
-    function _isForceCompliant(
+    function _isForceTransferCompliant(
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index a8a28a31..84cd52b1 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -167,11 +167,12 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) internal virtual returns (euint64 transferred) {
-        require(_isForceCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+        require(_isForceTransferCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
         _disableERC7984FreezableUpdateCheck(); // bypass frozen check
         _disableERC7984RestrictedUpdateCheck(); // bypass default restriction check
         if (to != address(0)) _checkRestriction(to); // only perform restriction check on `to`
         transferred = super._update(from, to, encryptedAmount); // bypass compliance check
+        _postForceTransferHook(from, to, encryptedAmount);
         _restoreERC7984FreezableUpdateCheck();
         _restoreERC7984RestrictedUpdateCheck();
     }
@@ -181,10 +182,11 @@ abstract contract ERC7984Rwa is
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal override(ERC7984Freezable, ERC7984Restricted, ERC7984) whenNotPaused returns (euint64) {
-        require(_isCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+    ) internal override(ERC7984Freezable, ERC7984Restricted, ERC7984) whenNotPaused returns (euint64 transferred) {
+        require(_isTransferCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
         // frozen and restriction checks performed through inheritance
-        return super._update(from, to, encryptedAmount);
+        transferred = super._update(from, to, encryptedAmount);
+        _postTransferHook(from, to, encryptedAmount);
     }
 
     /**
@@ -194,8 +196,18 @@ abstract contract ERC7984Rwa is
     function _checkFreezer() internal override onlyAdminOrAgent {}
 
     /// @dev Checks if a transfer follows compliance.
-    function _isCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+    function _isTransferCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
 
     /// @dev Checks if a force transfer follows compliance.
-    function _isForceCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+    function _isForceTransferCompliant(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) internal virtual returns (bool);
+
+    /// @dev Peforms operation after transfer.
+    function _postTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
+
+    /// @dev Peforms operation after force transfer.
+    function _postForceTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
index 24f9ab46..db779b0e 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaCompliance, IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IERC7984RwaTransferComplianceModule, IDENTITY_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984Rwa} from "./ERC7984Rwa.sol";
 
 /**
@@ -14,8 +14,8 @@ import {ERC7984Rwa} from "./ERC7984Rwa.sol";
 abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
     using EnumerableSet for *;
 
-    EnumerableSet.AddressSet private _identityComplianceModules;
     EnumerableSet.AddressSet private _transferComplianceModules;
+    EnumerableSet.AddressSet private _forceTransferComplianceModules;
 
     /// @dev Emitted when a module is installed.
     event ModuleInstalled(uint256 moduleTypeId, address module);
@@ -24,80 +24,123 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
 
     /// @dev The module type is not supported.
     error ERC7984RwaUnsupportedModuleType(uint256 moduleTypeId);
-    /// @dev The provided module doesn't match the provided module type.
-    error ERC7984RwaMismatchedModuleTypeId(uint256 moduleTypeId, address module);
+    /// @dev The address is not a transfer compliance module.
+    error ERC7984RwaNotTransferComplianceModule(address module);
     /// @dev The module is already installed.
     error ERC7984RwaAlreadyInstalledModule(uint256 moduleTypeId, address module);
+    /// @dev The module is already uninstalled.
+    error ERC7984RwaAlreadyUninstalledModule(uint256 moduleTypeId, address module);
 
     /**
      * @dev Check if a certain module typeId is supported.
      *
      * Supported module types:
      *
-     * * Identity compliance module
      * * Transfer compliance module
+     * * Force transfer compliance module
      */
     function supportsModule(uint256 moduleTypeId) public view virtual returns (bool) {
-        return moduleTypeId == IDENTITY_COMPLIANCE_MODULE_TYPE || moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE;
+        return moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE || moduleTypeId == FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE;
     }
 
+    /**
+     * @inheritdoc IERC7984RwaCompliance
+     * @dev Consider gas footprint of the module before adding it since all modules will perform
+     * all steps (pre-check, compliance check, post-hook) in a single transaction.
+     */
     function installModule(uint256 moduleTypeId, address module) public virtual onlyAdminOrAgent {
         _installModule(moduleTypeId, module);
     }
 
+    /// @inheritdoc IERC7984RwaCompliance
+    function uninstallModule(uint256 moduleTypeId, address module) public virtual onlyAdminOrAgent {
+        _uninstallModule(moduleTypeId, module);
+    }
+
+    /// @dev Internal function which installs a transfer compliance module.
     function _installModule(uint256 moduleTypeId, address module) internal virtual {
         require(supportsModule(moduleTypeId), ERC7984RwaUnsupportedModuleType(moduleTypeId));
         require(
-            IERC7984RwaComplianceModule(module).isModuleType(moduleTypeId),
-            ERC7984RwaMismatchedModuleTypeId(moduleTypeId, module)
+            IERC7984RwaTransferComplianceModule(module).isModule() ==
+                IERC7984RwaTransferComplianceModule.isModule.selector,
+            ERC7984RwaNotTransferComplianceModule(module)
         );
 
-        if (moduleTypeId == IDENTITY_COMPLIANCE_MODULE_TYPE) {
-            require(_identityComplianceModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
-        } else if (moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE) {
+        if (moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE) {
             require(_transferComplianceModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
+        } else if (moduleTypeId == FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE) {
+            require(
+                _forceTransferComplianceModules.add(module),
+                ERC7984RwaAlreadyInstalledModule(moduleTypeId, module)
+            );
         }
         emit ModuleInstalled(moduleTypeId, module);
     }
 
-    /// @dev Checks if an identity is compliant.
-    function _isCompliantIdentity(address identity) internal virtual returns (bool) {
-        address[] memory modules = _identityComplianceModules.values();
+    /// @dev Internal function which uninstalls a transfer compliance module.
+    function _uninstallModule(uint256 moduleTypeId, address module) internal virtual {
+        require(supportsModule(moduleTypeId), ERC7984RwaUnsupportedModuleType(moduleTypeId));
+        if (moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE) {
+            require(
+                _transferComplianceModules.remove(module),
+                ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module)
+            );
+        } else if (moduleTypeId == FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE) {
+            require(
+                _forceTransferComplianceModules.remove(module),
+                ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module)
+            );
+        }
+        emit ModuleUninstalled(moduleTypeId, module);
+    }
+
+    /// @dev Checks if a transfer is compliant.
+    function _isTransferCompliantTransfer(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) internal virtual returns (bool) {
+        address[] memory modules = _transferComplianceModules.values();
         uint256 modulesLength = modules.length;
-        for (uint256 index = 0; index < modulesLength; index++) {
-            address module = modules[index];
-            if (!IERC7984RwaIdentityComplianceModule(module).isCompliantIdentity(identity)) {
+        for (uint256 i = 0; i < modulesLength; i++) {
+            if (!IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)) {
                 return false;
             }
         }
         return true;
     }
 
-    /// @dev Checks if a transfer is compliant.
-    function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool) {
-        address[] memory modules = _transferComplianceModules.values();
+    /// @dev Checks if a force transfer is compliant.
+    function _isTransferCompliantForceTransfer(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) internal virtual returns (bool) {
+        address[] memory modules = _forceTransferComplianceModules.values();
         uint256 modulesLength = modules.length;
-        for (uint256 index = 0; index < modulesLength; index++) {
-            address module = modules[index];
-            if (!IERC7984RwaTransferComplianceModule(module).isCompliantTransfer(from, to, encryptedAmount)) {
+        for (uint256 i = 0; i < modulesLength; i++) {
+            if (!IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)) {
                 return false;
             }
         }
         return true;
     }
 
-    /// @dev Checks if a transfer follows compliance.
-    function _isCompliant(address from, address to, euint64 encryptedAmount) internal override returns (bool) {
-        return
-            _isCompliantIdentity(from) && _isCompliantIdentity(to) && _isCompliantTransfer(from, to, encryptedAmount);
+    /// @dev Peforms operation after transfer.
+    function _postTransferHook(address from, address to, euint64 encryptedAmount) internal override {
+        address[] memory modules = _transferComplianceModules.values();
+        uint256 modulesLength = modules.length;
+        for (uint256 i = 0; i < modulesLength; i++) {
+            IERC7984RwaTransferComplianceModule(modules[i]).postTransferHook(from, to, encryptedAmount);
+        }
     }
 
-    /// @dev Checks if a force transfer follows compliance.
-    function _isForceCompliant(
-        address /*from*/,
-        address to,
-        euint64 /*encryptedAmount*/
-    ) internal override returns (bool) {
-        return _isCompliantIdentity(to);
+    /// @dev Peforms operation after force transfer.
+    function _postForceTransferHook(address from, address to, euint64 encryptedAmount) internal override {
+        address[] memory modules = _forceTransferComplianceModules.values();
+        uint256 modulesLength = modules.length;
+        for (uint256 i = 0; i < modulesLength; i++) {
+            IERC7984RwaTransferComplianceModule(modules[i]).postTransferHook(from, to, encryptedAmount);
+        }
     }
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol
deleted file mode 100644
index 7c1e9abc..00000000
--- a/contracts/token/ERC7984/extensions/ERC7984RwaIdentityComplianceModule.sol
+++ /dev/null
@@ -1,25 +0,0 @@
-// SPDX-License-Identifier: MIT
-
-pragma solidity ^0.8.27;
-
-import {IERC7984RwaComplianceModule, IERC7984RwaIdentityComplianceModule, IDENTITY_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
-
-/**
- * @dev A contract which allows to build an identity compliance module for confidential Real World Assets (RWAs).
- */
-abstract contract ERC7984RwaIdentityComplianceModule is
-    IERC7984RwaComplianceModule,
-    IERC7984RwaIdentityComplianceModule
-{
-    /// @inheritdoc IERC7984RwaComplianceModule
-    function isModuleType(uint256 moduleTypeId) public pure override returns (bool) {
-        return moduleTypeId == IDENTITY_COMPLIANCE_MODULE_TYPE;
-    }
-
-    /// @inheritdoc IERC7984RwaIdentityComplianceModule
-    function isCompliantIdentity(address identity) public virtual returns (bool) {
-        return _isCompliantIdentity(identity);
-    }
-
-    function _isCompliantIdentity(address identity) internal virtual returns (bool);
-}
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
index ba8b0283..a245fc08 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
@@ -3,18 +3,15 @@
 pragma solidity ^0.8.27;
 
 import {euint64} from "@fhevm/solidity/lib/FHE.sol";
-import {IERC7984RwaComplianceModule, IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
 
 /**
- * @dev A contract which allows to build an transfer compliance module for confidential Real World Assets (RWAs).
+ * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).
  */
-abstract contract ERC7984RwaTransferComplianceModule is
-    IERC7984RwaComplianceModule,
-    IERC7984RwaTransferComplianceModule
-{
-    /// @inheritdoc IERC7984RwaComplianceModule
-    function isModuleType(uint256 moduleTypeId) public pure override returns (bool) {
-        return moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE;
+abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule {
+    /// @inheritdoc IERC7984RwaTransferComplianceModule
+    function isModule() public pure override returns (bytes4) {
+        return this.isModule.selector;
     }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
@@ -22,5 +19,23 @@ abstract contract ERC7984RwaTransferComplianceModule is
         return _isCompliantTransfer(from, to, encryptedAmount);
     }
 
-    function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+    /// @inheritdoc IERC7984RwaTransferComplianceModule
+    function postTransferHook(address from, address to, euint64 encryptedAmount) public virtual {
+        _postTransferHook(from, to, encryptedAmount);
+    }
+
+    /// @dev Internal function which checks if a transfer is compliant.
+    function _isCompliantTransfer(
+        address /*from*/,
+        address /*to*/,
+        euint64 /*encryptedAmount*/
+    ) internal virtual returns (bool) {
+        // default to non-compliant
+        return false;
+    }
+
+    /// @dev Internal function which peforms operation after transfer.
+    function _postTransferHook(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal virtual {
+        // default to no-op
+    }
 }

From 30ca7dfe897662dd7348dec2dc5fe785b5d1d71a Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Mon, 8 Sep 2025 16:35:44 +0200
Subject: [PATCH 35/52] Typo

---
 contracts/interfaces/IERC7984Rwa.sol                          | 2 +-
 contracts/token/ERC7984/extensions/ERC7984Rwa.sol             | 4 ++--
 contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol   | 4 ++--
 .../ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index f0c06628..2ce4b4ef 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -100,6 +100,6 @@ interface IERC7984RwaTransferComplianceModule {
     function isModule() external returns (bytes4);
     /// @dev Checks if a transfer is compliant. Should be non-mutating.
     function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
-    /// @dev Peforms operation after transfer.
+    /// @dev Performs operation after transfer.
     function postTransferHook(address from, address to, euint64 encryptedAmount) external;
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 84cd52b1..cddead24 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -205,9 +205,9 @@ abstract contract ERC7984Rwa is
         euint64 encryptedAmount
     ) internal virtual returns (bool);
 
-    /// @dev Peforms operation after transfer.
+    /// @dev Performs operation after transfer.
     function _postTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
 
-    /// @dev Peforms operation after force transfer.
+    /// @dev Performs operation after force transfer.
     function _postForceTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
index db779b0e..d97b8b6b 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
@@ -126,7 +126,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         return true;
     }
 
-    /// @dev Peforms operation after transfer.
+    /// @dev Performs operation after transfer.
     function _postTransferHook(address from, address to, euint64 encryptedAmount) internal override {
         address[] memory modules = _transferComplianceModules.values();
         uint256 modulesLength = modules.length;
@@ -135,7 +135,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         }
     }
 
-    /// @dev Peforms operation after force transfer.
+    /// @dev Performs operation after force transfer.
     function _postForceTransferHook(address from, address to, euint64 encryptedAmount) internal override {
         address[] memory modules = _forceTransferComplianceModules.values();
         uint256 modulesLength = modules.length;
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
index a245fc08..24a0b0a6 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
@@ -34,7 +34,7 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
         return false;
     }
 
-    /// @dev Internal function which peforms operation after transfer.
+    /// @dev Internal function which Performs operation after transfer.
     function _postTransferHook(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal virtual {
         // default to no-op
     }

From 57ab5005c371f9867ea14fe4dd37937d1ca78d67 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 9 Sep 2025 13:02:50 +0200
Subject: [PATCH 36/52] Init investor cap module

---
 .../ERC7984RwaInvestorCapModule.sol           | 54 +++++++++++++++++++
 .../ERC7984RwaTransferComplianceModule.sol    | 11 +++-
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol

diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol
new file mode 100644
index 00000000..4e804444
--- /dev/null
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol
@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+import {ERC7984RwaTransferComplianceModule} from "../ERC7984RwaTransferComplianceModule.sol";
+
+/**
+ * @dev A transfer compliance module for confidential Real World Assets (RWAs) which limits the number of investors.
+ */
+abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceModule {
+    using EnumerableSet for *;
+
+    uint256 private _maxInvestor;
+    EnumerableSet.AddressSet private _investors;
+
+    constructor(address compliance, uint256 maxInvestor) ERC7984RwaTransferComplianceModule(compliance) {
+        setMaxInvestor(maxInvestor);
+    }
+
+    /// @dev Sets max number of investors.
+    function setMaxInvestor(uint256 maxInvestor) public virtual onlyCompliance {
+        _maxInvestor = maxInvestor;
+    }
+
+    /// @dev Gets max number of investors.
+    function getMaxInvestor() public virtual returns (uint256) {
+        return _maxInvestor;
+    }
+
+    /// @dev Internal function which checks if a transfer is compliant.
+    function _isCompliantTransfer(
+        address /*from*/,
+        address to,
+        euint64 /*encryptedAmount*/
+    ) internal override returns (bool) {
+        if (
+            to == address(0) || // burning
+            _investors.contains(to) || // or already investor
+            _investors.length() < _maxInvestor // or not reached max investors limit
+        ) {
+            return true;
+        }
+        return false;
+    }
+
+    /// @dev Internal function which Performs operation after transfer.
+    function _postTransferHook(address /*from*/, address to, euint64 /*encryptedAmount*/) internal override {
+        if (!_investors.contains(to)) {
+            _investors.add(to);
+        }
+    }
+}
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
index 24a0b0a6..120a4e0c 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
@@ -3,12 +3,21 @@
 pragma solidity ^0.8.27;
 
 import {euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 import {IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
 
 /**
  * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).
  */
-abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule {
+abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule, Ownable {
+    /// @dev Throws if called by any account other than the compliance.
+    modifier onlyCompliance() {
+        _checkOwner();
+        _;
+    }
+
+    constructor(address compliance) Ownable(compliance) {}
+
     /// @inheritdoc IERC7984RwaTransferComplianceModule
     function isModule() public pure override returns (bytes4) {
         return this.isModule.selector;

From ed06057c48c0d4541b1a368662edf418f81d9023 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 9 Sep 2025 14:27:46 +0200
Subject: [PATCH 37/52] Move rwa compliance contracts

---
 .../rwa/ERC7984RwaBalanceCapModule.sol        | 53 +++++++++++++++++++
 .../{ => rwa}/ERC7984RwaCompliance.sol        |  4 +-
 .../ERC7984RwaInvestorCapModule.sol           |  2 +-
 .../ERC7984RwaTransferComplianceModule.sol    |  2 +-
 4 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
 rename contracts/token/ERC7984/extensions/{ => rwa}/ERC7984RwaCompliance.sol (98%)
 rename contracts/token/ERC7984/extensions/{ERC7984Rwa => rwa}/ERC7984RwaInvestorCapModule.sol (94%)
 rename contracts/token/ERC7984/extensions/{ => rwa}/ERC7984RwaTransferComplianceModule.sol (96%)

diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
new file mode 100644
index 00000000..18f7fe62
--- /dev/null
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+import {IERC7984} from "./../../../../interfaces/IERC7984.sol";
+import {FHESafeMath} from "../../../../utils/FHESafeMath.sol";
+import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
+
+/**
+ * @dev A transfer compliance module for confidential Real World Assets (RWAs) which limits the balance of an investor.
+ */
+abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModule {
+    using EnumerableSet for *;
+
+    euint64 private _maxBalance;
+    address private _token;
+
+    constructor(address compliance, euint64 maxBalance) ERC7984RwaTransferComplianceModule(compliance) {
+        setMaxBalance(maxBalance);
+    }
+
+    /// @dev Sets max balance of an investor.
+    function setMaxBalance(euint64 maxBalance) public virtual onlyCompliance {
+        _maxBalance = maxBalance;
+    }
+
+    /// @dev Gets max balance of an investor.
+    function getMaxBalance() public virtual returns (euint64) {
+        return _maxBalance;
+    }
+
+    /// @dev Internal function which checks if a transfer is compliant.
+    function _isCompliantTransfer(
+        address /*from*/,
+        address to,
+        euint64 encryptedAmount
+    ) internal override returns (bool) {
+        if (
+            to == address(0) // burning
+        ) {
+            return true;
+        }
+        (ebool increased, euint64 futureBalance) = FHESafeMath.tryIncrease(
+            IERC7984(_token).confidentialBalanceOf(to),
+            encryptedAmount
+        );
+        ebool isCompliant = FHE.and(increased, FHE.le(futureBalance, _maxBalance));
+        isCompliant;
+        return false;
+    }
+}
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
similarity index 98%
rename from contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
rename to contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
index d97b8b6b..0459550f 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
@@ -4,8 +4,8 @@ pragma solidity ^0.8.27;
 
 import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
-import {ERC7984Rwa} from "./ERC7984Rwa.sol";
+import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
+import {ERC7984Rwa} from "../ERC7984Rwa.sol";
 
 /**
  * @dev Extension of {ERC7984Rwa} that supports compliance modules for confidential Real World Assets (RWAs).
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
similarity index 94%
rename from contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol
rename to contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index 4e804444..48c7fc00 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {ERC7984RwaTransferComplianceModule} from "../ERC7984RwaTransferComplianceModule.sol";
+import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
 
 /**
  * @dev A transfer compliance module for confidential Real World Assets (RWAs) which limits the number of investors.
diff --git a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
similarity index 96%
rename from contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
rename to contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
index 120a4e0c..a5daf4a3 100644
--- a/contracts/token/ERC7984/extensions/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
-import {IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
 
 /**
  * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).

From 4b4d558a819c300306bdfde417f292cd391ff175 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 10 Sep 2025 10:45:10 +0200
Subject: [PATCH 38/52] Support confidential rwa module

---
 contracts/interfaces/IERC7984Rwa.sol          |  4 +-
 contracts/mocks/token/ERC7984RwaMock.sol      | 16 +++---
 .../ERC7984/extensions/ERC7984Freezable.sol   |  3 ++
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 26 +++++++---
 .../rwa/ERC7984RwaBalanceCapModule.sol        | 10 ++--
 .../extensions/rwa/ERC7984RwaCompliance.sol   | 30 +++++++-----
 .../rwa/ERC7984RwaInvestorCapModule.sol       | 13 ++---
 .../ERC7984RwaTransferComplianceModule.sol    |  8 +--
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 49 ++++++++++++-------
 9 files changed, 98 insertions(+), 61 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 2ce4b4ef..fa3767e8 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
-import {externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {ebool, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
 import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
@@ -99,7 +99,7 @@ interface IERC7984RwaTransferComplianceModule {
     /// @dev Returns magic number if it is a module.
     function isModule() external returns (bytes4);
     /// @dev Checks if a transfer is compliant. Should be non-mutating.
-    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (bool);
+    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (ebool);
     /// @dev Performs operation after transfer.
     function postTransferHook(address from, address to, euint64 encryptedAmount) external;
 }
diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index f3c93708..4ea777c7 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -3,7 +3,7 @@
 pragma solidity ^0.8.24;
 
 import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
-import {FHE, euint64, externalEuint64, ebool} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
 import {Impl} from "@fhevm/solidity/lib/Impl.sol";
 import {ERC7984Rwa} from "../../token/ERC7984/extensions/ERC7984Rwa.sol";
 import {FHESafeMath} from "../../utils/FHESafeMath.sol";
@@ -12,8 +12,8 @@ import {HandleAccessManager} from "../../utils/HandleAccessManager.sol";
 // solhint-disable func-name-mixedcase
 contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
     mapping(address account => euint64 encryptedAmount) private _frozenBalances;
-    bool public compliantTransfer;
-    bool public compliantForceTransfer;
+    bool public compliantTransfer = false;
+    bool public compliantForceTransfer = false;
 
     // TODO: Move modifiers to `ERC7984Rwa` or remove from mock if useless
     /// @dev Checks if the sender is an admin.
@@ -58,16 +58,18 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
-    ) internal override returns (bool) {
-        return compliantTransfer;
+    ) internal override returns (ebool compliant) {
+        compliant = FHE.asEbool(compliantTransfer);
+        FHE.allowThis(compliant);
     }
 
     function _isForceTransferCompliant(
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
-    ) internal override returns (bool) {
-        return compliantForceTransfer;
+    ) internal override returns (ebool compliant) {
+        compliant = FHE.asEbool(compliantForceTransfer);
+        FHE.allowThis(compliant);
     }
 
     function _validateHandleAllowance(bytes32 handle) internal view override onlyAdminOrAgent {}
diff --git a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
index 025be19a..c1a61ab3 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
@@ -38,6 +38,9 @@ abstract contract ERC7984Freezable is ERC7984 {
             confidentialBalanceOf(account),
             confidentialFrozen(account)
         );
+        if (!FHE.isInitialized(unfrozen)) {
+            return unfrozen;
+        }
         return FHE.select(success, unfrozen, FHE.asEuint64(0));
     }
 
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index cddead24..937409d8 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.27;
 
-import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
 import {ERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
 import {Multicall} from "@openzeppelin/contracts/utils/Multicall.sol";
@@ -30,8 +30,6 @@ abstract contract ERC7984Rwa is
 
     /// @dev The caller account is not authorized to perform the operation.
     error UnauthorizedSender(address account);
-    /// @dev The transfer does not follow token compliance.
-    error UncompliantTransfer(address from, address to, euint64 encryptedAmount);
 
     /// @dev Checks if the sender is an admin or an agent.
     modifier onlyAdminOrAgent() {
@@ -167,7 +165,14 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) internal virtual returns (euint64 transferred) {
-        require(_isForceTransferCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+        if (!FHE.isInitialized(encryptedAmount)) {
+            return encryptedAmount;
+        }
+        encryptedAmount = FHE.select(
+            _isForceTransferCompliant(from, to, encryptedAmount),
+            encryptedAmount,
+            FHE.asEuint64(0)
+        );
         _disableERC7984FreezableUpdateCheck(); // bypass frozen check
         _disableERC7984RestrictedUpdateCheck(); // bypass default restriction check
         if (to != address(0)) _checkRestriction(to); // only perform restriction check on `to`
@@ -183,7 +188,14 @@ abstract contract ERC7984Rwa is
         address to,
         euint64 encryptedAmount
     ) internal override(ERC7984Freezable, ERC7984Restricted, ERC7984) whenNotPaused returns (euint64 transferred) {
-        require(_isTransferCompliant(from, to, encryptedAmount), UncompliantTransfer(from, to, encryptedAmount));
+        if (!FHE.isInitialized(encryptedAmount)) {
+            return encryptedAmount;
+        }
+        encryptedAmount = FHE.select(
+            _isTransferCompliant(from, to, encryptedAmount),
+            encryptedAmount,
+            FHE.asEuint64(0)
+        );
         // frozen and restriction checks performed through inheritance
         transferred = super._update(from, to, encryptedAmount);
         _postTransferHook(from, to, encryptedAmount);
@@ -196,14 +208,14 @@ abstract contract ERC7984Rwa is
     function _checkFreezer() internal override onlyAdminOrAgent {}
 
     /// @dev Checks if a transfer follows compliance.
-    function _isTransferCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
+    function _isTransferCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (ebool);
 
     /// @dev Checks if a force transfer follows compliance.
     function _isForceTransferCompliant(
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal virtual returns (bool);
+    ) internal virtual returns (ebool);
 
     /// @dev Performs operation after transfer.
     function _postTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index 18f7fe62..5468f6c7 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -36,18 +36,16 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModu
         address /*from*/,
         address to,
         euint64 encryptedAmount
-    ) internal override returns (bool) {
+    ) internal override returns (ebool compliant) {
         if (
-            to == address(0) // burning
+            !FHE.isInitialized(encryptedAmount) || to == address(0) // if no amount or burning
         ) {
-            return true;
+            return FHE.asEbool(true);
         }
         (ebool increased, euint64 futureBalance) = FHESafeMath.tryIncrease(
             IERC7984(_token).confidentialBalanceOf(to),
             encryptedAmount
         );
-        ebool isCompliant = FHE.and(increased, FHE.le(futureBalance, _maxBalance));
-        isCompliant;
-        return false;
+        compliant = FHE.and(increased, FHE.le(futureBalance, _maxBalance));
     }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
index 0459550f..7d52e878 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.27;
 
-import {FHE, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984Rwa} from "../ERC7984Rwa.sol";
@@ -99,15 +99,19 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal virtual returns (bool) {
+    ) internal virtual returns (ebool compliant) {
+        if (!FHE.isInitialized(encryptedAmount)) {
+            return FHE.asEbool(true);
+        }
         address[] memory modules = _transferComplianceModules.values();
         uint256 modulesLength = modules.length;
+        compliant = FHE.asEbool(true);
         for (uint256 i = 0; i < modulesLength; i++) {
-            if (!IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)) {
-                return false;
-            }
+            compliant = FHE.and(
+                compliant,
+                IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)
+            );
         }
-        return true;
     }
 
     /// @dev Checks if a force transfer is compliant.
@@ -115,15 +119,19 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         address from,
         address to,
         euint64 encryptedAmount
-    ) internal virtual returns (bool) {
+    ) internal virtual returns (ebool compliant) {
+        if (!FHE.isInitialized(encryptedAmount)) {
+            return FHE.asEbool(true);
+        }
         address[] memory modules = _forceTransferComplianceModules.values();
         uint256 modulesLength = modules.length;
+        compliant = FHE.asEbool(true);
         for (uint256 i = 0; i < modulesLength; i++) {
-            if (!IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)) {
-                return false;
-            }
+            compliant = FHE.and(
+                compliant,
+                IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)
+            );
         }
-        return true;
     }
 
     /// @dev Performs operation after transfer.
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index 48c7fc00..e63fbd18 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.27;
 
-import {euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
 
@@ -33,16 +33,17 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceMod
     function _isCompliantTransfer(
         address /*from*/,
         address to,
-        euint64 /*encryptedAmount*/
-    ) internal override returns (bool) {
+        euint64 encryptedAmount
+    ) internal override returns (ebool) {
         if (
-            to == address(0) || // burning
+            FHE.isInitialized(encryptedAmount) || // no amount
+            to == address(0) || // or burning
             _investors.contains(to) || // or already investor
             _investors.length() < _maxInvestor // or not reached max investors limit
         ) {
-            return true;
+            return FHE.asEbool(true);
         }
-        return false;
+        return FHE.asEbool(false);
     }
 
     /// @dev Internal function which Performs operation after transfer.
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
index a5daf4a3..d0d53397 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.27;
 
-import {euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 import {IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
 
@@ -24,7 +24,7 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
     }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
-    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) public virtual returns (bool) {
+    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) public virtual returns (ebool) {
         return _isCompliantTransfer(from, to, encryptedAmount);
     }
 
@@ -38,9 +38,9 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
-    ) internal virtual returns (bool) {
+    ) internal virtual returns (ebool) {
         // default to non-compliant
-        return false;
+        return FHE.asEbool(false);
     }
 
     /// @dev Internal function which Performs operation after transfer.
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index 3f3ef38d..cf540458 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -188,13 +188,16 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
         .encrypt();
-      await expect(
+      const [, , transferred] = await callAndGetResult(
         token
           .connect(admin)
           ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
-      )
-        .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
-        .withArgs(ethers.ZeroAddress, recipient.address, encryptedInput.handles[0]);
+        transferEventSignature,
+      );
+      await token.getHandleAllowance(transferred, admin.address, true);
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferred, await token.getAddress(), admin),
+      ).to.eventually.equal(0);
     });
 
     it('should not mint if paused', async function () {
@@ -291,13 +294,19 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
         .encrypt();
-      await expect(
+      await token
+        .connect(admin)
+        ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
+      const [, , transferredHandle] = await callAndGetResult(
         token
           .connect(admin)
           ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
-      )
-        .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
-        .withArgs(recipient.address, ethers.ZeroAddress, encryptedInput.handles[0]);
+        transferEventSignature,
+      );
+      await token.getHandleAllowance(transferredHandle, admin.address, true);
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), admin),
+      ).to.eventually.equal(0);
     });
 
     it('should not burn if paused', async function () {
@@ -346,7 +355,6 @@ describe('ERC7984Rwa', function () {
               encryptedFrozenValueInput.inputProof,
             );
           await token.$_unsetCompliant();
-          expect(await token.compliantTransfer()).to.be.false;
           const amount = 25;
           let params = [recipient.address, anyone.address] as unknown as [
             from: AddressLike,
@@ -425,7 +433,6 @@ describe('ERC7984Rwa', function () {
             );
           // should force transfer even if not compliant
           await token.$_unsetCompliant();
-          expect(await token.compliantTransfer()).to.be.false;
           // should force transfer even if paused
           await token.connect(manager).pause();
           expect(await token.paused()).to.be.true;
@@ -578,7 +585,6 @@ describe('ERC7984Rwa', function () {
         .add64(amount)
         .encrypt();
       await token.$_setCompliant();
-      expect(await token.compliantTransfer()).to.be.true;
       const [from, to, transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)
@@ -631,13 +637,19 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not transfer if transfer not compliant', async function () {
-      const { token, recipient, anyone } = await fixture();
+      const { token, admin, recipient, anyone } = await fixture();
+      const encryptedMint = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(25)
+        .encrypt();
+      await token
+        .connect(admin)
+        ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedMint.handles[0], encryptedMint.inputProof);
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      expect(await token.compliantTransfer()).to.be.false;
-      await expect(
+      const [, , transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)
           ['confidentialTransfer(address,bytes32,bytes)'](
@@ -645,9 +657,11 @@ describe('ERC7984Rwa', function () {
             encryptedTransferValueInput.handles[0],
             encryptedTransferValueInput.inputProof,
           ),
-      )
-        .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
-        .withArgs(recipient.address, anyone.address, encryptedTransferValueInput.handles[0]);
+        transferEventSignature,
+      );
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+      ).to.eventually.equal(0);
     });
 
     it('should not transfer if frozen', async function () {
@@ -681,7 +695,6 @@ describe('ERC7984Rwa', function () {
         .add64(25)
         .encrypt();
       await token.$_setCompliant();
-      expect(await token.compliantTransfer()).to.be.true;
       const [, , transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)

From 4c459487840bdc65b18a67efc3257113eb62baf6 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 10 Sep 2025 11:05:06 +0200
Subject: [PATCH 39/52] Rename rwa mock functions

---
 contracts/mocks/token/ERC7984RwaMock.sol      |  8 ++---
 .../rwa/ERC7984RwaBalanceCapModule.sol        |  3 +-
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 30 +++++++++----------
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index 4ea777c7..cffb8e1f 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -34,19 +34,19 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         FHE.allow(encryptedAmount, msg.sender);
     }
 
-    function $_setCompliant() public {
+    function $_setCompliantTransfer() public {
         compliantTransfer = true;
     }
 
-    function $_unsetCompliant() public {
+    function $_unsetCompliantTransfer() public {
         compliantTransfer = false;
     }
 
-    function $_setForceCompliant() public {
+    function $_setCompliantForceTransfer() public {
         compliantForceTransfer = true;
     }
 
-    function $_unsetForceCompliant() public {
+    function $_unsetCompliantForceTransfer() public {
         compliantForceTransfer = false;
     }
 
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index 5468f6c7..f83ce1d5 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -17,7 +17,8 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModu
     euint64 private _maxBalance;
     address private _token;
 
-    constructor(address compliance, euint64 maxBalance) ERC7984RwaTransferComplianceModule(compliance) {
+    constructor(address compliance, address token, euint64 maxBalance) ERC7984RwaTransferComplianceModule(compliance) {
+        _token = token;
         setMaxBalance(maxBalance);
     }
 
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index cf540458..576aa8f5 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -131,7 +131,7 @@ describe('ERC7984Rwa', function () {
         const { admin, agent1, recipient } = await fixture();
         for (const manager of [admin, agent1]) {
           const { token } = await fixture();
-          await token.$_setCompliant();
+          await token.$_setCompliantTransfer();
           const amount = 100;
           let params = [recipient.address] as unknown as [
             account: AddressLike,
@@ -172,7 +172,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliant();
+      await token.$_setCompliantTransfer();
       await expect(
         token
           .connect(anyone)
@@ -225,7 +225,7 @@ describe('ERC7984Rwa', function () {
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
             .encrypt();
-          await token.$_setCompliant();
+          await token.$_setCompliantTransfer();
           await token
             .connect(manager)
             ['confidentialMint(address,bytes32,bytes)'](
@@ -278,7 +278,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), anyone.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliant();
+      await token.$_setCompliantTransfer();
       await expect(
         token
           .connect(anyone)
@@ -334,7 +334,7 @@ describe('ERC7984Rwa', function () {
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
             .encrypt();
-          await token.$_setCompliant();
+          await token.$_setCompliantTransfer();
           await token
             .connect(manager)
             ['confidentialMint(address,bytes32,bytes)'](
@@ -354,7 +354,7 @@ describe('ERC7984Rwa', function () {
               encryptedFrozenValueInput.handles[0],
               encryptedFrozenValueInput.inputProof,
             );
-          await token.$_unsetCompliant();
+          await token.$_unsetCompliantTransfer();
           const amount = 25;
           let params = [recipient.address, anyone.address] as unknown as [
             from: AddressLike,
@@ -372,7 +372,7 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
-          await token.$_setForceCompliant();
+          await token.$_setCompliantForceTransfer();
           const [from, to, transferredHandle] = await callAndGetResult(
             token
               .connect(manager)
@@ -411,7 +411,7 @@ describe('ERC7984Rwa', function () {
             .createEncryptedInput(await token.getAddress(), manager.address)
             .add64(100)
             .encrypt();
-          await token.$_setCompliant();
+          await token.$_setCompliantTransfer();
           await token
             .connect(manager)
             ['confidentialMint(address,bytes32,bytes)'](
@@ -432,7 +432,7 @@ describe('ERC7984Rwa', function () {
               encryptedFrozenValueInput.inputProof,
             );
           // should force transfer even if not compliant
-          await token.$_unsetCompliant();
+          await token.$_unsetCompliantTransfer();
           // should force transfer even if paused
           await token.connect(manager).pause();
           expect(await token.paused()).to.be.true;
@@ -453,7 +453,7 @@ describe('ERC7984Rwa', function () {
             await token.connect(manager).createEncryptedAmount(amount);
             params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
           }
-          await token.$_setForceCompliant();
+          await token.$_setCompliantForceTransfer();
           const [account, frozenAmountHandle] = await callAndGetResult(
             token
               .connect(manager)
@@ -559,7 +559,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliant();
+      await token.$_setCompliantTransfer();
       await token
         .connect(manager)
         ['confidentialMint(address,bytes32,bytes)'](
@@ -584,7 +584,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(amount)
         .encrypt();
-      await token.$_setCompliant();
+      await token.$_setCompliantTransfer();
       const [from, to, transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)
@@ -670,7 +670,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), manager.address)
         .add64(100)
         .encrypt();
-      await token.$_setCompliant();
+      await token.$_setCompliantTransfer();
       await token
         .connect(manager)
         ['confidentialMint(address,bytes32,bytes)'](
@@ -694,7 +694,7 @@ describe('ERC7984Rwa', function () {
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      await token.$_setCompliant();
+      await token.$_setCompliantTransfer();
       const [, , transferredHandle] = await callAndGetResult(
         token
           .connect(recipient)
@@ -723,7 +723,7 @@ describe('ERC7984Rwa', function () {
       it(`should not transfer if ${arg ? 'sender' : 'receiver'} blocked `, async function () {
         const { token, admin: manager, recipient, anyone } = await fixture();
         const account = arg ? recipient : anyone;
-        await token.$_setCompliant();
+        await token.$_setCompliantTransfer();
         const encryptedInput = await fhevm
           .createEncryptedInput(await token.getAddress(), recipient.address)
           .add64(25)

From 9974faeedb5daf2dc3e3d65c3566b3b2e5abed9f Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 10 Sep 2025 11:12:10 +0200
Subject: [PATCH 40/52] Immutable token in balance cap module

---
 .../token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index f83ce1d5..2b195313 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -15,7 +15,7 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModu
     using EnumerableSet for *;
 
     euint64 private _maxBalance;
-    address private _token;
+    address private immutable _token;
 
     constructor(address compliance, address token, euint64 maxBalance) ERC7984RwaTransferComplianceModule(compliance) {
         _token = token;

From 33325810c520fdf24e23518a9cd36012b9c5695d Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Thu, 11 Sep 2025 12:01:56 +0200
Subject: [PATCH 41/52] Switch to always-on/transfer-only compliance modules

---
 contracts/interfaces/IERC7984Rwa.sol          |  6 +-
 contracts/mocks/token/ERC7984RwaMock.sol      |  4 +-
 .../token/ERC7984/extensions/ERC7984Rwa.sol   | 24 ++---
 .../extensions/rwa/ERC7984RwaCompliance.sol   | 87 +++++++++++--------
 .../rwa/ERC7984RwaInvestorCapModule.sol       |  4 +-
 .../ERC7984RwaTransferComplianceModule.sol    |  9 +-
 6 files changed, 71 insertions(+), 63 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index fa3767e8..d87b0a2d 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -7,8 +7,8 @@ import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 import {IERC7984Restricted} from "./IERC7984Restricted.sol";
 
-uint256 constant TRANSFER_COMPLIANCE_MODULE_TYPE = 1;
-uint256 constant FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE = 2;
+uint256 constant ALWAYS_ON_MODULE_TYPE = 1;
+uint256 constant TRANSFER_ONLY_MODULE_TYPE = 2;
 
 /// @dev Base interface for confidential RWA contracts.
 interface IERC7984RwaBase {
@@ -101,5 +101,5 @@ interface IERC7984RwaTransferComplianceModule {
     /// @dev Checks if a transfer is compliant. Should be non-mutating.
     function isCompliantTransfer(address from, address to, euint64 encryptedAmount) external returns (ebool);
     /// @dev Performs operation after transfer.
-    function postTransferHook(address from, address to, euint64 encryptedAmount) external;
+    function postTransfer(address from, address to, euint64 encryptedAmount) external;
 }
diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index cffb8e1f..b89c9164 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -54,7 +54,7 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return _mint(to, FHE.asEuint64(amount));
     }
 
-    function _isTransferCompliant(
+    function _preCheckTransfer(
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
@@ -63,7 +63,7 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         FHE.allowThis(compliant);
     }
 
-    function _isForceTransferCompliant(
+    function _preCheckForceTransfer(
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 937409d8..50d7b1b7 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -169,7 +169,7 @@ abstract contract ERC7984Rwa is
             return encryptedAmount;
         }
         encryptedAmount = FHE.select(
-            _isForceTransferCompliant(from, to, encryptedAmount),
+            _preCheckForceTransfer(from, to, encryptedAmount),
             encryptedAmount,
             FHE.asEuint64(0)
         );
@@ -177,7 +177,7 @@ abstract contract ERC7984Rwa is
         _disableERC7984RestrictedUpdateCheck(); // bypass default restriction check
         if (to != address(0)) _checkRestriction(to); // only perform restriction check on `to`
         transferred = super._update(from, to, encryptedAmount); // bypass compliance check
-        _postForceTransferHook(from, to, encryptedAmount);
+        _postForceTransfer(from, to, encryptedAmount);
         _restoreERC7984FreezableUpdateCheck();
         _restoreERC7984RestrictedUpdateCheck();
     }
@@ -191,14 +191,10 @@ abstract contract ERC7984Rwa is
         if (!FHE.isInitialized(encryptedAmount)) {
             return encryptedAmount;
         }
-        encryptedAmount = FHE.select(
-            _isTransferCompliant(from, to, encryptedAmount),
-            encryptedAmount,
-            FHE.asEuint64(0)
-        );
+        encryptedAmount = FHE.select(_preCheckTransfer(from, to, encryptedAmount), encryptedAmount, FHE.asEuint64(0));
         // frozen and restriction checks performed through inheritance
         transferred = super._update(from, to, encryptedAmount);
-        _postTransferHook(from, to, encryptedAmount);
+        _postTransfer(from, to, encryptedAmount);
     }
 
     /**
@@ -208,18 +204,14 @@ abstract contract ERC7984Rwa is
     function _checkFreezer() internal override onlyAdminOrAgent {}
 
     /// @dev Checks if a transfer follows compliance.
-    function _isTransferCompliant(address from, address to, euint64 encryptedAmount) internal virtual returns (ebool);
+    function _preCheckTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (ebool);
 
     /// @dev Checks if a force transfer follows compliance.
-    function _isForceTransferCompliant(
-        address from,
-        address to,
-        euint64 encryptedAmount
-    ) internal virtual returns (ebool);
+    function _preCheckForceTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (ebool);
 
     /// @dev Performs operation after transfer.
-    function _postTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
+    function _postTransfer(address from, address to, euint64 encryptedAmount) internal virtual {}
 
     /// @dev Performs operation after force transfer.
-    function _postForceTransferHook(address from, address to, euint64 encryptedAmount) internal virtual {}
+    function _postForceTransfer(address from, address to, euint64 encryptedAmount) internal virtual {}
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
index 7d52e878..c66d1b65 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {FHE, ebool, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, TRANSFER_ONLY_MODULE_TYPE, ALWAYS_ON_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984Rwa} from "../ERC7984Rwa.sol";
 
 /**
@@ -14,8 +14,8 @@ import {ERC7984Rwa} from "../ERC7984Rwa.sol";
 abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
     using EnumerableSet for *;
 
-    EnumerableSet.AddressSet private _transferComplianceModules;
-    EnumerableSet.AddressSet private _forceTransferComplianceModules;
+    EnumerableSet.AddressSet private _alwaysOnModules;
+    EnumerableSet.AddressSet private _transferOnlyModules;
 
     /// @dev Emitted when a module is installed.
     event ModuleInstalled(uint256 moduleTypeId, address module);
@@ -40,7 +40,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
      * * Force transfer compliance module
      */
     function supportsModule(uint256 moduleTypeId) public view virtual returns (bool) {
-        return moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE || moduleTypeId == FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE;
+        return moduleTypeId == ALWAYS_ON_MODULE_TYPE || moduleTypeId == TRANSFER_ONLY_MODULE_TYPE;
     }
 
     /**
@@ -66,13 +66,10 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
             ERC7984RwaNotTransferComplianceModule(module)
         );
 
-        if (moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE) {
-            require(_transferComplianceModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
-        } else if (moduleTypeId == FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE) {
-            require(
-                _forceTransferComplianceModules.add(module),
-                ERC7984RwaAlreadyInstalledModule(moduleTypeId, module)
-            );
+        if (moduleTypeId == ALWAYS_ON_MODULE_TYPE) {
+            require(_alwaysOnModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
+        } else if (moduleTypeId == TRANSFER_ONLY_MODULE_TYPE) {
+            require(_transferOnlyModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
         }
         emit ModuleInstalled(moduleTypeId, module);
     }
@@ -80,22 +77,42 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
     /// @dev Internal function which uninstalls a transfer compliance module.
     function _uninstallModule(uint256 moduleTypeId, address module) internal virtual {
         require(supportsModule(moduleTypeId), ERC7984RwaUnsupportedModuleType(moduleTypeId));
-        if (moduleTypeId == TRANSFER_COMPLIANCE_MODULE_TYPE) {
-            require(
-                _transferComplianceModules.remove(module),
-                ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module)
-            );
-        } else if (moduleTypeId == FORCE_TRANSFER_COMPLIANCE_MODULE_TYPE) {
-            require(
-                _forceTransferComplianceModules.remove(module),
-                ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module)
-            );
+        if (moduleTypeId == ALWAYS_ON_MODULE_TYPE) {
+            require(_alwaysOnModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module));
+        } else if (moduleTypeId == TRANSFER_ONLY_MODULE_TYPE) {
+            require(_transferOnlyModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module));
         }
         emit ModuleUninstalled(moduleTypeId, module);
     }
 
-    /// @dev Checks if a transfer is compliant.
-    function _isTransferCompliantTransfer(
+    /// @dev Checks if a transfer follows compliance.
+    function _preCheckTransfer(address from, address to, euint64 encryptedAmount) internal override returns (ebool) {
+        return
+            FHE.and(_checkAlwaysBefore(from, to, encryptedAmount), _checkOnlyBeforeTransfer(from, to, encryptedAmount));
+    }
+
+    /// @dev Checks if a force transfer follows compliance.
+    function _preCheckForceTransfer(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) internal override returns (ebool) {
+        return _checkAlwaysBefore(from, to, encryptedAmount);
+    }
+
+    /// @dev Peforms operations after transfer.
+    function _postTransfer(address from, address to, euint64 encryptedAmount) internal override {
+        _runAlwaysAfter(from, to, encryptedAmount);
+        _runOnlyAfterTransfer(from, to, encryptedAmount);
+    }
+
+    /// @dev Peforms operations after force transfer.
+    function _postForceTransfer(address from, address to, euint64 encryptedAmount) internal override {
+        _runAlwaysAfter(from, to, encryptedAmount);
+    }
+
+    /// @dev Checks always-on compliance.
+    function _checkAlwaysBefore(
         address from,
         address to,
         euint64 encryptedAmount
@@ -103,7 +120,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         if (!FHE.isInitialized(encryptedAmount)) {
             return FHE.asEbool(true);
         }
-        address[] memory modules = _transferComplianceModules.values();
+        address[] memory modules = _alwaysOnModules.values();
         uint256 modulesLength = modules.length;
         compliant = FHE.asEbool(true);
         for (uint256 i = 0; i < modulesLength; i++) {
@@ -114,8 +131,8 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         }
     }
 
-    /// @dev Checks if a force transfer is compliant.
-    function _isTransferCompliantForceTransfer(
+    /// @dev Checks transfer-only compliance.
+    function _checkOnlyBeforeTransfer(
         address from,
         address to,
         euint64 encryptedAmount
@@ -123,7 +140,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         if (!FHE.isInitialized(encryptedAmount)) {
             return FHE.asEbool(true);
         }
-        address[] memory modules = _forceTransferComplianceModules.values();
+        address[] memory modules = _transferOnlyModules.values();
         uint256 modulesLength = modules.length;
         compliant = FHE.asEbool(true);
         for (uint256 i = 0; i < modulesLength; i++) {
@@ -134,21 +151,21 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         }
     }
 
-    /// @dev Performs operation after transfer.
-    function _postTransferHook(address from, address to, euint64 encryptedAmount) internal override {
-        address[] memory modules = _transferComplianceModules.values();
+    /// @dev Runs always.
+    function _runAlwaysAfter(address from, address to, euint64 encryptedAmount) internal virtual {
+        address[] memory modules = _alwaysOnModules.values();
         uint256 modulesLength = modules.length;
         for (uint256 i = 0; i < modulesLength; i++) {
-            IERC7984RwaTransferComplianceModule(modules[i]).postTransferHook(from, to, encryptedAmount);
+            IERC7984RwaTransferComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
         }
     }
 
-    /// @dev Performs operation after force transfer.
-    function _postForceTransferHook(address from, address to, euint64 encryptedAmount) internal override {
-        address[] memory modules = _forceTransferComplianceModules.values();
+    /// @dev Runs only after transfer.
+    function _runOnlyAfterTransfer(address from, address to, euint64 encryptedAmount) internal virtual {
+        address[] memory modules = _transferOnlyModules.values();
         uint256 modulesLength = modules.length;
         for (uint256 i = 0; i < modulesLength; i++) {
-            IERC7984RwaTransferComplianceModule(modules[i]).postTransferHook(from, to, encryptedAmount);
+            IERC7984RwaTransferComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
         }
     }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index e63fbd18..f0cc42f7 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -46,8 +46,8 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceMod
         return FHE.asEbool(false);
     }
 
-    /// @dev Internal function which Performs operation after transfer.
-    function _postTransferHook(address /*from*/, address to, euint64 /*encryptedAmount*/) internal override {
+    /// @dev Internal function which performs operation after transfer.
+    function _postTransfer(address /*from*/, address to, euint64 /*encryptedAmount*/) internal override {
         if (!_investors.contains(to)) {
             _investors.add(to);
         }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
index d0d53397..91870459 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
@@ -4,8 +4,7 @@ pragma solidity ^0.8.27;
 
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
-import {IERC7984RwaTransferComplianceModule, TRANSFER_COMPLIANCE_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
-
+import {IERC7984RwaTransferComplianceModule} from "./../../../../interfaces/IERC7984Rwa.sol";
 /**
  * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).
  */
@@ -29,8 +28,8 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
     }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
-    function postTransferHook(address from, address to, euint64 encryptedAmount) public virtual {
-        _postTransferHook(from, to, encryptedAmount);
+    function postTransfer(address from, address to, euint64 encryptedAmount) public virtual {
+        _postTransfer(from, to, encryptedAmount);
     }
 
     /// @dev Internal function which checks if a transfer is compliant.
@@ -44,7 +43,7 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
     }
 
     /// @dev Internal function which Performs operation after transfer.
-    function _postTransferHook(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal virtual {
+    function _postTransfer(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal virtual {
         // default to no-op
     }
 }

From c81b703f95c1396735304cd34e8d58ff09e2aace Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Thu, 11 Sep 2025 12:04:15 +0200
Subject: [PATCH 42/52] Typo

---
 .../token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol     | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
index c66d1b65..043337c3 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
@@ -100,13 +100,13 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         return _checkAlwaysBefore(from, to, encryptedAmount);
     }
 
-    /// @dev Peforms operations after transfer.
+    /// @dev Performs operations after transfer.
     function _postTransfer(address from, address to, euint64 encryptedAmount) internal override {
         _runAlwaysAfter(from, to, encryptedAmount);
         _runOnlyAfterTransfer(from, to, encryptedAmount);
     }
 
-    /// @dev Peforms operations after force transfer.
+    /// @dev Performs operations after force transfer.
     function _postForceTransfer(address from, address to, euint64 encryptedAmount) internal override {
         _runAlwaysAfter(from, to, encryptedAmount);
     }

From 7d438c320700e8cfa590830e189c7457e9396477 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Thu, 11 Sep 2025 12:27:14 +0200
Subject: [PATCH 43/52] Use enum for compliance module type

---
 contracts/interfaces/IERC7984Rwa.sol          | 12 +++--
 .../extensions/rwa/ERC7984RwaCompliance.sol   | 52 +++++++++----------
 2 files changed, 33 insertions(+), 31 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index d87b0a2d..f90a5c6f 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -7,9 +7,6 @@ import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol";
 import {IERC7984} from "./IERC7984.sol";
 import {IERC7984Restricted} from "./IERC7984Restricted.sol";
 
-uint256 constant ALWAYS_ON_MODULE_TYPE = 1;
-uint256 constant TRANSFER_ONLY_MODULE_TYPE = 2;
-
 /// @dev Base interface for confidential RWA contracts.
 interface IERC7984RwaBase {
     /// @dev Emitted when the contract is paused.
@@ -88,10 +85,15 @@ interface IERC7984Rwa is IERC7984, IERC7984RwaBase, IERC165, IAccessControl {}
 
 /// @dev Interface for confidential RWA compliance.
 interface IERC7984RwaCompliance {
+    enum ComplianceModuleType {
+        ALWAYS_ON,
+        TRANSFER_ONLY
+    }
+
     /// @dev Installs a transfer compliance module.
-    function installModule(uint256 moduleTypeId, address module) external;
+    function installModule(ComplianceModuleType moduleType, address module) external;
     /// @dev Uninstalls a transfer compliance module.
-    function uninstallModule(uint256 moduleTypeId, address module) external;
+    function uninstallModule(ComplianceModuleType moduleType, address module) external;
 }
 
 /// @dev Interface for confidential RWA transfer compliance module.
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
index 043337c3..0bc16790 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {FHE, ebool, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule, TRANSFER_ONLY_MODULE_TYPE, ALWAYS_ON_MODULE_TYPE} from "./../../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule} from "./../../../../interfaces/IERC7984Rwa.sol";
 import {ERC7984Rwa} from "../ERC7984Rwa.sol";
 
 /**
@@ -18,18 +18,18 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
     EnumerableSet.AddressSet private _transferOnlyModules;
 
     /// @dev Emitted when a module is installed.
-    event ModuleInstalled(uint256 moduleTypeId, address module);
+    event ModuleInstalled(ComplianceModuleType moduleType, address module);
     /// @dev Emitted when a module is uninstalled.
-    event ModuleUninstalled(uint256 moduleTypeId, address module);
+    event ModuleUninstalled(ComplianceModuleType moduleType, address module);
 
     /// @dev The module type is not supported.
-    error ERC7984RwaUnsupportedModuleType(uint256 moduleTypeId);
+    error ERC7984RwaUnsupportedModuleType(ComplianceModuleType moduleType);
     /// @dev The address is not a transfer compliance module.
     error ERC7984RwaNotTransferComplianceModule(address module);
     /// @dev The module is already installed.
-    error ERC7984RwaAlreadyInstalledModule(uint256 moduleTypeId, address module);
+    error ERC7984RwaAlreadyInstalledModule(ComplianceModuleType moduleType, address module);
     /// @dev The module is already uninstalled.
-    error ERC7984RwaAlreadyUninstalledModule(uint256 moduleTypeId, address module);
+    error ERC7984RwaAlreadyUninstalledModule(ComplianceModuleType moduleType, address module);
 
     /**
      * @dev Check if a certain module typeId is supported.
@@ -39,8 +39,8 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
      * * Transfer compliance module
      * * Force transfer compliance module
      */
-    function supportsModule(uint256 moduleTypeId) public view virtual returns (bool) {
-        return moduleTypeId == ALWAYS_ON_MODULE_TYPE || moduleTypeId == TRANSFER_ONLY_MODULE_TYPE;
+    function supportsModule(ComplianceModuleType moduleType) public view virtual returns (bool) {
+        return moduleType == ComplianceModuleType.ALWAYS_ON || moduleType == ComplianceModuleType.TRANSFER_ONLY;
     }
 
     /**
@@ -48,41 +48,41 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
      * @dev Consider gas footprint of the module before adding it since all modules will perform
      * all steps (pre-check, compliance check, post-hook) in a single transaction.
      */
-    function installModule(uint256 moduleTypeId, address module) public virtual onlyAdminOrAgent {
-        _installModule(moduleTypeId, module);
+    function installModule(ComplianceModuleType moduleType, address module) public virtual onlyAdminOrAgent {
+        _installModule(moduleType, module);
     }
 
     /// @inheritdoc IERC7984RwaCompliance
-    function uninstallModule(uint256 moduleTypeId, address module) public virtual onlyAdminOrAgent {
-        _uninstallModule(moduleTypeId, module);
+    function uninstallModule(ComplianceModuleType moduleType, address module) public virtual onlyAdminOrAgent {
+        _uninstallModule(moduleType, module);
     }
 
     /// @dev Internal function which installs a transfer compliance module.
-    function _installModule(uint256 moduleTypeId, address module) internal virtual {
-        require(supportsModule(moduleTypeId), ERC7984RwaUnsupportedModuleType(moduleTypeId));
+    function _installModule(ComplianceModuleType moduleType, address module) internal virtual {
+        require(supportsModule(moduleType), ERC7984RwaUnsupportedModuleType(moduleType));
         require(
             IERC7984RwaTransferComplianceModule(module).isModule() ==
                 IERC7984RwaTransferComplianceModule.isModule.selector,
             ERC7984RwaNotTransferComplianceModule(module)
         );
 
-        if (moduleTypeId == ALWAYS_ON_MODULE_TYPE) {
-            require(_alwaysOnModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
-        } else if (moduleTypeId == TRANSFER_ONLY_MODULE_TYPE) {
-            require(_transferOnlyModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleTypeId, module));
+        if (moduleType == ComplianceModuleType.ALWAYS_ON) {
+            require(_alwaysOnModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleType, module));
+        } else if (moduleType == ComplianceModuleType.TRANSFER_ONLY) {
+            require(_transferOnlyModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleType, module));
         }
-        emit ModuleInstalled(moduleTypeId, module);
+        emit ModuleInstalled(moduleType, module);
     }
 
     /// @dev Internal function which uninstalls a transfer compliance module.
-    function _uninstallModule(uint256 moduleTypeId, address module) internal virtual {
-        require(supportsModule(moduleTypeId), ERC7984RwaUnsupportedModuleType(moduleTypeId));
-        if (moduleTypeId == ALWAYS_ON_MODULE_TYPE) {
-            require(_alwaysOnModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module));
-        } else if (moduleTypeId == TRANSFER_ONLY_MODULE_TYPE) {
-            require(_transferOnlyModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleTypeId, module));
+    function _uninstallModule(ComplianceModuleType moduleType, address module) internal virtual {
+        require(supportsModule(moduleType), ERC7984RwaUnsupportedModuleType(moduleType));
+        if (moduleType == ComplianceModuleType.ALWAYS_ON) {
+            require(_alwaysOnModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleType, module));
+        } else if (moduleType == ComplianceModuleType.TRANSFER_ONLY) {
+            require(_transferOnlyModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleType, module));
         }
-        emit ModuleUninstalled(moduleTypeId, module);
+        emit ModuleUninstalled(moduleType, module);
     }
 
     /// @dev Checks if a transfer follows compliance.

From fea22af80c3f589a8abc38c08627ee9c649ea054 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Mon, 15 Sep 2025 12:18:54 +0200
Subject: [PATCH 44/52] Enable token handles access to modules

---
 contracts/interfaces/IERC7984Rwa.sol          |   6 +
 .../token/ERC7984RwaBalanceCapModuleMock.sol  |  10 +
 .../token/ERC7984RwaInvestorCapModuleMock.sol |  10 +
 .../token/ERC7984RwaModularComplianceMock.sol |  15 ++
 .../rwa/ERC7984RwaBalanceCapModule.sol        |  33 +--
 .../extensions/rwa/ERC7984RwaCompliance.sol   |  37 ++-
 .../rwa/ERC7984RwaInvestorCapModule.sol       |  14 +-
 .../ERC7984RwaTransferComplianceModule.sol    |  53 +++-
 .../extensions/ERC7984RwaCompliance.test.ts   | 238 ++++++++++++++++++
 9 files changed, 379 insertions(+), 37 deletions(-)
 create mode 100644 contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol
 create mode 100644 contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol
 create mode 100644 contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
 create mode 100644 test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index f90a5c6f..93b965de 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -21,6 +21,10 @@ interface IERC7984RwaBase {
     /// @dev The operation failed because the contract is paused.
     error EnforcedPause();
 
+    /// @dev Returns true if has admin role, false otherwise.
+    function isAdmin(address account) external view returns (bool);
+    /// @dev Returns true if agent, false otherwise.
+    function isAgent(address account) external view returns (bool);
     /// @dev Returns true if the contract is paused, and false otherwise.
     function paused() external view returns (bool);
     /// @dev Pauses contract.
@@ -94,6 +98,8 @@ interface IERC7984RwaCompliance {
     function installModule(ComplianceModuleType moduleType, address module) external;
     /// @dev Uninstalls a transfer compliance module.
     function uninstallModule(ComplianceModuleType moduleType, address module) external;
+    /// @dev Checks if a compliance module is installed.
+    function isModuleInstalled(ComplianceModuleType moduleType, address module) external view returns (bool);
 }
 
 /// @dev Interface for confidential RWA transfer compliance module.
diff --git a/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol b/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol
new file mode 100644
index 00000000..8b2610ca
--- /dev/null
+++ b/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+import {ERC7984RwaBalanceCapModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol";
+
+contract ERC7984RwaBalanceCapModuleMock is ERC7984RwaBalanceCapModule, SepoliaConfig {
+    constructor(address compliance) ERC7984RwaBalanceCapModule(compliance) {}
+}
diff --git a/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol b/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol
new file mode 100644
index 00000000..c2306d60
--- /dev/null
+++ b/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+import {ERC7984RwaInvestorCapModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol";
+
+contract ERC7984RwaInvestorCapModuleMock is ERC7984RwaInvestorCapModule, SepoliaConfig {
+    constructor(address compliance, uint256 maxInvestor) ERC7984RwaInvestorCapModule(compliance, maxInvestor) {}
+}
diff --git a/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol b/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
new file mode 100644
index 00000000..b3f1409f
--- /dev/null
+++ b/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+import {FHE, ebool, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
+import {Impl} from "@fhevm/solidity/lib/Impl.sol";
+import {ERC7984Rwa} from "../../token/ERC7984/extensions/ERC7984Rwa.sol";
+import {ERC7984RwaCompliance} from "../../token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol";
+import {FHESafeMath} from "../../utils/FHESafeMath.sol";
+import {HandleAccessManager} from "../../utils/HandleAccessManager.sol";
+
+contract ERC7984RwaModularComplianceMock is ERC7984RwaCompliance, SepoliaConfig {
+    constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
+}
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index 2b195313..4cf59907 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -2,9 +2,9 @@
 
 pragma solidity ^0.8.27;
 
-import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984} from "./../../../../interfaces/IERC7984.sol";
+import {IERC7984} from "../../../../interfaces/IERC7984.sol";
 import {FHESafeMath} from "../../../../utils/FHESafeMath.sol";
 import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
 
@@ -14,17 +14,21 @@ import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferCompliance
 abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModule {
     using EnumerableSet for *;
 
-    euint64 private _maxBalance;
     address private immutable _token;
+    euint64 private _maxBalance;
 
-    constructor(address compliance, address token, euint64 maxBalance) ERC7984RwaTransferComplianceModule(compliance) {
+    constructor(address token) ERC7984RwaTransferComplianceModule(token) {
         _token = token;
-        setMaxBalance(maxBalance);
+    }
+
+    /// @dev Sets max balance of an investor with proof.
+    function setMaxBalance(externalEuint64 maxBalance, bytes calldata inputProof) public virtual onlyTokenAdmin {
+        FHE.allowThis(_maxBalance = FHE.fromExternal(maxBalance, inputProof));
     }
 
     /// @dev Sets max balance of an investor.
-    function setMaxBalance(euint64 maxBalance) public virtual onlyCompliance {
-        _maxBalance = maxBalance;
+    function setMaxBalance(euint64 maxBalance) public virtual onlyTokenAdmin {
+        FHE.allowThis(_maxBalance = maxBalance);
     }
 
     /// @dev Gets max balance of an investor.
@@ -38,15 +42,16 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModu
         address to,
         euint64 encryptedAmount
     ) internal override returns (ebool compliant) {
-        if (
-            !FHE.isInitialized(encryptedAmount) || to == address(0) // if no amount or burning
-        ) {
+        if (!FHE.isInitialized(encryptedAmount) || to == address(0)) {
+            // if no amount or burning
             return FHE.asEbool(true);
         }
-        (ebool increased, euint64 futureBalance) = FHESafeMath.tryIncrease(
-            IERC7984(_token).confidentialBalanceOf(to),
-            encryptedAmount
-        );
+        euint64 balance = IERC7984(_token).confidentialBalanceOf(to);
+        if (FHE.isInitialized(balance)) {
+            _allowTokenHandleToThis(balance);
+        }
+        _allowTokenHandleToThis(encryptedAmount);
+        (ebool increased, euint64 futureBalance) = FHESafeMath.tryIncrease(balance, encryptedAmount);
         compliant = FHE.and(increased, FHE.le(futureBalance, _maxBalance));
     }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
index 0bc16790..f751007f 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
@@ -2,16 +2,17 @@
 
 pragma solidity ^0.8.27;
 
-import {FHE, ebool, externalEuint64, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule} from "./../../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule} from "../../../../interfaces/IERC7984Rwa.sol";
+import {HandleAccessManager} from "../../../../utils/HandleAccessManager.sol";
 import {ERC7984Rwa} from "../ERC7984Rwa.sol";
 
 /**
  * @dev Extension of {ERC7984Rwa} that supports compliance modules for confidential Real World Assets (RWAs).
  * Inspired by ERC-7579 modules.
  */
-abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
+abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, HandleAccessManager {
     using EnumerableSet for *;
 
     EnumerableSet.AddressSet private _alwaysOnModules;
@@ -30,6 +31,8 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
     error ERC7984RwaAlreadyInstalledModule(ComplianceModuleType moduleType, address module);
     /// @dev The module is already uninstalled.
     error ERC7984RwaAlreadyUninstalledModule(ComplianceModuleType moduleType, address module);
+    /// @dev The sender is not a compliance module.
+    error SenderNotComplianceModule(address account);
 
     /**
      * @dev Check if a certain module typeId is supported.
@@ -57,12 +60,19 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         _uninstallModule(moduleType, module);
     }
 
+    /// @inheritdoc IERC7984RwaCompliance
+    function isModuleInstalled(ComplianceModuleType moduleType, address module) public view virtual returns (bool) {
+        return _isModuleInstalled(moduleType, module);
+    }
+
     /// @dev Internal function which installs a transfer compliance module.
     function _installModule(ComplianceModuleType moduleType, address module) internal virtual {
         require(supportsModule(moduleType), ERC7984RwaUnsupportedModuleType(moduleType));
+        (bool success, bytes memory returnData) = module.staticcall(
+            abi.encodePacked(IERC7984RwaTransferComplianceModule.isModule.selector)
+        );
         require(
-            IERC7984RwaTransferComplianceModule(module).isModule() ==
-                IERC7984RwaTransferComplianceModule.isModule.selector,
+            success && bytes4(returnData) == IERC7984RwaTransferComplianceModule.isModule.selector,
             ERC7984RwaNotTransferComplianceModule(module)
         );
 
@@ -85,6 +95,13 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         emit ModuleUninstalled(moduleType, module);
     }
 
+    /// @dev Checks if a compliance module is installed.
+    function _isModuleInstalled(ComplianceModuleType moduleType, address module) internal view virtual returns (bool) {
+        if (moduleType == ComplianceModuleType.ALWAYS_ON) return _alwaysOnModules.contains(module);
+        if (moduleType == ComplianceModuleType.TRANSFER_ONLY) return _transferOnlyModules.contains(module);
+        return false;
+    }
+
     /// @dev Checks if a transfer follows compliance.
     function _preCheckTransfer(address from, address to, euint64 encryptedAmount) internal override returns (ebool) {
         return
@@ -151,7 +168,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
         }
     }
 
-    /// @dev Runs always.
+    /// @dev Runs always after.
     function _runAlwaysAfter(address from, address to, euint64 encryptedAmount) internal virtual {
         address[] memory modules = _alwaysOnModules.values();
         uint256 modulesLength = modules.length;
@@ -168,4 +185,12 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance {
             IERC7984RwaTransferComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
         }
     }
+
+    /// @dev Allow modules to get access to token handles over {HandleAccessManager-getHandleAllowance}.
+    function _validateHandleAllowance(bytes32) internal view override {
+        require(
+            _alwaysOnModules.contains(msg.sender) || _transferOnlyModules.contains(msg.sender),
+            SenderNotComplianceModule(msg.sender)
+        );
+    }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index f0cc42f7..4aa19a99 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -16,19 +16,24 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceMod
     EnumerableSet.AddressSet private _investors;
 
     constructor(address compliance, uint256 maxInvestor) ERC7984RwaTransferComplianceModule(compliance) {
-        setMaxInvestor(maxInvestor);
+        _maxInvestor = maxInvestor;
     }
 
     /// @dev Sets max number of investors.
-    function setMaxInvestor(uint256 maxInvestor) public virtual onlyCompliance {
+    function setMaxInvestor(uint256 maxInvestor) public virtual onlyTokenAdmin {
         _maxInvestor = maxInvestor;
     }
 
     /// @dev Gets max number of investors.
-    function getMaxInvestor() public virtual returns (uint256) {
+    function getMaxInvestor() public view virtual returns (uint256) {
         return _maxInvestor;
     }
 
+    /// @dev Gets current number of investors.
+    function getCurrentInvestor() public view virtual returns (uint256) {
+        return _investors.length();
+    }
+
     /// @dev Internal function which checks if a transfer is compliant.
     function _isCompliantTransfer(
         address /*from*/,
@@ -36,13 +41,14 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceMod
         euint64 encryptedAmount
     ) internal override returns (ebool) {
         if (
-            FHE.isInitialized(encryptedAmount) || // no amount
+            !FHE.isInitialized(encryptedAmount) || // no amount
             to == address(0) || // or burning
             _investors.contains(to) || // or already investor
             _investors.length() < _maxInvestor // or not reached max investors limit
         ) {
             return FHE.asEbool(true);
         }
+
         return FHE.asEbool(false);
     }
 
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
index 91870459..13f79b01 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
@@ -3,19 +3,35 @@
 pragma solidity ^0.8.27;
 
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
-import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
-import {IERC7984RwaTransferComplianceModule} from "./../../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984Rwa, IERC7984RwaTransferComplianceModule} from "../../../../interfaces/IERC7984Rwa.sol";
+import {HandleAccessManager} from "../../../../utils/HandleAccessManager.sol";
+
 /**
  * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).
  */
-abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule, Ownable {
-    /// @dev Throws if called by any account other than the compliance.
-    modifier onlyCompliance() {
-        _checkOwner();
+abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule {
+    address private immutable _token;
+
+    /// @dev The sender is not the token.
+    error SenderNotToken(address account);
+    /// @dev The sender is not the token admin.
+    error SenderNotTokenAdmin(address account);
+
+    /// @dev Throws if called by any account other than the token.
+    modifier onlyToken() {
+        require(msg.sender == _token, SenderNotToken(msg.sender));
+        _;
+    }
+
+    /// @dev Throws if called by any account other than the token admin.
+    modifier onlyTokenAdmin() {
+        require(IERC7984Rwa(_token).isAdmin(msg.sender), SenderNotTokenAdmin(msg.sender));
         _;
     }
 
-    constructor(address compliance) Ownable(compliance) {}
+    constructor(address token) {
+        _token = token;
+    }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
     function isModule() public pure override returns (bytes4) {
@@ -23,8 +39,12 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
     }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
-    function isCompliantTransfer(address from, address to, euint64 encryptedAmount) public virtual returns (ebool) {
-        return _isCompliantTransfer(from, to, encryptedAmount);
+    function isCompliantTransfer(
+        address from,
+        address to,
+        euint64 encryptedAmount
+    ) public virtual onlyToken returns (ebool compliant) {
+        FHE.allow(compliant = _isCompliantTransfer(from, to, encryptedAmount), msg.sender);
     }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
@@ -37,13 +57,20 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
         address /*from*/,
         address /*to*/,
         euint64 /*encryptedAmount*/
-    ) internal virtual returns (ebool) {
-        // default to non-compliant
-        return FHE.asEbool(false);
-    }
+    ) internal virtual returns (ebool);
 
     /// @dev Internal function which Performs operation after transfer.
     function _postTransfer(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal virtual {
         // default to no-op
     }
+
+    /// @dev Allow modules to get access to token handles within transaction time.
+    function _allowTokenHandleToThis(euint64 handle) internal virtual {
+        _allowTokenHandleToThis(handle, false);
+    }
+
+    /// @dev Allow modules to get access to token handles.
+    function _allowTokenHandleToThis(euint64 handle, bool persistent) internal virtual {
+        HandleAccessManager(_token).getHandleAllowance(euint64.unwrap(handle), address(this), persistent);
+    }
 }
diff --git a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
new file mode 100644
index 00000000..68898b41
--- /dev/null
+++ b/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
@@ -0,0 +1,238 @@
+import { callAndGetResult } from '../../../helpers/event';
+import { FhevmType } from '@fhevm/hardhat-plugin';
+import { expect } from 'chai';
+import { ethers, fhevm } from 'hardhat';
+
+const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
+const alwaysOnType = 0;
+const transferOnlyType = 1;
+const moduleTypes = [alwaysOnType, transferOnlyType];
+const maxInverstor = 2;
+const maxBalance = 100;
+
+const fixture = async () => {
+  const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
+  const token = (await ethers.deployContract('ERC7984RwaModularComplianceMock', ['name', 'symbol', 'uri'])).connect(
+    anyone,
+  );
+  await token.connect(admin).addAgent(agent1);
+  const investorCapModule = await ethers.deployContract('ERC7984RwaInvestorCapModuleMock', [
+    await token.getAddress(),
+    maxInverstor,
+  ]);
+  const balanceCapModule = await ethers.deployContract('ERC7984RwaBalanceCapModuleMock', [await token.getAddress()]);
+  const encryptedInput = await fhevm
+    .createEncryptedInput(await balanceCapModule.getAddress(), admin.address)
+    .add64(maxBalance)
+    .encrypt();
+  await balanceCapModule
+    .connect(admin)
+    ['setMaxBalance(bytes32,bytes)'](encryptedInput.handles[0], encryptedInput.inputProof);
+  return {
+    token,
+    investorCapModule,
+    balanceCapModule,
+    admin,
+    agent1,
+    agent2,
+    recipient,
+    anyone,
+  };
+};
+
+describe('ERC7984RwaModularCompliance', function () {
+  describe('Support module', async function () {
+    for (const type of moduleTypes) {
+      it(`should support module type ${type}`, async function () {
+        const { token } = await fixture();
+        await expect(token.supportsModule(type)).to.eventually.be.true;
+      });
+    }
+  });
+
+  describe('Instal module', async function () {
+    for (const type of moduleTypes) {
+      it(`should install module type ${type}`, async function () {
+        const { token, investorCapModule, admin } = await fixture();
+        await expect(token.connect(admin).installModule(type, investorCapModule))
+          .to.emit(token, 'ModuleInstalled')
+          .withArgs(type, investorCapModule);
+        await expect(token.isModuleInstalled(type, investorCapModule)).to.eventually.be.true;
+      });
+    }
+
+    it('should not install module if not admin or agent', async function () {
+      const { token, investorCapModule, anyone } = await fixture();
+      await expect(token.connect(anyone).installModule(alwaysOnType, investorCapModule))
+        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
+        .withArgs(anyone.address);
+    });
+
+    for (const type of moduleTypes) {
+      it('should not install module if not module', async function () {
+        const { token, admin } = await fixture();
+        const notModule = '0x0000000000000000000000000000000000000001';
+        await expect(token.connect(admin).installModule(type, notModule))
+          .to.be.revertedWithCustomError(token, 'ERC7984RwaNotTransferComplianceModule')
+          .withArgs(notModule);
+        await expect(token.isModuleInstalled(type, notModule)).to.eventually.be.false;
+      });
+    }
+
+    for (const type of moduleTypes) {
+      it(`should not install module type ${type} if already installed`, async function () {
+        const { token, investorCapModule, admin } = await fixture();
+        await token.connect(admin).installModule(type, investorCapModule);
+        await expect(token.connect(admin).installModule(type, investorCapModule))
+          .to.be.revertedWithCustomError(token, 'ERC7984RwaAlreadyInstalledModule')
+          .withArgs(type, await investorCapModule.getAddress());
+      });
+    }
+  });
+
+  describe('Uninstal module', async function () {
+    for (const type of moduleTypes) {
+      it(`should remove module type ${type}`, async function () {
+        const { token, investorCapModule, admin } = await fixture();
+        await token.connect(admin).installModule(type, investorCapModule);
+        await expect(token.connect(admin).uninstallModule(type, investorCapModule))
+          .to.emit(token, 'ModuleUninstalled')
+          .withArgs(type, investorCapModule);
+        await expect(token.isModuleInstalled(type, investorCapModule)).to.eventually.be.false;
+      });
+    }
+  });
+
+  describe('Modules', async function () {
+    for (const type of moduleTypes) {
+      it(`should transfer if compliant to balance cap module with type ${type}`, async function () {
+        const { token, admin, balanceCapModule, recipient, anyone } = await fixture();
+        await token.connect(admin).installModule(type, balanceCapModule);
+        const encryptedMint = await fhevm
+          .createEncryptedInput(await token.getAddress(), admin.address)
+          .add64(100)
+          .encrypt();
+        await token
+          .connect(admin)
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedMint.handles[0], encryptedMint.inputProof);
+        const amount = 25;
+        const encryptedTransferValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), recipient.address)
+          .add64(amount)
+          .encrypt();
+        const [, , transferredHandle] = await callAndGetResult(
+          token
+            .connect(recipient)
+            ['confidentialTransfer(address,bytes32,bytes)'](
+              anyone,
+              encryptedTransferValueInput.handles[0],
+              encryptedTransferValueInput.inputProof,
+            ),
+          transferEventSignature,
+        );
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+        ).to.eventually.equal(amount);
+        await expect(
+          fhevm.userDecryptEuint(
+            FhevmType.euint64,
+            await token.confidentialBalanceOf(recipient),
+            await token.getAddress(),
+            recipient,
+          ),
+        ).to.eventually.equal(75);
+      });
+    }
+
+    it(`should transfer zero if not compliant to balance cap module`, async function () {
+      const { token, admin, balanceCapModule, recipient, anyone } = await fixture();
+      await token.connect(admin).installModule(transferOnlyType, balanceCapModule);
+      const encryptedMint = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(100)
+        .encrypt();
+      await token
+        .connect(admin)
+        ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedMint.handles[0], encryptedMint.inputProof);
+      await token
+        .connect(admin)
+        ['confidentialMint(address,bytes32,bytes)'](anyone, encryptedMint.handles[0], encryptedMint.inputProof);
+      const amount = 25;
+      const encryptedTransferValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), recipient.address)
+        .add64(amount)
+        .encrypt();
+      const [, , transferredHandle] = await callAndGetResult(
+        token
+          .connect(recipient)
+          ['confidentialTransfer(address,bytes32,bytes)'](
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          ),
+        transferEventSignature,
+      );
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+      ).to.eventually.equal(0);
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await token.confidentialBalanceOf(recipient),
+          await token.getAddress(),
+          recipient,
+        ),
+        // balance is unchanged
+      ).to.eventually.equal(100);
+    });
+  });
+
+  for (const type of moduleTypes) {
+    it(`should transfer if compliant to investor cap module else zero with type ${type}`, async function () {
+      const { token, admin, investorCapModule, recipient, anyone } = await fixture();
+      await token.connect(admin).installModule(type, investorCapModule);
+      const encryptedMint = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
+        .add64(100)
+        .encrypt();
+      for (const investor of [
+        recipient.address, // investor#1
+        ethers.Wallet.createRandom().address, //investor#2
+      ]) {
+        await token
+          .connect(admin)
+          ['confidentialMint(address,bytes32,bytes)'](investor, encryptedMint.handles[0], encryptedMint.inputProof);
+      }
+      await expect(investorCapModule.getCurrentInvestor())
+        .to.eventually.equal(await investorCapModule.getMaxInvestor())
+        .to.equal(2);
+      const amount = 25;
+      const encryptedTransferValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), recipient.address)
+        .add64(amount)
+        .encrypt();
+      // trying to transfer to investor#3 (anyone) but number of investors is capped
+      const [, , transferredHandle] = await callAndGetResult(
+        token
+          .connect(recipient)
+          ['confidentialTransfer(address,bytes32,bytes)'](
+            anyone,
+            encryptedTransferValueInput.handles[0],
+            encryptedTransferValueInput.inputProof,
+          ),
+        transferEventSignature,
+      );
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+      ).to.eventually.equal(0);
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await token.confidentialBalanceOf(recipient),
+          await token.getAddress(),
+          recipient,
+        ),
+      ).to.eventually.equal(100);
+    });
+  }
+});

From b74c5ba80fe1a679dd3696828d0185817debbbeb Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 16 Sep 2025 11:56:47 +0200
Subject: [PATCH 45/52] Increase coverage on modular compliance flow

---
 .../token/ERC7984RwaComplianceModuleMock.sol  | 41 +++++++++
 .../ERC7984/extensions/ERC7984Freezable.sol   |  6 +-
 .../ERC7984RwaTransferComplianceModule.sol    |  2 +-
 .../extensions/ERC7984RwaCompliance.test.ts   | 90 +++++++++++++++++++
 4 files changed, 137 insertions(+), 2 deletions(-)
 create mode 100644 contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol

diff --git a/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol b/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
new file mode 100644
index 00000000..35b64213
--- /dev/null
+++ b/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
+import {ERC7984RwaTransferComplianceModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol";
+
+// solhint-disable func-name-mixedcase
+contract ERC7984RwaComplianceModuleMock is ERC7984RwaTransferComplianceModule, SepoliaConfig {
+    bool private _compliant = false;
+    string private _name;
+
+    event PostTransfer(string name);
+    event PreTransfer(string name);
+
+    constructor(address compliance, string memory name) ERC7984RwaTransferComplianceModule(compliance) {
+        _name = name;
+    }
+
+    function $_setCompliant() public {
+        _compliant = true;
+    }
+
+    function $_unsetCompliant() public {
+        _compliant = false;
+    }
+
+    function _isCompliantTransfer(
+        address /*from*/,
+        address /*to*/,
+        euint64 /*encryptedAmount*/
+    ) internal override returns (ebool) {
+        emit PreTransfer(_name);
+        return FHE.asEbool(_compliant);
+    }
+
+    function _postTransfer(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal override {
+        emit PostTransfer(_name);
+    }
+}
diff --git a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
index c1a61ab3..5794ada3 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
@@ -110,9 +110,13 @@ abstract contract ERC7984Freezable is ERC7984 {
         euint64 transferred = super._update(from, to, encryptedAmount);
         if (from != address(0) && _skipUpdateCheck) {
             // Reset frozen to balance if transferred more than available
+            euint64 frozen = confidentialFrozen(from);
+            if (!FHE.isInitialized(frozen)) {
+                frozen = FHE.asEuint64(0);
+            }
             _setConfidentialFrozen(
                 from,
-                FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), confidentialFrozen(from)),
+                FHE.select(FHE.gt(transferred, available), confidentialBalanceOf(from), frozen),
                 false
             );
         }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
index 13f79b01..af6a6ffc 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
@@ -48,7 +48,7 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
     }
 
     /// @inheritdoc IERC7984RwaTransferComplianceModule
-    function postTransfer(address from, address to, euint64 encryptedAmount) public virtual {
+    function postTransfer(address from, address to, euint64 encryptedAmount) public virtual onlyToken {
         _postTransfer(from, to, encryptedAmount);
     }
 
diff --git a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
index 68898b41..8acd5378 100644
--- a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
@@ -1,5 +1,6 @@
 import { callAndGetResult } from '../../../helpers/event';
 import { FhevmType } from '@fhevm/hardhat-plugin';
+import { time } from '@nomicfoundation/hardhat-network-helpers';
 import { expect } from 'chai';
 import { ethers, fhevm } from 'hardhat';
 
@@ -7,6 +8,8 @@ const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
 const alwaysOnType = 0;
 const transferOnlyType = 1;
 const moduleTypes = [alwaysOnType, transferOnlyType];
+const alwaysOn = 'always-on';
+const transferOnly = 'transfer-only';
 const maxInverstor = 2;
 const maxBalance = 100;
 
@@ -16,6 +19,14 @@ const fixture = async () => {
     anyone,
   );
   await token.connect(admin).addAgent(agent1);
+  const alwaysOnModule = await ethers.deployContract('ERC7984RwaComplianceModuleMock', [
+    await token.getAddress(),
+    alwaysOn,
+  ]);
+  const transferOnlyModule = await ethers.deployContract('ERC7984RwaComplianceModuleMock', [
+    await token.getAddress(),
+    transferOnly,
+  ]);
   const investorCapModule = await ethers.deployContract('ERC7984RwaInvestorCapModuleMock', [
     await token.getAddress(),
     maxInverstor,
@@ -30,6 +41,8 @@ const fixture = async () => {
     ['setMaxBalance(bytes32,bytes)'](encryptedInput.handles[0], encryptedInput.inputProof);
   return {
     token,
+    alwaysOnModule,
+    transferOnlyModule,
     investorCapModule,
     balanceCapModule,
     admin,
@@ -104,6 +117,82 @@ describe('ERC7984RwaModularCompliance', function () {
   });
 
   describe('Modules', async function () {
+    for (const forceTransfer of [false, true]) {
+      for (const compliant of [true, false]) {
+        it(`should ${forceTransfer ? 'force transfer' : 'transfer'} ${
+          compliant ? 'if' : 'zero if not'
+        } compliant`, async function () {
+          const { token, alwaysOnModule, transferOnlyModule, admin, recipient, anyone } = await fixture();
+          await token.connect(admin).installModule(alwaysOnType, alwaysOnModule);
+          await token.connect(admin).installModule(transferOnlyType, transferOnlyModule);
+          const amount = 100;
+          const encryptedMint = await fhevm
+            .createEncryptedInput(await token.getAddress(), admin.address)
+            .add64(amount)
+            .encrypt();
+          // set compliant for initial mint
+          await alwaysOnModule.$_setCompliant();
+          await transferOnlyModule.$_setCompliant();
+          await token
+            .connect(admin)
+            ['confidentialMint(address,bytes32,bytes)'](
+              recipient.address,
+              encryptedMint.handles[0],
+              encryptedMint.inputProof,
+            );
+          await expect(
+            fhevm.userDecryptEuint(
+              FhevmType.euint64,
+              await token.confidentialBalanceOf(recipient.address),
+              await token.getAddress(),
+              recipient,
+            ),
+          ).to.eventually.equal(amount);
+          const encryptedMint2 = await fhevm
+            .createEncryptedInput(await token.getAddress(), admin.address)
+            .add64(amount)
+            .encrypt();
+          if (compliant) {
+            await alwaysOnModule.$_setCompliant();
+            await transferOnlyModule.$_setCompliant();
+          } else {
+            await alwaysOnModule.$_unsetCompliant();
+            await transferOnlyModule.$_unsetCompliant();
+          }
+          if (!forceTransfer) {
+            await token.connect(recipient).setOperator(admin.address, (await time.latest()) + 1000);
+          }
+          const tx = token
+            .connect(admin)
+            [
+              forceTransfer
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'confidentialTransferFrom(address,address,bytes32,bytes)'
+            ](recipient.address, anyone.address, encryptedMint2.handles[0], encryptedMint2.inputProof);
+          const [, , transferredHandle] = await callAndGetResult(tx, transferEventSignature);
+          await expect(
+            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+          ).to.eventually.equal(compliant ? amount : 0);
+          await expect(tx)
+            .to.emit(alwaysOnModule, 'PreTransfer')
+            .withArgs(alwaysOn)
+            .to.emit(alwaysOnModule, 'PostTransfer')
+            .withArgs(alwaysOn);
+          if (forceTransfer) {
+            await expect(tx)
+              .to.not.emit(transferOnlyModule, 'PreTransfer')
+              .to.not.emit(transferOnlyModule, 'PostTransfer');
+          } else {
+            await expect(tx)
+              .to.emit(transferOnlyModule, 'PreTransfer')
+              .withArgs(transferOnly)
+              .to.emit(transferOnlyModule, 'PostTransfer')
+              .withArgs(transferOnly);
+          }
+        });
+      }
+    }
+
     for (const type of moduleTypes) {
       it(`should transfer if compliant to balance cap module with type ${type}`, async function () {
         const { token, admin, balanceCapModule, recipient, anyone } = await fixture();
@@ -233,6 +322,7 @@ describe('ERC7984RwaModularCompliance', function () {
           recipient,
         ),
       ).to.eventually.equal(100);
+      await expect(investorCapModule.getCurrentInvestor()).to.eventually.equal(3); //TODO: Should be 2
     });
   }
 });

From b964c8b2b0ba4b2cf638fc3140b43bab481d8da4 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 17 Sep 2025 10:32:19 +0200
Subject: [PATCH 46/52] Should not post update investors if not compliant

---
 .../token/ERC7984RwaInvestorCapModuleMock.sol |  2 +-
 .../token/ERC7984RwaModularComplianceMock.sol |  4 --
 .../rwa/ERC7984RwaBalanceCapModule.sol        |  7 +--
 .../rwa/ERC7984RwaInvestorCapModule.sol       | 58 +++++++++++--------
 .../ERC7984RwaTransferComplianceModule.sol    | 18 +++---
 .../extensions/ERC7984RwaCompliance.test.ts   | 25 +++++++-
 6 files changed, 71 insertions(+), 43 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol b/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol
index c2306d60..b3d5631a 100644
--- a/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol
+++ b/contracts/mocks/token/ERC7984RwaInvestorCapModuleMock.sol
@@ -6,5 +6,5 @@ import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
 import {ERC7984RwaInvestorCapModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol";
 
 contract ERC7984RwaInvestorCapModuleMock is ERC7984RwaInvestorCapModule, SepoliaConfig {
-    constructor(address compliance, uint256 maxInvestor) ERC7984RwaInvestorCapModule(compliance, maxInvestor) {}
+    constructor(address token, uint64 maxInvestor) ERC7984RwaInvestorCapModule(token, maxInvestor) {}
 }
diff --git a/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol b/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
index b3f1409f..ce36df3d 100644
--- a/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
+++ b/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
@@ -3,12 +3,8 @@
 pragma solidity ^0.8.24;
 
 import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
-import {FHE, ebool, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol";
-import {Impl} from "@fhevm/solidity/lib/Impl.sol";
 import {ERC7984Rwa} from "../../token/ERC7984/extensions/ERC7984Rwa.sol";
 import {ERC7984RwaCompliance} from "../../token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol";
-import {FHESafeMath} from "../../utils/FHESafeMath.sol";
-import {HandleAccessManager} from "../../utils/HandleAccessManager.sol";
 
 contract ERC7984RwaModularComplianceMock is ERC7984RwaCompliance, SepoliaConfig {
     constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index 4cf59907..7317ecb2 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -14,7 +14,6 @@ import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferCompliance
 abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModule {
     using EnumerableSet for *;
 
-    address private immutable _token;
     euint64 private _maxBalance;
 
     constructor(address token) ERC7984RwaTransferComplianceModule(token) {
@@ -47,10 +46,8 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModu
             return FHE.asEbool(true);
         }
         euint64 balance = IERC7984(_token).confidentialBalanceOf(to);
-        if (FHE.isInitialized(balance)) {
-            _allowTokenHandleToThis(balance);
-        }
-        _allowTokenHandleToThis(encryptedAmount);
+        _getTokenHandleAllowance(balance);
+        _getTokenHandleAllowance(encryptedAmount);
         (ebool increased, euint64 futureBalance) = FHESafeMath.tryIncrease(balance, encryptedAmount);
         compliant = FHE.and(increased, FHE.le(futureBalance, _maxBalance));
     }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index 4aa19a99..3558b116 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -3,35 +3,33 @@
 pragma solidity ^0.8.27;
 
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
-import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+import {IERC7984} from "../../../../interfaces/IERC7984.sol";
 import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
 
 /**
  * @dev A transfer compliance module for confidential Real World Assets (RWAs) which limits the number of investors.
  */
 abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceModule {
-    using EnumerableSet for *;
+    uint64 private _maxInvestor;
+    euint64 private _investors;
 
-    uint256 private _maxInvestor;
-    EnumerableSet.AddressSet private _investors;
-
-    constructor(address compliance, uint256 maxInvestor) ERC7984RwaTransferComplianceModule(compliance) {
+    constructor(address token, uint64 maxInvestor) ERC7984RwaTransferComplianceModule(token) {
         _maxInvestor = maxInvestor;
     }
 
     /// @dev Sets max number of investors.
-    function setMaxInvestor(uint256 maxInvestor) public virtual onlyTokenAdmin {
+    function setMaxInvestor(uint64 maxInvestor) public virtual onlyTokenAdmin {
         _maxInvestor = maxInvestor;
     }
 
     /// @dev Gets max number of investors.
-    function getMaxInvestor() public view virtual returns (uint256) {
+    function getMaxInvestor() public view virtual returns (uint64) {
         return _maxInvestor;
     }
 
     /// @dev Gets current number of investors.
-    function getCurrentInvestor() public view virtual returns (uint256) {
-        return _investors.length();
+    function getCurrentInvestor() public view virtual returns (euint64) {
+        return _investors;
     }
 
     /// @dev Internal function which checks if a transfer is compliant.
@@ -39,23 +37,35 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceMod
         address /*from*/,
         address to,
         euint64 encryptedAmount
-    ) internal override returns (ebool) {
-        if (
-            !FHE.isInitialized(encryptedAmount) || // no amount
-            to == address(0) || // or burning
-            _investors.contains(to) || // or already investor
-            _investors.length() < _maxInvestor // or not reached max investors limit
-        ) {
-            return FHE.asEbool(true);
-        }
-
-        return FHE.asEbool(false);
+    ) internal override returns (ebool compliant) {
+        euint64 balance = IERC7984(_token).confidentialBalanceOf(to);
+        _getTokenHandleAllowance(balance);
+        _getTokenHandleAllowance(encryptedAmount);
+        compliant = FHE.or(
+            FHE.or(
+                FHE.asEbool(
+                    to == address(0) || // return true if burning
+                        !FHE.isInitialized(encryptedAmount) // or no amount
+                ),
+                FHE.eq(encryptedAmount, FHE.asEuint64(0)) // or zero amount
+            ),
+            FHE.or(
+                FHE.gt(balance, FHE.asEuint64(0)), // or already investor
+                FHE.lt(_investors, FHE.asEuint64(_maxInvestor)) // or not reached max investors limit
+            )
+        );
     }
 
     /// @dev Internal function which performs operation after transfer.
-    function _postTransfer(address /*from*/, address to, euint64 /*encryptedAmount*/) internal override {
-        if (!_investors.contains(to)) {
-            _investors.add(to);
+    function _postTransfer(address /*from*/, address to, euint64 encryptedAmount) internal override {
+        euint64 balance = IERC7984(_token).confidentialBalanceOf(to);
+        _getTokenHandleAllowance(balance);
+        _getTokenHandleAllowance(encryptedAmount);
+        if (!FHE.isInitialized(_investors)) {
+            _investors = FHE.asEuint64(0);
         }
+        _investors = FHE.select(FHE.eq(balance, encryptedAmount), FHE.add(_investors, FHE.asEuint64(1)), _investors);
+        _investors = FHE.select(FHE.eq(balance, FHE.asEuint64(0)), FHE.sub(_investors, FHE.asEuint64(1)), _investors);
+        FHE.allowThis(_investors);
     }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
index af6a6ffc..1210a565 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
@@ -9,8 +9,8 @@ import {HandleAccessManager} from "../../../../utils/HandleAccessManager.sol";
 /**
  * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).
  */
-abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule {
-    address private immutable _token;
+abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule, HandleAccessManager {
+    address internal immutable _token;
 
     /// @dev The sender is not the token.
     error SenderNotToken(address account);
@@ -64,13 +64,17 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
         // default to no-op
     }
 
-    /// @dev Allow modules to get access to token handles within transaction time.
-    function _allowTokenHandleToThis(euint64 handle) internal virtual {
-        _allowTokenHandleToThis(handle, false);
+    /// @dev Allow modules to get access to token handles during transaction.
+    function _getTokenHandleAllowance(euint64 handle) internal virtual {
+        _getTokenHandleAllowance(handle, false);
     }
 
     /// @dev Allow modules to get access to token handles.
-    function _allowTokenHandleToThis(euint64 handle, bool persistent) internal virtual {
-        HandleAccessManager(_token).getHandleAllowance(euint64.unwrap(handle), address(this), persistent);
+    function _getTokenHandleAllowance(euint64 handle, bool persistent) internal virtual {
+        if (FHE.isInitialized(handle)) {
+            HandleAccessManager(_token).getHandleAllowance(euint64.unwrap(handle), address(this), persistent);
+        }
     }
+
+    function _validateHandleAllowance(bytes32 handle) internal view override onlyTokenAdmin {}
 }
diff --git a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
index 8acd5378..286e900b 100644
--- a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
@@ -292,7 +292,17 @@ describe('ERC7984RwaModularCompliance', function () {
           .connect(admin)
           ['confidentialMint(address,bytes32,bytes)'](investor, encryptedMint.handles[0], encryptedMint.inputProof);
       }
-      await expect(investorCapModule.getCurrentInvestor())
+      await investorCapModule
+        .connect(admin)
+        .getHandleAllowance(await investorCapModule.getCurrentInvestor(), admin.address, true);
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await investorCapModule.getCurrentInvestor(),
+          await investorCapModule.getAddress(),
+          admin,
+        ),
+      )
         .to.eventually.equal(await investorCapModule.getMaxInvestor())
         .to.equal(2);
       const amount = 25;
@@ -322,7 +332,18 @@ describe('ERC7984RwaModularCompliance', function () {
           recipient,
         ),
       ).to.eventually.equal(100);
-      await expect(investorCapModule.getCurrentInvestor()).to.eventually.equal(3); //TODO: Should be 2
+      // current investor should be unchanged
+      await investorCapModule
+        .connect(admin)
+        .getHandleAllowance(await investorCapModule.getCurrentInvestor(), admin.address, true);
+      await expect(
+        fhevm.userDecryptEuint(
+          FhevmType.euint64,
+          await investorCapModule.getCurrentInvestor(),
+          await investorCapModule.getAddress(),
+          admin,
+        ),
+      ).to.eventually.equal(2);
     });
   }
 });

From 3be8b00f25bf6e12a7ac07b6083f4e24120a940d Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 17 Sep 2025 10:42:42 +0200
Subject: [PATCH 47/52] Rename to `ModularCompliance` & `ComplianceModule`

---
 .changeset/wet-results-doubt.md               |  2 +-
 contracts/interfaces/IERC7984Rwa.sol          |  6 ++---
 .../token/ERC7984RwaComplianceModuleMock.sol  |  6 ++---
 .../token/ERC7984RwaModularComplianceMock.sol |  4 ++--
 .../rwa/ERC7984RwaBalanceCapModule.sol        |  6 ++---
 ...ule.sol => ERC7984RwaComplianceModule.sol} | 10 ++++-----
 .../rwa/ERC7984RwaInvestorCapModule.sol       |  6 ++---
 ...ce.sol => ERC7984RwaModularCompliance.sol} | 22 +++++++++----------
 ...ts => ERC7984RwaModularCompliance.test.ts} |  4 ++--
 9 files changed, 33 insertions(+), 33 deletions(-)
 rename contracts/token/ERC7984/extensions/rwa/{ERC7984RwaTransferComplianceModule.sol => ERC7984RwaComplianceModule.sol} (87%)
 rename contracts/token/ERC7984/extensions/rwa/{ERC7984RwaCompliance.sol => ERC7984RwaModularCompliance.sol} (89%)
 rename test/token/ERC7984/extensions/{ERC7984RwaCompliance.test.ts => ERC7984RwaModularCompliance.test.ts} (99%)

diff --git a/.changeset/wet-results-doubt.md b/.changeset/wet-results-doubt.md
index e4b5af97..4b731310 100644
--- a/.changeset/wet-results-doubt.md
+++ b/.changeset/wet-results-doubt.md
@@ -2,4 +2,4 @@
 'openzeppelin-confidential-contracts': minor
 ---
 
-`ERC7984RwaCompliance`: Support compliance modules for confidential RWAs.
+`ERC7984RwaModularCompliance`: Support compliance modules for confidential RWAs.
diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 93b965de..922b1eb5 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -87,8 +87,8 @@ interface IERC7984RwaBase {
 /// @dev Full interface for confidential RWA contracts.
 interface IERC7984Rwa is IERC7984, IERC7984RwaBase, IERC165, IAccessControl {}
 
-/// @dev Interface for confidential RWA compliance.
-interface IERC7984RwaCompliance {
+/// @dev Interface for confidential RWA with modular compliance.
+interface IERC7984RwaModularCompliance {
     enum ComplianceModuleType {
         ALWAYS_ON,
         TRANSFER_ONLY
@@ -103,7 +103,7 @@ interface IERC7984RwaCompliance {
 }
 
 /// @dev Interface for confidential RWA transfer compliance module.
-interface IERC7984RwaTransferComplianceModule {
+interface IERC7984RwaComplianceModule {
     /// @dev Returns magic number if it is a module.
     function isModule() external returns (bytes4);
     /// @dev Checks if a transfer is compliant. Should be non-mutating.
diff --git a/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol b/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
index 35b64213..955f3260 100644
--- a/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
+++ b/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
@@ -4,17 +4,17 @@ pragma solidity ^0.8.24;
 
 import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
-import {ERC7984RwaTransferComplianceModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol";
+import {ERC7984RwaComplianceModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaComplianceModule.sol";
 
 // solhint-disable func-name-mixedcase
-contract ERC7984RwaComplianceModuleMock is ERC7984RwaTransferComplianceModule, SepoliaConfig {
+contract ERC7984RwaModularComplianceModuleMock is ERC7984RwaComplianceModule, SepoliaConfig {
     bool private _compliant = false;
     string private _name;
 
     event PostTransfer(string name);
     event PreTransfer(string name);
 
-    constructor(address compliance, string memory name) ERC7984RwaTransferComplianceModule(compliance) {
+    constructor(address compliance, string memory name) ERC7984RwaComplianceModule(compliance) {
         _name = name;
     }
 
diff --git a/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol b/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
index ce36df3d..483c7e29 100644
--- a/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
+++ b/contracts/mocks/token/ERC7984RwaModularComplianceMock.sol
@@ -4,8 +4,8 @@ pragma solidity ^0.8.24;
 
 import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
 import {ERC7984Rwa} from "../../token/ERC7984/extensions/ERC7984Rwa.sol";
-import {ERC7984RwaCompliance} from "../../token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol";
+import {ERC7984RwaModularCompliance} from "../../token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol";
 
-contract ERC7984RwaModularComplianceMock is ERC7984RwaCompliance, SepoliaConfig {
+contract ERC7984RwaModularComplianceMock is ERC7984RwaModularCompliance, SepoliaConfig {
     constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index 7317ecb2..365f202a 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -6,17 +6,17 @@ import {FHE, ebool, euint64, externalEuint64} from "@fhevm/solidity/lib/FHE.sol"
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {IERC7984} from "../../../../interfaces/IERC7984.sol";
 import {FHESafeMath} from "../../../../utils/FHESafeMath.sol";
-import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
+import {ERC7984RwaComplianceModule} from "./ERC7984RwaComplianceModule.sol";
 
 /**
  * @dev A transfer compliance module for confidential Real World Assets (RWAs) which limits the balance of an investor.
  */
-abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaTransferComplianceModule {
+abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaComplianceModule {
     using EnumerableSet for *;
 
     euint64 private _maxBalance;
 
-    constructor(address token) ERC7984RwaTransferComplianceModule(token) {
+    constructor(address token) ERC7984RwaComplianceModule(token) {
         _token = token;
     }
 
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaComplianceModule.sol
similarity index 87%
rename from contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
rename to contracts/token/ERC7984/extensions/rwa/ERC7984RwaComplianceModule.sol
index 1210a565..577deba7 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaTransferComplianceModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaComplianceModule.sol
@@ -3,13 +3,13 @@
 pragma solidity ^0.8.27;
 
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
-import {IERC7984Rwa, IERC7984RwaTransferComplianceModule} from "../../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984Rwa, IERC7984RwaComplianceModule} from "../../../../interfaces/IERC7984Rwa.sol";
 import {HandleAccessManager} from "../../../../utils/HandleAccessManager.sol";
 
 /**
  * @dev A contract which allows to build a transfer compliance module for confidential Real World Assets (RWAs).
  */
-abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferComplianceModule, HandleAccessManager {
+abstract contract ERC7984RwaComplianceModule is IERC7984RwaComplianceModule, HandleAccessManager {
     address internal immutable _token;
 
     /// @dev The sender is not the token.
@@ -33,12 +33,12 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
         _token = token;
     }
 
-    /// @inheritdoc IERC7984RwaTransferComplianceModule
+    /// @inheritdoc IERC7984RwaComplianceModule
     function isModule() public pure override returns (bytes4) {
         return this.isModule.selector;
     }
 
-    /// @inheritdoc IERC7984RwaTransferComplianceModule
+    /// @inheritdoc IERC7984RwaComplianceModule
     function isCompliantTransfer(
         address from,
         address to,
@@ -47,7 +47,7 @@ abstract contract ERC7984RwaTransferComplianceModule is IERC7984RwaTransferCompl
         FHE.allow(compliant = _isCompliantTransfer(from, to, encryptedAmount), msg.sender);
     }
 
-    /// @inheritdoc IERC7984RwaTransferComplianceModule
+    /// @inheritdoc IERC7984RwaComplianceModule
     function postTransfer(address from, address to, euint64 encryptedAmount) public virtual onlyToken {
         _postTransfer(from, to, encryptedAmount);
     }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index 3558b116..4a5e463e 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -4,16 +4,16 @@ pragma solidity ^0.8.27;
 
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {IERC7984} from "../../../../interfaces/IERC7984.sol";
-import {ERC7984RwaTransferComplianceModule} from "./ERC7984RwaTransferComplianceModule.sol";
+import {ERC7984RwaComplianceModule} from "./ERC7984RwaComplianceModule.sol";
 
 /**
  * @dev A transfer compliance module for confidential Real World Assets (RWAs) which limits the number of investors.
  */
-abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaTransferComplianceModule {
+abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaComplianceModule {
     uint64 private _maxInvestor;
     euint64 private _investors;
 
-    constructor(address token, uint64 maxInvestor) ERC7984RwaTransferComplianceModule(token) {
+    constructor(address token, uint64 maxInvestor) ERC7984RwaComplianceModule(token) {
         _maxInvestor = maxInvestor;
     }
 
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol
similarity index 89%
rename from contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
rename to contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol
index f751007f..9c4c9db9 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.27;
 
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {IERC7984RwaCompliance, IERC7984RwaTransferComplianceModule} from "../../../../interfaces/IERC7984Rwa.sol";
+import {IERC7984RwaModularCompliance, IERC7984RwaComplianceModule} from "../../../../interfaces/IERC7984Rwa.sol";
 import {HandleAccessManager} from "../../../../utils/HandleAccessManager.sol";
 import {ERC7984Rwa} from "../ERC7984Rwa.sol";
 
@@ -12,7 +12,7 @@ import {ERC7984Rwa} from "../ERC7984Rwa.sol";
  * @dev Extension of {ERC7984Rwa} that supports compliance modules for confidential Real World Assets (RWAs).
  * Inspired by ERC-7579 modules.
  */
-abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, HandleAccessManager {
+abstract contract ERC7984RwaModularCompliance is ERC7984Rwa, IERC7984RwaModularCompliance, HandleAccessManager {
     using EnumerableSet for *;
 
     EnumerableSet.AddressSet private _alwaysOnModules;
@@ -47,7 +47,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
     }
 
     /**
-     * @inheritdoc IERC7984RwaCompliance
+     * @inheritdoc IERC7984RwaModularCompliance
      * @dev Consider gas footprint of the module before adding it since all modules will perform
      * all steps (pre-check, compliance check, post-hook) in a single transaction.
      */
@@ -55,12 +55,12 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
         _installModule(moduleType, module);
     }
 
-    /// @inheritdoc IERC7984RwaCompliance
+    /// @inheritdoc IERC7984RwaModularCompliance
     function uninstallModule(ComplianceModuleType moduleType, address module) public virtual onlyAdminOrAgent {
         _uninstallModule(moduleType, module);
     }
 
-    /// @inheritdoc IERC7984RwaCompliance
+    /// @inheritdoc IERC7984RwaModularCompliance
     function isModuleInstalled(ComplianceModuleType moduleType, address module) public view virtual returns (bool) {
         return _isModuleInstalled(moduleType, module);
     }
@@ -69,10 +69,10 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
     function _installModule(ComplianceModuleType moduleType, address module) internal virtual {
         require(supportsModule(moduleType), ERC7984RwaUnsupportedModuleType(moduleType));
         (bool success, bytes memory returnData) = module.staticcall(
-            abi.encodePacked(IERC7984RwaTransferComplianceModule.isModule.selector)
+            abi.encodePacked(IERC7984RwaComplianceModule.isModule.selector)
         );
         require(
-            success && bytes4(returnData) == IERC7984RwaTransferComplianceModule.isModule.selector,
+            success && bytes4(returnData) == IERC7984RwaComplianceModule.isModule.selector,
             ERC7984RwaNotTransferComplianceModule(module)
         );
 
@@ -143,7 +143,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
         for (uint256 i = 0; i < modulesLength; i++) {
             compliant = FHE.and(
                 compliant,
-                IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)
+                IERC7984RwaComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)
             );
         }
     }
@@ -163,7 +163,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
         for (uint256 i = 0; i < modulesLength; i++) {
             compliant = FHE.and(
                 compliant,
-                IERC7984RwaTransferComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)
+                IERC7984RwaComplianceModule(modules[i]).isCompliantTransfer(from, to, encryptedAmount)
             );
         }
     }
@@ -173,7 +173,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
         address[] memory modules = _alwaysOnModules.values();
         uint256 modulesLength = modules.length;
         for (uint256 i = 0; i < modulesLength; i++) {
-            IERC7984RwaTransferComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
+            IERC7984RwaComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
         }
     }
 
@@ -182,7 +182,7 @@ abstract contract ERC7984RwaCompliance is ERC7984Rwa, IERC7984RwaCompliance, Han
         address[] memory modules = _transferOnlyModules.values();
         uint256 modulesLength = modules.length;
         for (uint256 i = 0; i < modulesLength; i++) {
-            IERC7984RwaTransferComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
+            IERC7984RwaComplianceModule(modules[i]).postTransfer(from, to, encryptedAmount);
         }
     }
 
diff --git a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
similarity index 99%
rename from test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
rename to test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
index 286e900b..757e0a5f 100644
--- a/test/token/ERC7984/extensions/ERC7984RwaCompliance.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
@@ -19,11 +19,11 @@ const fixture = async () => {
     anyone,
   );
   await token.connect(admin).addAgent(agent1);
-  const alwaysOnModule = await ethers.deployContract('ERC7984RwaComplianceModuleMock', [
+  const alwaysOnModule = await ethers.deployContract('ERC7984RwaModularComplianceModuleMock', [
     await token.getAddress(),
     alwaysOn,
   ]);
-  const transferOnlyModule = await ethers.deployContract('ERC7984RwaComplianceModuleMock', [
+  const transferOnlyModule = await ethers.deployContract('ERC7984RwaModularComplianceModuleMock', [
     await token.getAddress(),
     transferOnly,
   ]);

From 1bc4c3cf95120a5afa007faf3b96675e93fd39c6 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 17 Sep 2025 15:18:37 +0200
Subject: [PATCH 48/52] Add balance cap module tests

---
 .../token/ERC7984RwaBalanceCapModuleMock.sol  |   9 +
 .../rwa/ERC7984RwaBalanceCapModule.sol        |  14 +-
 .../rwa/ERC7984RwaInvestorCapModule.sol       |   5 +-
 .../ERC7984RwaModularCompliance.test.ts       | 208 +++++++++++++-----
 4 files changed, 170 insertions(+), 66 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol b/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol
index 8b2610ca..fbf84e0d 100644
--- a/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol
+++ b/contracts/mocks/token/ERC7984RwaBalanceCapModuleMock.sol
@@ -3,8 +3,17 @@
 pragma solidity ^0.8.24;
 
 import {SepoliaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+import {FHE, euint64} from "@fhevm/solidity/lib/FHE.sol";
 import {ERC7984RwaBalanceCapModule} from "../../token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol";
 
 contract ERC7984RwaBalanceCapModuleMock is ERC7984RwaBalanceCapModule, SepoliaConfig {
+    event AmountEncrypted(euint64 amount);
+
     constructor(address compliance) ERC7984RwaBalanceCapModule(compliance) {}
+
+    function createEncryptedAmount(uint64 amount) public returns (euint64 encryptedAmount) {
+        FHE.allowThis(encryptedAmount = FHE.asEuint64(amount));
+        FHE.allow(encryptedAmount, msg.sender);
+        emit AmountEncrypted(encryptedAmount);
+    }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
index 365f202a..6a45bf18 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaBalanceCapModule.sol
@@ -16,22 +16,27 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaComplianceModule {
 
     euint64 private _maxBalance;
 
+    event MaxBalanceSet(euint64 newMaxBalance);
+
     constructor(address token) ERC7984RwaComplianceModule(token) {
         _token = token;
     }
 
     /// @dev Sets max balance of an investor with proof.
     function setMaxBalance(externalEuint64 maxBalance, bytes calldata inputProof) public virtual onlyTokenAdmin {
-        FHE.allowThis(_maxBalance = FHE.fromExternal(maxBalance, inputProof));
+        euint64 maxBalance_ = FHE.fromExternal(maxBalance, inputProof);
+        FHE.allowThis(_maxBalance = maxBalance_);
+        emit MaxBalanceSet(maxBalance_);
     }
 
     /// @dev Sets max balance of an investor.
     function setMaxBalance(euint64 maxBalance) public virtual onlyTokenAdmin {
         FHE.allowThis(_maxBalance = maxBalance);
+        emit MaxBalanceSet(maxBalance);
     }
 
     /// @dev Gets max balance of an investor.
-    function getMaxBalance() public virtual returns (euint64) {
+    function getMaxBalance() public view virtual returns (euint64) {
         return _maxBalance;
     }
 
@@ -41,9 +46,8 @@ abstract contract ERC7984RwaBalanceCapModule is ERC7984RwaComplianceModule {
         address to,
         euint64 encryptedAmount
     ) internal override returns (ebool compliant) {
-        if (!FHE.isInitialized(encryptedAmount) || to == address(0)) {
-            // if no amount or burning
-            return FHE.asEbool(true);
+        if (to == address(0)) {
+            return FHE.asEbool(true); // if burning
         }
         euint64 balance = IERC7984(_token).confidentialBalanceOf(to);
         _getTokenHandleAllowance(balance);
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index 4a5e463e..dd0160c3 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -43,10 +43,7 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaComplianceModule {
         _getTokenHandleAllowance(encryptedAmount);
         compliant = FHE.or(
             FHE.or(
-                FHE.asEbool(
-                    to == address(0) || // return true if burning
-                        !FHE.isInitialized(encryptedAmount) // or no amount
-                ),
+                FHE.asEbool(to == address(0)), // return true if burning
                 FHE.eq(encryptedAmount, FHE.asEuint64(0)) // or zero amount
             ),
             FHE.or(
diff --git a/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
index 757e0a5f..0bb500ed 100644
--- a/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
@@ -2,6 +2,7 @@ import { callAndGetResult } from '../../../helpers/event';
 import { FhevmType } from '@fhevm/hardhat-plugin';
 import { time } from '@nomicfoundation/hardhat-network-helpers';
 import { expect } from 'chai';
+import { BytesLike } from 'ethers';
 import { ethers, fhevm } from 'hardhat';
 
 const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
@@ -192,6 +193,68 @@ describe('ERC7984RwaModularCompliance', function () {
         });
       }
     }
+  });
+
+  describe('Balance cap module', async function () {
+    for (const withProof of [false, true]) {
+      it(`should set max balance ${withProof ? 'with proof' : ''}`, async function () {
+        const { admin, balanceCapModule } = await fixture();
+        let params = [] as unknown as [encryptedAmount: BytesLike, inputProof: BytesLike];
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await balanceCapModule.getAddress(), admin.address)
+            .add64(maxBalance + 100)
+            .encrypt();
+          params.push(handles[0], inputProof);
+        } else {
+          const [newBalance] = await callAndGetResult(
+            balanceCapModule.connect(admin).createEncryptedAmount(maxBalance + 100),
+            'AmountEncrypted(bytes32)',
+          );
+          params.push(newBalance);
+        }
+        await expect(
+          balanceCapModule
+            .connect(admin)
+            [withProof ? 'setMaxBalance(bytes32,bytes)' : 'setMaxBalance(bytes32)'](...params),
+        )
+          .to.emit(balanceCapModule, 'MaxBalanceSet')
+          .withArgs(params[0]);
+        await balanceCapModule
+          .connect(admin)
+          .getHandleAllowance(await balanceCapModule.getMaxBalance(), admin.address, true);
+        await expect(
+          fhevm.userDecryptEuint(
+            FhevmType.euint64,
+            await balanceCapModule.getMaxBalance(),
+            await balanceCapModule.getAddress(),
+            admin,
+          ),
+        ).to.eventually.equal(maxBalance + 100);
+      });
+    }
+    for (const withProof of [false, true]) {
+      it(`should not set max balance if not admin ${withProof ? 'with proof' : ''}`, async function () {
+        const { admin, balanceCapModule, anyone } = await fixture();
+        const [newBalance] = await callAndGetResult(
+          balanceCapModule.connect(admin).createEncryptedAmount(maxBalance + 100),
+          'AmountEncrypted(bytes32)',
+        );
+        const oldBalance = await balanceCapModule.getMaxBalance();
+        const params = [newBalance];
+        if (withProof) {
+          params.push('0x');
+        }
+        await expect(
+          balanceCapModule
+            .connect(anyone)
+            [withProof ? 'setMaxBalance(bytes32,bytes)' : 'setMaxBalance(bytes32)'](...params),
+        )
+          .to.be.revertedWithCustomError(balanceCapModule, 'SenderNotTokenAdmin')
+          .withArgs(anyone.address);
+        await expect(balanceCapModule.getMaxBalance()).to.eventually.equal(oldBalance);
+      });
+    }
 
     for (const type of moduleTypes) {
       it(`should transfer if compliant to balance cap module with type ${type}`, async function () {
@@ -264,66 +327,36 @@ describe('ERC7984RwaModularCompliance', function () {
       await expect(
         fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
       ).to.eventually.equal(0);
-      await expect(
-        fhevm.userDecryptEuint(
-          FhevmType.euint64,
-          await token.confidentialBalanceOf(recipient),
-          await token.getAddress(),
-          recipient,
-        ),
-        // balance is unchanged
-      ).to.eventually.equal(100);
     });
-  });
 
-  for (const type of moduleTypes) {
-    it(`should transfer if compliant to investor cap module else zero with type ${type}`, async function () {
-      const { token, admin, investorCapModule, recipient, anyone } = await fixture();
-      await token.connect(admin).installModule(type, investorCapModule);
+    it('should transfer if compliant because burning', async function () {
+      const { token, admin, balanceCapModule, recipient } = await fixture();
+      await token.connect(admin).installModule(alwaysOnType, balanceCapModule);
       const encryptedMint = await fhevm
         .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(100)
         .encrypt();
-      for (const investor of [
-        recipient.address, // investor#1
-        ethers.Wallet.createRandom().address, //investor#2
-      ]) {
-        await token
-          .connect(admin)
-          ['confidentialMint(address,bytes32,bytes)'](investor, encryptedMint.handles[0], encryptedMint.inputProof);
-      }
-      await investorCapModule
+      await token
         .connect(admin)
-        .getHandleAllowance(await investorCapModule.getCurrentInvestor(), admin.address, true);
-      await expect(
-        fhevm.userDecryptEuint(
-          FhevmType.euint64,
-          await investorCapModule.getCurrentInvestor(),
-          await investorCapModule.getAddress(),
-          admin,
-        ),
-      )
-        .to.eventually.equal(await investorCapModule.getMaxInvestor())
-        .to.equal(2);
+        ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedMint.handles[0], encryptedMint.inputProof);
       const amount = 25;
-      const encryptedTransferValueInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), recipient.address)
+      const encryptedBurnValueInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), admin.address)
         .add64(amount)
         .encrypt();
-      // trying to transfer to investor#3 (anyone) but number of investors is capped
       const [, , transferredHandle] = await callAndGetResult(
         token
-          .connect(recipient)
-          ['confidentialTransfer(address,bytes32,bytes)'](
-            anyone,
-            encryptedTransferValueInput.handles[0],
-            encryptedTransferValueInput.inputProof,
+          .connect(admin)
+          ['confidentialBurn(address,bytes32,bytes)'](
+            recipient,
+            encryptedBurnValueInput.handles[0],
+            encryptedBurnValueInput.inputProof,
           ),
         transferEventSignature,
       );
       await expect(
         fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
-      ).to.eventually.equal(0);
+      ).to.eventually.equal(25);
       await expect(
         fhevm.userDecryptEuint(
           FhevmType.euint64,
@@ -331,19 +364,80 @@ describe('ERC7984RwaModularCompliance', function () {
           await token.getAddress(),
           recipient,
         ),
-      ).to.eventually.equal(100);
-      // current investor should be unchanged
-      await investorCapModule
-        .connect(admin)
-        .getHandleAllowance(await investorCapModule.getCurrentInvestor(), admin.address, true);
-      await expect(
-        fhevm.userDecryptEuint(
-          FhevmType.euint64,
-          await investorCapModule.getCurrentInvestor(),
-          await investorCapModule.getAddress(),
-          admin,
-        ),
-      ).to.eventually.equal(2);
+      ).to.eventually.equal(75);
     });
-  }
+  });
+
+  describe('Investor cap module', async function () {
+    for (const type of moduleTypes) {
+      it(`should transfer if compliant to investor cap module else zero with type ${type}`, async function () {
+        const { token, admin, investorCapModule, recipient, anyone } = await fixture();
+        await token.connect(admin).installModule(type, investorCapModule);
+        const encryptedMint = await fhevm
+          .createEncryptedInput(await token.getAddress(), admin.address)
+          .add64(100)
+          .encrypt();
+        for (const investor of [
+          recipient.address, // investor#1
+          ethers.Wallet.createRandom().address, //investor#2
+        ]) {
+          await token
+            .connect(admin)
+            ['confidentialMint(address,bytes32,bytes)'](investor, encryptedMint.handles[0], encryptedMint.inputProof);
+        }
+        await investorCapModule
+          .connect(admin)
+          .getHandleAllowance(await investorCapModule.getCurrentInvestor(), admin.address, true);
+        await expect(
+          fhevm.userDecryptEuint(
+            FhevmType.euint64,
+            await investorCapModule.getCurrentInvestor(),
+            await investorCapModule.getAddress(),
+            admin,
+          ),
+        )
+          .to.eventually.equal(await investorCapModule.getMaxInvestor())
+          .to.equal(2);
+        const amount = 25;
+        const encryptedTransferValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), recipient.address)
+          .add64(amount)
+          .encrypt();
+        // trying to transfer to investor#3 (anyone) but number of investors is capped
+        const [, , transferredHandle] = await callAndGetResult(
+          token
+            .connect(recipient)
+            ['confidentialTransfer(address,bytes32,bytes)'](
+              anyone,
+              encryptedTransferValueInput.handles[0],
+              encryptedTransferValueInput.inputProof,
+            ),
+          transferEventSignature,
+        );
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+        ).to.eventually.equal(0);
+        await expect(
+          fhevm.userDecryptEuint(
+            FhevmType.euint64,
+            await token.confidentialBalanceOf(recipient),
+            await token.getAddress(),
+            recipient,
+          ),
+        ).to.eventually.equal(100);
+        // current investor should be unchanged
+        await investorCapModule
+          .connect(admin)
+          .getHandleAllowance(await investorCapModule.getCurrentInvestor(), admin.address, true);
+        await expect(
+          fhevm.userDecryptEuint(
+            FhevmType.euint64,
+            await investorCapModule.getCurrentInvestor(),
+            await investorCapModule.getAddress(),
+            admin,
+          ),
+        ).to.eventually.equal(2);
+      });
+    }
+  });
 });

From 1336085c9a0ba2aaa07af50e148121ebb1f69dfb Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 17 Sep 2025 15:38:34 +0200
Subject: [PATCH 49/52] Add max investor tests

---
 .../rwa/ERC7984RwaInvestorCapModule.sol       |  3 +++
 .../ERC7984RwaModularCompliance.test.ts       | 20 +++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
index dd0160c3..099f3cd9 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaInvestorCapModule.sol
@@ -13,6 +13,8 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaComplianceModule {
     uint64 private _maxInvestor;
     euint64 private _investors;
 
+    event MaxInvestorSet(uint64 maxInvestor);
+
     constructor(address token, uint64 maxInvestor) ERC7984RwaComplianceModule(token) {
         _maxInvestor = maxInvestor;
     }
@@ -20,6 +22,7 @@ abstract contract ERC7984RwaInvestorCapModule is ERC7984RwaComplianceModule {
     /// @dev Sets max number of investors.
     function setMaxInvestor(uint64 maxInvestor) public virtual onlyTokenAdmin {
         _maxInvestor = maxInvestor;
+        emit MaxInvestorSet(maxInvestor);
     }
 
     /// @dev Gets max number of investors.
diff --git a/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
index 0bb500ed..d8281db8 100644
--- a/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
@@ -439,5 +439,25 @@ describe('ERC7984RwaModularCompliance', function () {
         ).to.eventually.equal(2);
       });
     }
+
+    it('should set max investor', async function () {
+      const { token, admin, investorCapModule } = await fixture();
+      const newMaxInvestor = 100;
+      await token.connect(admin).installModule(alwaysOnType, investorCapModule);
+      await expect(investorCapModule.connect(admin).setMaxInvestor(newMaxInvestor))
+        .to.emit(investorCapModule, 'MaxInvestorSet')
+        .withArgs(newMaxInvestor);
+      await expect(investorCapModule.getMaxInvestor()).to.eventually.equal(newMaxInvestor);
+    });
+
+    it('should not set max investor if not admin', async function () {
+      const { token, admin, anyone, investorCapModule } = await fixture();
+      const newMaxInvestor = maxInverstor + 10;
+      await token.connect(admin).installModule(alwaysOnType, investorCapModule);
+      await expect(investorCapModule.connect(anyone).setMaxInvestor(newMaxInvestor))
+        .to.be.revertedWithCustomError(investorCapModule, 'SenderNotTokenAdmin')
+        .withArgs(anyone.address);
+      await expect(investorCapModule.getMaxInvestor()).to.eventually.equal(maxInverstor);
+    });
   });
 });

From 4f04222ec405e5ba888a1913709c52aa0a296b36 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Wed, 17 Sep 2025 18:35:50 +0200
Subject: [PATCH 50/52] Use agent for operations

---
 contracts/mocks/token/ERC7984RwaMock.sol      |  14 +-
 .../token/ERC7984/extensions/ERC7984Rwa.sol   |  61 +-
 .../ERC7984/extensions/ERC7984Rwa.test.ts     | 600 +++++++++---------
 3 files changed, 319 insertions(+), 356 deletions(-)

diff --git a/contracts/mocks/token/ERC7984RwaMock.sol b/contracts/mocks/token/ERC7984RwaMock.sol
index 0ae0200e..8267808f 100644
--- a/contracts/mocks/token/ERC7984RwaMock.sol
+++ b/contracts/mocks/token/ERC7984RwaMock.sol
@@ -14,18 +14,6 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
     mapping(address account => euint64 encryptedAmount) private _frozenBalances;
     bool public compliantTransfer;
 
-    // TODO: Move modifiers to `ERC7984Rwa` or remove from mock if useless
-    /// @dev Checks if the sender is an admin.
-    modifier onlyAdmin() {
-        require(isAdmin(_msgSender()), UnauthorizedSender(_msgSender()));
-        _;
-    }
-    /// @dev Checks if the sender is an agent.
-    modifier onlyAgent() {
-        require(isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
-        _;
-    }
-
     constructor(string memory name, string memory symbol, string memory tokenUri) ERC7984Rwa(name, symbol, tokenUri) {}
 
     function createEncryptedAmount(uint64 amount) public returns (euint64 encryptedAmount) {
@@ -53,5 +41,5 @@ contract ERC7984RwaMock is ERC7984Rwa, HandleAccessManager, SepoliaConfig {
         return compliantTransfer;
     }
 
-    function _validateHandleAllowance(bytes32 handle) internal view override onlyAdminOrAgent {}
+    function _validateHandleAllowance(bytes32 handle) internal view override onlyAgent {}
 }
diff --git a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
index 99c3ad09..1e70fefe 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Rwa.sol
@@ -28,14 +28,18 @@ abstract contract ERC7984Rwa is
 {
     bytes32 public constant AGENT_ROLE = keccak256("AGENT_ROLE");
 
-    /// @dev The caller account is not authorized to perform the operation.
-    error UnauthorizedSender(address account);
     /// @dev The transfer does not follow token compliance.
     error UncompliantTransfer(address from, address to, euint64 encryptedAmount);
 
-    /// @dev Checks if the sender is an admin or an agent.
-    modifier onlyAdminOrAgent() {
-        require(isAdmin(_msgSender()) || isAgent(_msgSender()), UnauthorizedSender(_msgSender()));
+    /// @dev Checks if the sender is an admin.
+    modifier onlyAdmin() {
+        _checkRole(DEFAULT_ADMIN_ROLE);
+        _;
+    }
+
+    /// @dev Checks if the sender is an agent.
+    modifier onlyAgent() {
+        _checkRole(AGENT_ROLE);
         _;
     }
 
@@ -52,12 +56,12 @@ abstract contract ERC7984Rwa is
     }
 
     /// @dev Pauses contract.
-    function pause() public virtual onlyAdminOrAgent {
+    function pause() public virtual onlyAgent {
         _pause();
     }
 
     /// @dev Unpauses contract.
-    function unpause() public virtual onlyAdminOrAgent {
+    function unpause() public virtual onlyAgent {
         _unpause();
     }
 
@@ -72,27 +76,27 @@ abstract contract ERC7984Rwa is
     }
 
     /// @dev Adds agent.
-    function addAgent(address account) public virtual onlyAdminOrAgent {
-        _addAgent(account);
+    function addAgent(address account) public virtual onlyAdmin {
+        _grantRole(AGENT_ROLE, account);
     }
 
     /// @dev Removes agent.
-    function removeAgent(address account) public virtual onlyAdminOrAgent {
-        _removeAgent(account);
+    function removeAgent(address account) public virtual onlyAdmin {
+        _revokeRole(AGENT_ROLE, account);
     }
 
     /// @dev Blocks a user account.
-    function blockUser(address account) public virtual onlyAdminOrAgent {
+    function blockUser(address account) public virtual onlyAgent {
         _blockUser(account);
     }
 
     /// @dev Unblocks a user account.
-    function unblockUser(address account) public virtual onlyAdminOrAgent {
+    function unblockUser(address account) public virtual onlyAgent {
         _allowUser(account);
     }
 
     /// @dev Sets confidential frozen with proof.
-    function setConfidentialFrozen(address account, euint64 encryptedAmount) public virtual onlyAdminOrAgent {
+    function setConfidentialFrozen(address account, euint64 encryptedAmount) public virtual onlyAgent {
         require(
             FHE.isAllowed(encryptedAmount, account),
             ERC7984UnauthorizedUseOfEncryptedAmount(encryptedAmount, msg.sender)
@@ -105,7 +109,7 @@ abstract contract ERC7984Rwa is
         address account,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual onlyAdminOrAgent {
+    ) public virtual onlyAgent {
         _setConfidentialFrozen(account, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
@@ -114,12 +118,12 @@ abstract contract ERC7984Rwa is
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual onlyAdminOrAgent returns (euint64) {
+    ) public virtual onlyAgent returns (euint64) {
         return _confidentialMint(to, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Mints confidential amount of tokens to account.
-    function confidentialMint(address to, euint64 encryptedAmount) public virtual onlyAdminOrAgent returns (euint64) {
+    function confidentialMint(address to, euint64 encryptedAmount) public virtual onlyAgent returns (euint64) {
         return _confidentialMint(to, encryptedAmount);
     }
 
@@ -128,15 +132,12 @@ abstract contract ERC7984Rwa is
         address account,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual onlyAdminOrAgent returns (euint64) {
+    ) public virtual onlyAgent returns (euint64) {
         return _confidentialBurn(account, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
     /// @dev Burns confidential amount of tokens from account.
-    function confidentialBurn(
-        address account,
-        euint64 encryptedAmount
-    ) public virtual onlyAdminOrAgent returns (euint64) {
+    function confidentialBurn(address account, euint64 encryptedAmount) public virtual onlyAgent returns (euint64) {
         return _confidentialBurn(account, encryptedAmount);
     }
 
@@ -146,7 +147,7 @@ abstract contract ERC7984Rwa is
         address to,
         externalEuint64 encryptedAmount,
         bytes calldata inputProof
-    ) public virtual onlyAdminOrAgent returns (euint64) {
+    ) public virtual onlyAgent returns (euint64) {
         return _forceConfidentialTransferFrom(from, to, FHE.fromExternal(encryptedAmount, inputProof));
     }
 
@@ -155,20 +156,10 @@ abstract contract ERC7984Rwa is
         address from,
         address to,
         euint64 encryptedAmount
-    ) public virtual onlyAdminOrAgent returns (euint64 transferred) {
+    ) public virtual onlyAgent returns (euint64 transferred) {
         return _forceConfidentialTransferFrom(from, to, encryptedAmount);
     }
 
-    /// @dev Internal function which adds an agent.
-    function _addAgent(address account) internal virtual {
-        _grantRole(AGENT_ROLE, account);
-    }
-
-    /// @dev Internal function which removes an agent.
-    function _removeAgent(address account) internal virtual {
-        _revokeRole(AGENT_ROLE, account);
-    }
-
     /// @dev Internal function which mints confidential amount of tokens to account.
     function _confidentialMint(address to, euint64 encryptedAmount) internal virtual returns (euint64) {
         return _mint(to, encryptedAmount);
@@ -208,7 +199,7 @@ abstract contract ERC7984Rwa is
      * @dev Internal function which reverts if `msg.sender` is not authorized as a freezer.
      * This freezer role is only granted to admin or agent.
      */
-    function _checkFreezer() internal override onlyAdminOrAgent {}
+    function _checkFreezer() internal override onlyAgent {}
 
     /// @dev Checks if a transfer follows token compliance.
     function _isCompliantTransfer(address from, address to, euint64 encryptedAmount) internal virtual returns (bool);
diff --git a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
index ae9de6fb..8a022b94 100644
--- a/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984Rwa.test.ts
@@ -13,6 +13,8 @@ import { ethers, fhevm } from 'hardhat';
 
 const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
 const frozenEventSignature = 'TokensFrozen(address,bytes32)';
+const adminRole = ethers.ZeroHash;
+const agentRole = ethers.id('AGENT_ROLE');
 
 const fixture = async () => {
   const [admin, agent1, agent2, recipient, anyone] = await ethers.getSigners();
@@ -45,28 +47,26 @@ describe('ERC7984Rwa', function () {
 
   describe('Pausable', async function () {
     it('should pause & unpause', async function () {
-      const { token, admin, agent1 } = await fixture();
-      for (const manager of [admin, agent1]) {
-        expect(await token.paused()).is.false;
-        await token.connect(manager).pause();
-        expect(await token.paused()).is.true;
-        await token.connect(manager).unpause();
-        expect(await token.paused()).is.false;
-      }
+      const { token, agent1 } = await fixture();
+      expect(await token.paused()).is.false;
+      await token.connect(agent1).pause();
+      expect(await token.paused()).is.true;
+      await token.connect(agent1).unpause();
+      expect(await token.paused()).is.false;
     });
 
-    it('should not pause if neither admin nor agent', async function () {
+    it('should not pause if not agent', async function () {
       const { token, anyone } = await fixture();
       await expect(token.connect(anyone).pause())
-        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone.address);
+        .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+        .withArgs(anyone.address, agentRole);
     });
 
-    it('should not unpause if neither admin nor agent', async function () {
+    it('should not unpause if not agent', async function () {
       const { token, anyone } = await fixture();
       await expect(token.connect(anyone).unpause())
-        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone.address);
+        .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+        .withArgs(anyone.address, agentRole);
     });
   });
 
@@ -78,95 +78,89 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should check/add/remove agent', async function () {
-      const { token, admin, agent1, agent2 } = await fixture();
-      for (const manager of [admin, agent1]) {
-        expect(await token.isAgent(agent2)).is.false;
-        await token.connect(manager).addAgent(agent2);
-        expect(await token.isAgent(agent2)).is.true;
-        await token.connect(manager).removeAgent(agent2);
-        expect(await token.isAgent(agent2)).is.false;
-      }
+      const { token, admin, agent2 } = await fixture();
+      expect(await token.isAgent(agent2)).is.false;
+      await token.connect(admin).addAgent(agent2);
+      expect(await token.isAgent(agent2)).is.true;
+      await token.connect(admin).removeAgent(agent2);
+      expect(await token.isAgent(agent2)).is.false;
     });
 
-    it('should not add agent if neither admin nor agent', async function () {
+    it('should not add agent if not admin', async function () {
       const { token, agent1, anyone } = await fixture();
       await expect(token.connect(anyone).addAgent(agent1))
-        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone.address);
+        .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+        .withArgs(anyone.address, adminRole);
     });
 
-    it('should not remove agent if neither admin nor agent', async function () {
+    it('should not remove agent if not admin', async function () {
       const { token, agent1, anyone } = await fixture();
       await expect(token.connect(anyone).removeAgent(agent1))
-        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone.address);
+        .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+        .withArgs(anyone.address, adminRole);
     });
   });
 
   describe('ERC7984Restricted', async function () {
     it('should block & unblock', async function () {
-      const { token, admin, agent1, recipient } = await fixture();
-      for (const manager of [admin, agent1]) {
-        await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
-        await token.connect(manager).blockUser(recipient);
-        await expect(token.isUserAllowed(recipient)).to.eventually.be.false;
-        await token.connect(manager).unblockUser(recipient);
-        await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
-      }
+      const { token, agent1, recipient } = await fixture();
+      await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
+      await token.connect(agent1).blockUser(recipient);
+      await expect(token.isUserAllowed(recipient)).to.eventually.be.false;
+      await token.connect(agent1).unblockUser(recipient);
+      await expect(token.isUserAllowed(recipient)).to.eventually.be.true;
     });
 
     for (const arg of [true, false]) {
-      it(`should not ${arg ? 'block' : 'unblock'} if neither admin nor agent`, async function () {
+      it(`should not ${arg ? 'block' : 'unblock'} if not agent`, async function () {
         const { token, anyone } = await fixture();
         await expect(token.connect(anyone)[arg ? 'blockUser' : 'unblockUser'](anyone))
-          .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-          .withArgs(anyone.address);
+          .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+          .withArgs(anyone.address, agentRole);
       });
     }
   });
 
   describe('Mintable', async function () {
     for (const withProof of [true, false]) {
-      it(`should mint by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient } = await fixture();
-        for (const manager of [admin, agent1]) {
-          const { token } = await fixture();
-          await token.$_setCompliantTransfer();
-          const amount = 100;
-          let params = [recipient.address] as unknown as [
-            account: AddressLike,
-            encryptedAmount: BytesLike,
-            inputProof: BytesLike,
-          ];
-          if (withProof) {
-            const { handles, inputProof } = await fhevm
-              .createEncryptedInput(await token.getAddress(), manager.address)
-              .add64(amount)
-              .encrypt();
-            params.push(handles[0], inputProof);
-          } else {
-            await token.connect(manager).createEncryptedAmount(amount);
-            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
-          }
-          const [, , transferredHandle] = await callAndGetResult(
-            token
-              .connect(manager)
-              [withProof ? 'confidentialMint(address,bytes32,bytes)' : 'confidentialMint(address,bytes32)'](...params),
-            transferEventSignature,
-          );
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
-          ).to.eventually.equal(amount);
-          const balanceHandle = await token.confidentialBalanceOf(recipient);
-          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(amount);
+      it(`should mint by agent ${withProof ? 'with proof' : ''}`, async function () {
+        const { agent1, recipient } = await fixture();
+        const { token } = await fixture();
+        await token.$_setCompliantTransfer();
+        const amount = 100;
+        let params = [recipient.address] as unknown as [
+          account: AddressLike,
+          encryptedAmount: BytesLike,
+          inputProof: BytesLike,
+        ];
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await token.getAddress(), agent1.address)
+            .add64(amount)
+            .encrypt();
+          params.push(handles[0], inputProof);
+        } else {
+          await token.connect(agent1).createEncryptedAmount(amount);
+          params.push(await token.connect(agent1).createEncryptedAmount.staticCall(amount));
         }
+        const [, , transferredHandle] = await callAndGetResult(
+          token
+            .connect(agent1)
+            [withProof ? 'confidentialMint(address,bytes32,bytes)' : 'confidentialMint(address,bytes32)'](...params),
+          transferEventSignature,
+        );
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+        ).to.eventually.equal(amount);
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token.connect(agent1).getHandleAllowance(balanceHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), agent1),
+        ).to.eventually.equal(amount);
       });
     }
 
-    it('should not mint if neither admin nor agent', async function () {
+    it('should not mint if not agent', async function () {
       const { token, recipient, anyone } = await fixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), anyone.address)
@@ -178,19 +172,19 @@ describe('ERC7984Rwa', function () {
           .connect(anyone)
           ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
-        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone.address);
+        .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+        .withArgs(anyone.address, agentRole);
     });
 
     it('should not mint if transfer not compliant', async function () {
-      const { token, admin, recipient } = await fixture();
+      const { token, agent1, recipient } = await fixture();
       const encryptedInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), admin.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(100)
         .encrypt();
       await expect(
         token
-          .connect(admin)
+          .connect(agent1)
           ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
@@ -198,15 +192,15 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not mint if paused', async function () {
-      const { token, admin, recipient } = await fixture();
-      await token.connect(admin).pause();
+      const { token, agent1, recipient } = await fixture();
+      await token.connect(agent1).pause();
       const encryptedInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), admin.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(100)
         .encrypt();
       await expect(
         token
-          .connect(admin)
+          .connect(agent1)
           ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       ).to.be.revertedWithCustomError(token, 'EnforcedPause');
     });
@@ -214,62 +208,56 @@ describe('ERC7984Rwa', function () {
 
   describe('Burnable', async function () {
     for (const withProof of [true, false]) {
-      it(`should burn by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient } = await fixture();
-        for (const manager of [admin, agent1]) {
-          const { token } = await fixture();
-          const encryptedInput = await fhevm
-            .createEncryptedInput(await token.getAddress(), manager.address)
-            .add64(100)
+      it(`should burn agent ${withProof ? 'with proof' : ''}`, async function () {
+        const { agent1, recipient } = await fixture();
+        const { token } = await fixture();
+        const encryptedInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), agent1.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer();
+        await token
+          .connect(agent1)
+          ['confidentialMint(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof);
+        const balanceBeforeHandle = await token.confidentialBalanceOf(recipient);
+        await token.connect(agent1).getHandleAllowance(balanceBeforeHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceBeforeHandle, await token.getAddress(), agent1),
+        ).to.eventually.greaterThan(0);
+        const amount = 100;
+        let params = [recipient.address] as unknown as [
+          account: AddressLike,
+          encryptedAmount: BytesLike,
+          inputProof: BytesLike,
+        ];
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await token.getAddress(), agent1.address)
+            .add64(amount)
             .encrypt();
-          await token.$_setCompliantTransfer();
-          await token
-            .connect(manager)
-            ['confidentialMint(address,bytes32,bytes)'](
-              recipient,
-              encryptedInput.handles[0],
-              encryptedInput.inputProof,
-            );
-          const balanceBeforeHandle = await token.confidentialBalanceOf(recipient);
-          await token.connect(manager).getHandleAllowance(balanceBeforeHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, balanceBeforeHandle, await token.getAddress(), manager),
-          ).to.eventually.greaterThan(0);
-          const amount = 100;
-          let params = [recipient.address] as unknown as [
-            account: AddressLike,
-            encryptedAmount: BytesLike,
-            inputProof: BytesLike,
-          ];
-          if (withProof) {
-            const { handles, inputProof } = await fhevm
-              .createEncryptedInput(await token.getAddress(), manager.address)
-              .add64(amount)
-              .encrypt();
-            params.push(handles[0], inputProof);
-          } else {
-            await token.connect(manager).createEncryptedAmount(amount);
-            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
-          }
-          const [, , transferredHandle] = await callAndGetResult(
-            token
-              .connect(manager)
-              [withProof ? 'confidentialBurn(address,bytes32,bytes)' : 'confidentialBurn(address,bytes32)'](...params),
-            transferEventSignature,
-          );
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
-          ).to.eventually.equal(amount);
-          const balanceHandle = await token.confidentialBalanceOf(recipient);
-          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(0);
+          params.push(handles[0], inputProof);
+        } else {
+          await token.connect(agent1).createEncryptedAmount(amount);
+          params.push(await token.connect(agent1).createEncryptedAmount.staticCall(amount));
         }
+        const [, , transferredHandle] = await callAndGetResult(
+          token
+            .connect(agent1)
+            [withProof ? 'confidentialBurn(address,bytes32,bytes)' : 'confidentialBurn(address,bytes32)'](...params),
+          transferEventSignature,
+        );
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
+        ).to.eventually.equal(amount);
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token.connect(agent1).getHandleAllowance(balanceHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), agent1),
+        ).to.eventually.equal(0);
       });
     }
 
-    it('should not burn if neither admin nor agent', async function () {
+    it('should not burn if not agent', async function () {
       const { token, recipient, anyone } = await fixture();
       const encryptedInput = await fhevm
         .createEncryptedInput(await token.getAddress(), anyone.address)
@@ -281,19 +269,19 @@ describe('ERC7984Rwa', function () {
           .connect(anyone)
           ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
-        .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-        .withArgs(anyone.address);
+        .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+        .withArgs(anyone.address, agentRole);
     });
 
     it('should not burn if transfer not compliant', async function () {
-      const { token, admin, recipient } = await fixture();
+      const { token, agent1, recipient } = await fixture();
       const encryptedInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), admin.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(100)
         .encrypt();
       await expect(
         token
-          .connect(admin)
+          .connect(agent1)
           ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       )
         .to.be.revertedWithCustomError(token, 'UncompliantTransfer')
@@ -301,15 +289,15 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not burn if paused', async function () {
-      const { token, admin, recipient } = await fixture();
-      await token.connect(admin).pause();
+      const { token, agent1, recipient } = await fixture();
+      await token.connect(agent1).pause();
       const encryptedInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), admin.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(100)
         .encrypt();
       await expect(
         token
-          .connect(admin)
+          .connect(agent1)
           ['confidentialBurn(address,bytes32,bytes)'](recipient, encryptedInput.handles[0], encryptedInput.inputProof),
       ).to.be.revertedWithCustomError(token, 'EnforcedPause');
     });
@@ -318,163 +306,159 @@ describe('ERC7984Rwa', function () {
   describe('Force transfer', async function () {
     for (const withProof of [true, false]) {
       it(`should force transfer by admin or agent ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient, anyone } = await fixture();
-        for (const manager of [admin, agent1]) {
-          const { token } = await fixture();
-          const encryptedMintValueInput = await fhevm
-            .createEncryptedInput(await token.getAddress(), manager.address)
-            .add64(100)
-            .encrypt();
-          await token.$_setCompliantTransfer();
-          await token
-            .connect(manager)
-            ['confidentialMint(address,bytes32,bytes)'](
-              recipient,
-              encryptedMintValueInput.handles[0],
-              encryptedMintValueInput.inputProof,
-            );
-          // set frozen (50 available and about to force transfer 25)
-          const encryptedFrozenValueInput = await fhevm
-            .createEncryptedInput(await token.getAddress(), manager.address)
-            .add64(50)
-            .encrypt();
-          await token
-            .connect(manager)
-            ['setConfidentialFrozen(address,bytes32,bytes)'](
-              recipient,
-              encryptedFrozenValueInput.handles[0],
-              encryptedFrozenValueInput.inputProof,
-            );
-          await token.$_unsetCompliantTransfer();
-          expect(await token.compliantTransfer()).to.be.false;
-          const amount = 25;
-          let params = [recipient.address, anyone.address] as unknown as [
-            from: AddressLike,
-            to: AddressLike,
-            encryptedAmount: BytesLike,
-            inputProof: BytesLike,
-          ];
-          if (withProof) {
-            const { handles, inputProof } = await fhevm
-              .createEncryptedInput(await token.getAddress(), manager.address)
-              .add64(amount)
-              .encrypt();
-            params.push(handles[0], inputProof);
-          } else {
-            await token.connect(manager).createEncryptedAmount(amount);
-            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
-          }
-          const [from, to, transferredHandle] = await callAndGetResult(
-            token
-              .connect(manager)
-              [
-                withProof
-                  ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
-                  : 'forceConfidentialTransferFrom(address,address,bytes32)'
-              ](...params),
-            transferEventSignature,
+        const { agent1, recipient, anyone } = await fixture();
+        const { token } = await fixture();
+        const encryptedMintValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), agent1.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer();
+        await token
+          .connect(agent1)
+          ['confidentialMint(address,bytes32,bytes)'](
+            recipient,
+            encryptedMintValueInput.handles[0],
+            encryptedMintValueInput.inputProof,
+          );
+        // set frozen (50 available and about to force transfer 25)
+        const encryptedFrozenValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), agent1.address)
+          .add64(50)
+          .encrypt();
+        await token
+          .connect(agent1)
+          ['setConfidentialFrozen(address,bytes32,bytes)'](
+            recipient,
+            encryptedFrozenValueInput.handles[0],
+            encryptedFrozenValueInput.inputProof,
           );
-          expect(from).equal(recipient.address);
-          expect(to).equal(anyone.address);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), anyone),
-          ).to.eventually.equal(amount);
-          const balanceHandle = await token.confidentialBalanceOf(recipient);
-          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(75);
-          const frozenHandle = await token.confidentialFrozen(recipient);
-          await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(50); // frozen is left unchanged
+        await token.$_unsetCompliantTransfer();
+        expect(await token.compliantTransfer()).to.be.false;
+        const amount = 25;
+        let params = [recipient.address, anyone.address] as unknown as [
+          from: AddressLike,
+          to: AddressLike,
+          encryptedAmount: BytesLike,
+          inputProof: BytesLike,
+        ];
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await token.getAddress(), agent1.address)
+            .add64(amount)
+            .encrypt();
+          params.push(handles[0], inputProof);
+        } else {
+          await token.connect(agent1).createEncryptedAmount(amount);
+          params.push(await token.connect(agent1).createEncryptedAmount.staticCall(amount));
         }
+        const [from, to, transferredHandle] = await callAndGetResult(
+          token
+            .connect(agent1)
+            [
+              withProof
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'forceConfidentialTransferFrom(address,address,bytes32)'
+            ](...params),
+          transferEventSignature,
+        );
+        expect(from).equal(recipient.address);
+        expect(to).equal(anyone.address);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), anyone),
+        ).to.eventually.equal(amount);
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token.connect(agent1).getHandleAllowance(balanceHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), agent1),
+        ).to.eventually.equal(75);
+        const frozenHandle = await token.confidentialFrozen(recipient);
+        await token.connect(agent1).getHandleAllowance(frozenHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), agent1),
+        ).to.eventually.equal(50); // frozen is left unchanged
       });
     }
 
     for (const withProof of [true, false]) {
       it(`should force transfer even if frozen ${withProof ? 'with proof' : ''}`, async function () {
-        const { admin, agent1, recipient, anyone } = await fixture();
-        for (const manager of [admin, agent1]) {
-          const { token } = await fixture();
-          const encryptedMintValueInput = await fhevm
-            .createEncryptedInput(await token.getAddress(), manager.address)
-            .add64(100)
-            .encrypt();
-          await token.$_setCompliantTransfer();
-          await token
-            .connect(manager)
-            ['confidentialMint(address,bytes32,bytes)'](
-              recipient,
-              encryptedMintValueInput.handles[0],
-              encryptedMintValueInput.inputProof,
-            );
-          // set frozen (only 20 available but about to force transfer 25)
-          const encryptedFrozenValueInput = await fhevm
-            .createEncryptedInput(await token.getAddress(), manager.address)
-            .add64(80)
-            .encrypt();
-          await token
-            .connect(manager)
-            ['setConfidentialFrozen(address,bytes32,bytes)'](
-              recipient,
-              encryptedFrozenValueInput.handles[0],
-              encryptedFrozenValueInput.inputProof,
-            );
-          // should force transfer even if not compliant
-          await token.$_unsetCompliantTransfer();
-          expect(await token.compliantTransfer()).to.be.false;
-          // should force transfer even if paused
-          await token.connect(manager).pause();
-          expect(await token.paused()).to.be.true;
-          const amount = 25;
-          let params = [recipient.address, anyone.address] as unknown as [
-            from: AddressLike,
-            to: AddressLike,
-            encryptedAmount: BytesLike,
-            inputProof: BytesLike,
-          ];
-          if (withProof) {
-            const { handles, inputProof } = await fhevm
-              .createEncryptedInput(await token.getAddress(), manager.address)
-              .add64(amount)
-              .encrypt();
-            params.push(handles[0], inputProof);
-          } else {
-            await token.connect(manager).createEncryptedAmount(amount);
-            params.push(await token.connect(manager).createEncryptedAmount.staticCall(amount));
-          }
-          const [account, frozenAmountHandle] = await callAndGetResult(
-            token
-              .connect(manager)
-              [
-                withProof
-                  ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
-                  : 'forceConfidentialTransferFrom(address,address,bytes32)'
-              ](...params),
-            frozenEventSignature,
+        const { agent1, recipient, anyone } = await fixture();
+        const { token } = await fixture();
+        const encryptedMintValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), agent1.address)
+          .add64(100)
+          .encrypt();
+        await token.$_setCompliantTransfer();
+        await token
+          .connect(agent1)
+          ['confidentialMint(address,bytes32,bytes)'](
+            recipient,
+            encryptedMintValueInput.handles[0],
+            encryptedMintValueInput.inputProof,
+          );
+        // set frozen (only 20 available but about to force transfer 25)
+        const encryptedFrozenValueInput = await fhevm
+          .createEncryptedInput(await token.getAddress(), agent1.address)
+          .add64(80)
+          .encrypt();
+        await token
+          .connect(agent1)
+          ['setConfidentialFrozen(address,bytes32,bytes)'](
+            recipient,
+            encryptedFrozenValueInput.handles[0],
+            encryptedFrozenValueInput.inputProof,
           );
-          expect(account).equal(recipient.address);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, frozenAmountHandle, await token.getAddress(), recipient),
-          ).to.eventually.equal(75);
-          const balanceHandle = await token.confidentialBalanceOf(recipient);
-          await token.connect(manager).getHandleAllowance(balanceHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(75);
-          const frozenHandle = await token.confidentialFrozen(recipient);
-          await token.connect(manager).getHandleAllowance(frozenHandle, manager, true);
-          await expect(
-            fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), manager),
-          ).to.eventually.equal(75); // frozen got reset to balance
+        // should force transfer even if not compliant
+        await token.$_unsetCompliantTransfer();
+        expect(await token.compliantTransfer()).to.be.false;
+        // should force transfer even if paused
+        await token.connect(agent1).pause();
+        expect(await token.paused()).to.be.true;
+        const amount = 25;
+        let params = [recipient.address, anyone.address] as unknown as [
+          from: AddressLike,
+          to: AddressLike,
+          encryptedAmount: BytesLike,
+          inputProof: BytesLike,
+        ];
+        if (withProof) {
+          const { handles, inputProof } = await fhevm
+            .createEncryptedInput(await token.getAddress(), agent1.address)
+            .add64(amount)
+            .encrypt();
+          params.push(handles[0], inputProof);
+        } else {
+          await token.connect(agent1).createEncryptedAmount(amount);
+          params.push(await token.connect(agent1).createEncryptedAmount.staticCall(amount));
         }
+        const [account, frozenAmountHandle] = await callAndGetResult(
+          token
+            .connect(agent1)
+            [
+              withProof
+                ? 'forceConfidentialTransferFrom(address,address,bytes32,bytes)'
+                : 'forceConfidentialTransferFrom(address,address,bytes32)'
+            ](...params),
+          frozenEventSignature,
+        );
+        expect(account).equal(recipient.address);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, frozenAmountHandle, await token.getAddress(), recipient),
+        ).to.eventually.equal(75);
+        const balanceHandle = await token.confidentialBalanceOf(recipient);
+        await token.connect(agent1).getHandleAllowance(balanceHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, balanceHandle, await token.getAddress(), agent1),
+        ).to.eventually.equal(75);
+        const frozenHandle = await token.confidentialFrozen(recipient);
+        await token.connect(agent1).getHandleAllowance(frozenHandle, agent1, true);
+        await expect(
+          fhevm.userDecryptEuint(FhevmType.euint64, frozenHandle, await token.getAddress(), agent1),
+        ).to.eventually.equal(75); // frozen got reset to balance
       });
     }
 
     for (const withProof of [true, false]) {
-      it(`should not force transfer if neither admin nor agent ${withProof ? 'with proof' : ''}`, async function () {
+      it(`should not force transfer if not agent ${withProof ? 'with proof' : ''}`, async function () {
         const { token, recipient, anyone } = await fixture();
         let params = [recipient.address, anyone.address] as unknown as [
           from: AddressLike,
@@ -502,14 +486,14 @@ describe('ERC7984Rwa', function () {
                 : 'forceConfidentialTransferFrom(address,address,bytes32)'
             ](...params),
         )
-          .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-          .withArgs(anyone.address);
+          .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+          .withArgs(anyone.address, agentRole);
       });
     }
 
     for (const withProof of [true, false]) {
       it(`should not force transfer if receiver blocked ${withProof ? 'with proof' : ''}`, async function () {
-        const { token, recipient, anyone } = await fixture();
+        const { token, agent1, recipient, anyone } = await fixture();
         let params = [recipient.address, anyone.address] as unknown as [
           from: AddressLike,
           to: AddressLike,
@@ -527,7 +511,7 @@ describe('ERC7984Rwa', function () {
           await token.connect(anyone).createEncryptedAmount(amount);
           params.push(await token.connect(anyone).createEncryptedAmount.staticCall(amount));
         }
-        await token.blockUser(anyone);
+        await token.connect(agent1).blockUser(anyone);
         await expect(
           token
             .connect(anyone)
@@ -537,22 +521,22 @@ describe('ERC7984Rwa', function () {
                 : 'forceConfidentialTransferFrom(address,address,bytes32)'
             ](...params),
         )
-          .to.be.revertedWithCustomError(token, 'UnauthorizedSender')
-          .withArgs(anyone.address);
+          .to.be.revertedWithCustomError(token, 'AccessControlUnauthorizedAccount')
+          .withArgs(anyone.address, agentRole);
       });
     }
   });
 
   describe('Transfer', async function () {
     it('should transfer', async function () {
-      const { token, admin: manager, recipient, anyone } = await fixture();
+      const { token, agent1, recipient, anyone } = await fixture();
       const encryptedMintValueInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), manager.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(100)
         .encrypt();
       await token.$_setCompliantTransfer();
       await token
-        .connect(manager)
+        .connect(agent1)
         ['confidentialMint(address,bytes32,bytes)'](
           recipient,
           encryptedMintValueInput.handles[0],
@@ -560,11 +544,11 @@ describe('ERC7984Rwa', function () {
         );
       // set frozen (50 available and about to transfer 25)
       const encryptedFrozenValueInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), manager.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(50)
         .encrypt();
       await token
-        .connect(manager)
+        .connect(agent1)
         ['setConfidentialFrozen(address,bytes32,bytes)'](
           recipient,
           encryptedFrozenValueInput.handles[0],
@@ -611,12 +595,12 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not transfer if paused', async function () {
-      const { token, admin: manager, recipient, anyone } = await fixture();
+      const { token, agent1, recipient, anyone } = await fixture();
       const encryptedTransferValueInput = await fhevm
         .createEncryptedInput(await token.getAddress(), recipient.address)
         .add64(25)
         .encrypt();
-      await token.connect(manager).pause();
+      await token.connect(agent1).pause();
       await expect(
         token
           .connect(recipient)
@@ -649,14 +633,14 @@ describe('ERC7984Rwa', function () {
     });
 
     it('should not transfer if frozen', async function () {
-      const { token, admin: manager, recipient, anyone } = await fixture();
+      const { token, agent1, recipient, anyone } = await fixture();
       const encryptedMintValueInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), manager.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(100)
         .encrypt();
       await token.$_setCompliantTransfer();
       await token
-        .connect(manager)
+        .connect(agent1)
         ['confidentialMint(address,bytes32,bytes)'](
           recipient,
           encryptedMintValueInput.handles[0],
@@ -664,11 +648,11 @@ describe('ERC7984Rwa', function () {
         );
       // set frozen (20 available but about to transfer 25)
       const encryptedFrozenValueInput = await fhevm
-        .createEncryptedInput(await token.getAddress(), manager.address)
+        .createEncryptedInput(await token.getAddress(), agent1.address)
         .add64(80)
         .encrypt();
       await token
-        .connect(manager)
+        .connect(agent1)
         ['setConfidentialFrozen(address,bytes32,bytes)'](
           recipient,
           encryptedFrozenValueInput.handles[0],
@@ -706,14 +690,14 @@ describe('ERC7984Rwa', function () {
 
     for (const arg of [true, false]) {
       it(`should not transfer if ${arg ? 'sender' : 'receiver'} blocked `, async function () {
-        const { token, admin: manager, recipient, anyone } = await fixture();
+        const { token, agent1, recipient, anyone } = await fixture();
         const account = arg ? recipient : anyone;
         await token.$_setCompliantTransfer();
         const encryptedInput = await fhevm
           .createEncryptedInput(await token.getAddress(), recipient.address)
           .add64(25)
           .encrypt();
-        await token.connect(manager).blockUser(account);
+        await token.connect(agent1).blockUser(account);
 
         await expect(
           token

From 53ec92670a56858c76fa1811201cab15d5d71be9 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 30 Sep 2025 11:56:41 +0200
Subject: [PATCH 51/52] Restore restricted and freezable

---
 contracts/interfaces/IERC7984Restricted.sol   | 23 --------------
 .../ERC7984/extensions/ERC7984Freezable.sol   |  9 +++---
 .../ERC7984/extensions/ERC7984Restricted.sol  | 31 ++++++++-----------
 3 files changed, 18 insertions(+), 45 deletions(-)
 delete mode 100644 contracts/interfaces/IERC7984Restricted.sol

diff --git a/contracts/interfaces/IERC7984Restricted.sol b/contracts/interfaces/IERC7984Restricted.sol
deleted file mode 100644
index c129460f..00000000
--- a/contracts/interfaces/IERC7984Restricted.sol
+++ /dev/null
@@ -1,23 +0,0 @@
-// SPDX-License-Identifier: MIT
-
-pragma solidity ^0.8.24;
-
-/// @dev Interface for contracts that implements user account transfer restrictions.
-interface IERC7984Restricted {
-    enum Restriction {
-        DEFAULT, // User has no explicit restriction
-        BLOCKED, // User is explicitly blocked
-        ALLOWED // User is explicitly allowed
-    }
-
-    /// @dev Emitted when a user account's restriction is updated.
-    event UserRestrictionUpdated(address indexed account, Restriction restriction);
-
-    /// @dev The operation failed because the user account is restricted.
-    error UserRestricted(address account);
-
-    /// @dev Returns the restriction of a user account.
-    function getRestriction(address account) external view returns (Restriction);
-    /// @dev Returns whether a user account is allowed to interact with the token.
-    function isUserAllowed(address account) external view returns (bool);
-}
diff --git a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
index ec86c2be..52cbfd76 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Freezable.sol
@@ -36,9 +36,6 @@ abstract contract ERC7984Freezable is ERC7984 {
             confidentialBalanceOf(account),
             confidentialFrozen(account)
         );
-        if (!FHE.isInitialized(unfrozen)) {
-            return unfrozen;
-        }
         return FHE.select(success, unfrozen, FHE.asEuint64(0));
     }
 
@@ -51,8 +48,12 @@ abstract contract ERC7984Freezable is ERC7984 {
     }
 
     /**
-     * @dev See {ERC7984-_update}. The `from` account must have sufficient unfrozen balance,
+     * @dev See {ERC7984-_update}.
+     *
+     * The `from` account must have sufficient unfrozen balance,
      * otherwise 0 tokens are transferred.
+     * The default freezing behavior can be changed (for a pass-through for instance) by overriding
+     * {confidentialAvailable}.
      */
     function _update(address from, address to, euint64 encryptedAmount) internal virtual override returns (euint64) {
         if (from != address(0)) {
diff --git a/contracts/token/ERC7984/extensions/ERC7984Restricted.sol b/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
index 86e5c4d3..886be4b2 100644
--- a/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
+++ b/contracts/token/ERC7984/extensions/ERC7984Restricted.sol
@@ -2,7 +2,6 @@
 
 pragma solidity ^0.8.27;
 
-import {IERC7984Restricted} from "../../../interfaces/IERC7984Restricted.sol";
 import {ERC7984, euint64} from "../ERC7984.sol";
 
 /**
@@ -14,10 +13,20 @@ import {ERC7984, euint64} from "../ERC7984.sol";
  * a blocklist. Developers can override {isUserAllowed} to check that `restriction == ALLOWED`
  * to implement an allowlist.
  */
-abstract contract ERC7984Restricted is ERC7984, IERC7984Restricted {
+abstract contract ERC7984Restricted is ERC7984 {
+    enum Restriction {
+        DEFAULT, // User has no explicit restriction
+        BLOCKED, // User is explicitly blocked
+        ALLOWED // User is explicitly allowed
+    }
+
     mapping(address account => Restriction) private _restrictions;
-    /// @dev Skips restriction checks in {_update}.
-    bool private _skipUpdateCheck;
+
+    /// @dev Emitted when a user account's restriction is updated.
+    event UserRestrictionUpdated(address indexed account, Restriction restriction);
+
+    /// @dev The operation failed because the user account is restricted.
+    error UserRestricted(address account);
 
     /// @dev Returns the restriction of a user account.
     function getRestriction(address account) public view virtual returns (Restriction) {
@@ -41,20 +50,6 @@ abstract contract ERC7984Restricted is ERC7984, IERC7984Restricted {
         return getRestriction(account) != Restriction.BLOCKED; // i.e. DEFAULT && ALLOWED
     }
 
-    /// @dev Internal function to skip update check. Check can be restored with {_restoreERC7984RestrictedUpdateCheck}.
-    function _disableERC7984RestrictedUpdateCheck() internal virtual {
-        if (!_skipUpdateCheck) {
-            _skipUpdateCheck = true;
-        }
-    }
-
-    /// @dev Internal function to restore update check previously disabled by {_disableERC7984RestrictedUpdateCheck}.
-    function _restoreERC7984RestrictedUpdateCheck() internal virtual {
-        if (_skipUpdateCheck) {
-            _skipUpdateCheck = false;
-        }
-    }
-
     /**
      * @dev See {ERC7984-_update}. Enforces transfer restrictions (excluding minting and burning).
      *

From d3240e1d0470b2500d9474fea7a0e66b7f14b736 Mon Sep 17 00:00:00 2001
From: James Toussaint <33313130+james-toussaint@users.noreply.github.com>
Date: Tue, 30 Sep 2025 15:40:35 +0200
Subject: [PATCH 52/52] Update styling

---
 contracts/interfaces/IERC7984Rwa.sol          |  8 ++---
 .../token/ERC7984RwaComplianceModuleMock.sol  | 13 +++-----
 .../rwa/ERC7984RwaModularCompliance.sol       | 32 ++++++++++---------
 .../ERC7984RwaModularCompliance.test.ts       | 16 ++--------
 4 files changed, 28 insertions(+), 41 deletions(-)

diff --git a/contracts/interfaces/IERC7984Rwa.sol b/contracts/interfaces/IERC7984Rwa.sol
index 6faf5b85..35e730f8 100644
--- a/contracts/interfaces/IERC7984Rwa.sol
+++ b/contracts/interfaces/IERC7984Rwa.sol
@@ -71,16 +71,16 @@ interface IERC7984Rwa is IERC7984, IERC165 {
 /// @dev Interface for confidential RWA with modular compliance.
 interface IERC7984RwaModularCompliance {
     enum ComplianceModuleType {
-        ALWAYS_ON,
-        TRANSFER_ONLY
+        AlwaysOn,
+        TransferOnly
     }
 
+    /// @dev Checks if a compliance module is installed.
+    function isModuleInstalled(ComplianceModuleType moduleType, address module) external view returns (bool);
     /// @dev Installs a transfer compliance module.
     function installModule(ComplianceModuleType moduleType, address module) external;
     /// @dev Uninstalls a transfer compliance module.
     function uninstallModule(ComplianceModuleType moduleType, address module) external;
-    /// @dev Checks if a compliance module is installed.
-    function isModuleInstalled(ComplianceModuleType moduleType, address module) external view returns (bool);
 }
 
 /// @dev Interface for confidential RWA transfer compliance module.
diff --git a/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol b/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
index 955f3260..73d5d163 100644
--- a/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
+++ b/contracts/mocks/token/ERC7984RwaComplianceModuleMock.sol
@@ -9,14 +9,11 @@ import {ERC7984RwaComplianceModule} from "../../token/ERC7984/extensions/rwa/ERC
 // solhint-disable func-name-mixedcase
 contract ERC7984RwaModularComplianceModuleMock is ERC7984RwaComplianceModule, SepoliaConfig {
     bool private _compliant = false;
-    string private _name;
 
-    event PostTransfer(string name);
-    event PreTransfer(string name);
+    event PostTransfer();
+    event PreTransfer();
 
-    constructor(address compliance, string memory name) ERC7984RwaComplianceModule(compliance) {
-        _name = name;
-    }
+    constructor(address compliance) ERC7984RwaComplianceModule(compliance) {}
 
     function $_setCompliant() public {
         _compliant = true;
@@ -31,11 +28,11 @@ contract ERC7984RwaModularComplianceModuleMock is ERC7984RwaComplianceModule, Se
         address /*to*/,
         euint64 /*encryptedAmount*/
     ) internal override returns (ebool) {
-        emit PreTransfer(_name);
+        emit PreTransfer();
         return FHE.asEbool(_compliant);
     }
 
     function _postTransfer(address /*from*/, address /*to*/, euint64 /*encryptedAmount*/) internal override {
-        emit PostTransfer(_name);
+        emit PostTransfer();
     }
 }
diff --git a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol
index f342fce8..79fe7717 100644
--- a/contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol
+++ b/contracts/token/ERC7984/extensions/rwa/ERC7984RwaModularCompliance.sol
@@ -43,7 +43,12 @@ abstract contract ERC7984RwaModularCompliance is ERC7984Rwa, IERC7984RwaModularC
      * * Force transfer compliance module
      */
     function supportsModule(ComplianceModuleType moduleType) public view virtual returns (bool) {
-        return moduleType == ComplianceModuleType.ALWAYS_ON || moduleType == ComplianceModuleType.TRANSFER_ONLY;
+        return moduleType == ComplianceModuleType.AlwaysOn || moduleType == ComplianceModuleType.TransferOnly;
+    }
+
+    /// @inheritdoc IERC7984RwaModularCompliance
+    function isModuleInstalled(ComplianceModuleType moduleType, address module) public view virtual returns (bool) {
+        return _isModuleInstalled(moduleType, module);
     }
 
     /**
@@ -60,9 +65,13 @@ abstract contract ERC7984RwaModularCompliance is ERC7984Rwa, IERC7984RwaModularC
         _uninstallModule(moduleType, module);
     }
 
-    /// @inheritdoc IERC7984RwaModularCompliance
-    function isModuleInstalled(ComplianceModuleType moduleType, address module) public view virtual returns (bool) {
-        return _isModuleInstalled(moduleType, module);
+    /// @dev Checks if a compliance module is installed.
+    function _isModuleInstalled(
+        ComplianceModuleType moduleType,
+        address module
+    ) internal view virtual returns (bool installed) {
+        if (moduleType == ComplianceModuleType.AlwaysOn) return _alwaysOnModules.contains(module);
+        if (moduleType == ComplianceModuleType.TransferOnly) return _transferOnlyModules.contains(module);
     }
 
     /// @dev Internal function which installs a transfer compliance module.
@@ -76,9 +85,9 @@ abstract contract ERC7984RwaModularCompliance is ERC7984Rwa, IERC7984RwaModularC
             ERC7984RwaNotTransferComplianceModule(module)
         );
 
-        if (moduleType == ComplianceModuleType.ALWAYS_ON) {
+        if (moduleType == ComplianceModuleType.AlwaysOn) {
             require(_alwaysOnModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleType, module));
-        } else if (moduleType == ComplianceModuleType.TRANSFER_ONLY) {
+        } else if (moduleType == ComplianceModuleType.TransferOnly) {
             require(_transferOnlyModules.add(module), ERC7984RwaAlreadyInstalledModule(moduleType, module));
         }
         emit ModuleInstalled(moduleType, module);
@@ -87,21 +96,14 @@ abstract contract ERC7984RwaModularCompliance is ERC7984Rwa, IERC7984RwaModularC
     /// @dev Internal function which uninstalls a transfer compliance module.
     function _uninstallModule(ComplianceModuleType moduleType, address module) internal virtual {
         require(supportsModule(moduleType), ERC7984RwaUnsupportedModuleType(moduleType));
-        if (moduleType == ComplianceModuleType.ALWAYS_ON) {
+        if (moduleType == ComplianceModuleType.AlwaysOn) {
             require(_alwaysOnModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleType, module));
-        } else if (moduleType == ComplianceModuleType.TRANSFER_ONLY) {
+        } else if (moduleType == ComplianceModuleType.TransferOnly) {
             require(_transferOnlyModules.remove(module), ERC7984RwaAlreadyUninstalledModule(moduleType, module));
         }
         emit ModuleUninstalled(moduleType, module);
     }
 
-    /// @dev Checks if a compliance module is installed.
-    function _isModuleInstalled(ComplianceModuleType moduleType, address module) internal view virtual returns (bool) {
-        if (moduleType == ComplianceModuleType.ALWAYS_ON) return _alwaysOnModules.contains(module);
-        if (moduleType == ComplianceModuleType.TRANSFER_ONLY) return _transferOnlyModules.contains(module);
-        return false;
-    }
-
     /**
      * @dev Updates confidential balances. It transfers zero if it does not follow
      * transfer compliance. Runs hooks after the transfer.
diff --git a/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
index f13ac7aa..52ffa2e9 100644
--- a/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
+++ b/test/token/ERC7984/extensions/ERC7984RwaModularCompliance.test.ts
@@ -9,8 +9,6 @@ const transferEventSignature = 'ConfidentialTransfer(address,address,bytes32)';
 const alwaysOnType = 0;
 const transferOnlyType = 1;
 const moduleTypes = [alwaysOnType, transferOnlyType];
-const alwaysOn = 'always-on';
-const transferOnly = 'transfer-only';
 const maxInverstor = 2;
 const maxBalance = 100;
 const adminRole = ethers.ZeroHash;
@@ -23,11 +21,9 @@ const fixture = async () => {
   await token.connect(admin).addAgent(agent1);
   const alwaysOnModule = await ethers.deployContract('ERC7984RwaModularComplianceModuleMock', [
     await token.getAddress(),
-    alwaysOn,
   ]);
   const transferOnlyModule = await ethers.deployContract('ERC7984RwaModularComplianceModuleMock', [
     await token.getAddress(),
-    transferOnly,
   ]);
   const investorCapModule = await ethers.deployContract('ERC7984RwaInvestorCapModuleMock', [
     await token.getAddress(),
@@ -175,21 +171,13 @@ describe('ERC7984RwaModularCompliance', function () {
           await expect(
             fhevm.userDecryptEuint(FhevmType.euint64, transferredHandle, await token.getAddress(), recipient),
           ).to.eventually.equal(compliant ? amount : 0);
-          await expect(tx)
-            .to.emit(alwaysOnModule, 'PreTransfer')
-            .withArgs(alwaysOn)
-            .to.emit(alwaysOnModule, 'PostTransfer')
-            .withArgs(alwaysOn);
+          await expect(tx).to.emit(alwaysOnModule, 'PreTransfer').to.emit(alwaysOnModule, 'PostTransfer');
           if (forceTransfer) {
             await expect(tx)
               .to.not.emit(transferOnlyModule, 'PreTransfer')
               .to.not.emit(transferOnlyModule, 'PostTransfer');
           } else {
-            await expect(tx)
-              .to.emit(transferOnlyModule, 'PreTransfer')
-              .withArgs(transferOnly)
-              .to.emit(transferOnlyModule, 'PostTransfer')
-              .withArgs(transferOnly);
+            await expect(tx).to.emit(transferOnlyModule, 'PreTransfer').to.emit(transferOnlyModule, 'PostTransfer');
           }
         });
       }