Make associate more lenient

Signed-off-by: Luis Mastrangelo <luis@swirldslabs.com>

Author: acuarica

contracts/HtsSystemContract.sol

@@ -584,14 +584,14 @@ contract HtsSystemContract is IHederaTokenService {
             if (selector == this.associateToken.selector) {
                 require(msg.data.length >= 48, "associateToken: Not enough calldata");
                 address account = address(bytes20(msg.data[40:60]));
-                bytes32 slot = _isAssociatedSlot(account);
+                bytes32 slot = _isAssociatedSlot(account, false);
                 assembly { sstore(slot, true) }
                 return abi.encode(HederaResponseCodes.SUCCESS);
             }
             if (selector == this.dissociateToken.selector) {
                 require(msg.data.length >= 48, "dissociateToken: Not enough calldata");
                 address account = address(bytes20(msg.data[40:60]));
-                bytes32 slot = _isAssociatedSlot(account);
+                bytes32 slot = _isAssociatedSlot(account, false);
                 assembly { sstore(slot, false) }
                 return abi.encode(HederaResponseCodes.SUCCESS);
             }
@@ -797,17 +797,17 @@ contract HtsSystemContract is IHederaTokenService {
 
     function _redirectForHRC719(bytes4 selector) private returns (bytes memory) {
         if (selector == IHRC719.associate.selector) {
-            bytes32 slot = _isAssociatedSlot(msg.sender);
+            bytes32 slot = _isAssociatedSlot(msg.sender, false);
             assembly { sstore(slot, true) }
             return abi.encode(true);
         }
         if (selector == IHRC719.dissociate.selector) {
-            bytes32 slot = _isAssociatedSlot(msg.sender);
+            bytes32 slot = _isAssociatedSlot(msg.sender, false);
             assembly { sstore(slot, false) }
             return abi.encode(true);
         }
         if (selector == IHRC719.isAssociated.selector) {
-            bytes32 slot = _isAssociatedSlot(msg.sender);
+            bytes32 slot = _isAssociatedSlot(msg.sender, true);
             bool res;
             assembly { res := sload(slot) }
             return abi.encode(res);
@@ -892,11 +892,11 @@ contract HtsSystemContract is IHederaTokenService {
         return bytes32(abi.encodePacked(selector, pad, spenderId, ownerId));
     }
 
-    function _isAssociatedSlot(address account) internal virtual returns (bytes32) {
+    function _isAssociatedSlot(address account, bool revertIfNotExists) internal virtual returns (bytes32) {
         bytes4 selector = IHRC719.isAssociated.selector;
         uint192 pad = 0x0;
         (uint32 accountId, bool exists) = HtsSystemContract(HTS_ADDRESS).getAccountId(account);
-        require(exists);
+        require(!revertIfNotExists || exists, "_isAssociatedSlot: account does not exist");
         return bytes32(abi.encodePacked(selector, pad, accountId));
     }
 

contracts/HtsSystemContractJson.sol

@@ -487,8 +487,8 @@ contract HtsSystemContractJson is HtsSystemContract {
         return slot;
     }
 
-    function _isAssociatedSlot(address account) internal override returns (bytes32) {
-        bytes32 slot = super._isAssociatedSlot(account);
+    function _isAssociatedSlot(address account, bool revertIfNotExists) internal override returns (bytes32) {
+        bytes32 slot = super._isAssociatedSlot(account, revertIfNotExists);
         if (_shouldFetch(slot)) {
             bool associated = mirrorNode().isAssociated(address(this), account);
             _setValue(slot, bytes32(uint256(associated ? 1 : 0)));

examples/foundry-hts/script/ApproveToken.s.sol

@@ -7,32 +7,24 @@ import {Hsc} from "hedera-forking/Hsc.sol";
 import {Approver} from "../src/Approver.sol";
 
 /**
- * Usage
- *
- * forge script ApproveTokenScript --sig "deployApprover()" --rpc-url testnet --broadcast
+ * Usage 
  * 
- * Notice that the `--skip-simulation` is not needed because there is no HTS involved in this step.
- * 
- * forge script ApproveTokenScript --sig "run(address,address)" $TOKEN_ADDR $APPROVER --rpc-url testnet --broadcast --skip-simulation
- * 
- * where $APPROVER is the address returned in the previous step.
+ * forge script ApproveTokenScript --sig "run(address)" $TOKEN_ADDR --rpc-url testnet --broadcast --skip-simulation --slow
  */
 contract ApproveTokenScript is Script {
-    uint256 private PRIVATE_KEY = vm.envUint("PRIVATE_KEY");
+    function run(address tokenAddress) external returns (Approver approver, address payable account) {
+        Hsc.htsSetup();
+
+        uint256 PRIVATE_KEY = vm.envUint("PRIVATE_KEY");
 
-    function deployApprover() external returns (Approver approver) {
         vm.startBroadcast(PRIVATE_KEY);
         approver = new Approver();
         vm.stopBroadcast();
-    }
-
-    function run(address tokenAddress, address approver) external returns (address payable account) {
-        Hsc.htsSetup();
 
         vm.startBroadcast(PRIVATE_KEY);
         Approver(approver).associate(tokenAddress);
 
-        IERC20(tokenAddress).transfer(approver, 50000);
+        IERC20(tokenAddress).transfer(address(approver), 50000);
 
         VmSafe.Wallet memory wallet = vm.createWallet("wallet");
         account = payable(wallet.addr);