Add support for ERC721 and ERC1155

Author: k06a

contracts/ERC1155Pods.sol

@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC1155/ERC1155.sol";
+
+import "./interfaces/IERC1155Pods.sol";
+import "./TokenPodsLib.sol";
+import "./libs/ReentrancyGuard.sol";
+
+abstract contract ERC1155Pods is ERC1155, IERC1155Pods, ReentrancyGuardExt {
+    using TokenPodsLib for TokenPodsLib.Data;
+    using ReentrancyGuardLib for ReentrancyGuardLib.Data;
+
+    error PodsLimitReachedForAccount();
+
+    uint256 public immutable podsLimit;
+
+    ReentrancyGuardLib.Data private _guard;
+    mapping(uint256 => TokenPodsLib.Data) private _pods;
+
+    constructor(uint256 podsLimit_) {
+        podsLimit = podsLimit_;
+        _guard.init();
+    }
+
+    function hasPod(address account, address pod, uint256 id) public view virtual returns(bool) {
+        return _pods[id].hasPod(account, pod);
+    }
+
+    function podsCount(address account, uint256 id) public view virtual returns(uint256) {
+        return _pods[id].podsCount(account);
+    }
+
+    function podAt(address account, uint256 index, uint256 id) public view virtual returns(address) {
+        return _pods[id].podAt(account, index);
+    }
+
+    function pods(address account, uint256 id) public view virtual returns(address[] memory) {
+        return _pods[id].pods(account);
+    }
+
+    function balanceOf(address account, uint256 id) public nonReentrantView(_guard) view override(IERC1155, ERC1155) virtual returns(uint256) {
+        return super.balanceOf(account, id);
+    }
+
+    function podBalanceOf(address pod, address account, uint256 id) public nonReentrantView(_guard) view returns(uint256) {
+        return _pods[id].podBalanceOf(account, pod, balanceOf(msg.sender, id));
+    }
+
+    function addPod(address pod, uint256 id) public virtual {
+        if (_pods[id].addPod(msg.sender, pod, balanceOf(msg.sender, id)) > podsLimit) revert PodsLimitReachedForAccount();
+    }
+
+    function removePod(address pod, uint256 id) public virtual {
+        _pods[id].removePod(msg.sender, pod, balanceOf(msg.sender, id));
+    }
+
+    function removeAllPods(uint256 id) public virtual {
+        _pods[id].removeAllPods(msg.sender, balanceOf(msg.sender, id));
+    }
+
+    // ERC1155 Overrides
+
+    function _afterTokenTransfer(
+        address operator,
+        address from,
+        address to,
+        uint256[] memory ids,
+        uint256[] memory amounts,
+        bytes memory data
+    ) internal nonReentrant(_guard) override virtual {
+        super._afterTokenTransfer(operator, from, to, ids, amounts, data);
+
+        unchecked {
+            for (uint256 i = 0; i < ids.length; i++) {
+                _pods[ids[i]].updateBalancesWithTokenId(from, to, amounts[i], ids[i]);
+            }
+        }
+    }
+}

contracts/ERC20Pods.sol

@@ -6,164 +6,65 @@ import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "@1inch/solidity-utils/contracts/libraries/AddressSet.sol";
 
 import "./interfaces/IERC20Pods.sol";
-import "./interfaces/IPod.sol";
+import "./TokenPodsLib.sol";
 import "./libs/ReentrancyGuard.sol";
 
 abstract contract ERC20Pods is ERC20, IERC20Pods, ReentrancyGuardExt {
-    using AddressSet for AddressSet.Data;
-    using AddressArray for AddressArray.Data;
+    using TokenPodsLib for TokenPodsLib.Data;
     using ReentrancyGuardLib for ReentrancyGuardLib.Data;
 
-    error PodAlreadyAdded();
-    error PodNotFound();
-    error InvalidPodAddress();
     error PodsLimitReachedForAccount();
-    error InsufficientGas();
-
-    uint256 private constant _POD_CALL_GAS_LIMIT = 200_000;
 
     uint256 public immutable podsLimit;
 
     ReentrancyGuardLib.Data private _guard;
-    mapping(address => AddressSet.Data) private _pods;
+    TokenPodsLib.Data private _pods;
 
     constructor(uint256 podsLimit_) {
         podsLimit = podsLimit_;
         _guard.init();
     }
 
     function hasPod(address account, address pod) public view virtual returns(bool) {
-        return _pods[account].contains(pod);
+        return _pods.hasPod(account, pod);
     }
 
     function podsCount(address account) public view virtual returns(uint256) {
-        return _pods[account].length();
+        return _pods.podsCount(account);
     }
 
     function podAt(address account, uint256 index) public view virtual returns(address) {
-        return _pods[account].at(index);
+        return _pods.podAt(account, index);
     }
 
     function pods(address account) public view virtual returns(address[] memory) {
-        return _pods[account].items.get();
+        return _pods.pods(account);
     }
 
-    function balanceOf(address account) public nonReentrantView(_guard) view override(IERC20, ERC20) returns(uint256) {
+    function balanceOf(address account) public nonReentrantView(_guard) view override(IERC20, ERC20) virtual returns(uint256) {
         return super.balanceOf(account);
     }
 
-    function podBalanceOf(address pod, address account) public nonReentrantView(_guard) view returns(uint256) {
-        if (hasPod(account, pod)) {
-            return balanceOf(account);
-        }
-        return 0;
+    function podBalanceOf(address pod, address account) public nonReentrantView(_guard) view virtual returns(uint256) {
+        return _pods.podBalanceOf(account, pod, balanceOf(account));
     }
 
     function addPod(address pod) public virtual {
-        _addPod(msg.sender, pod);
+        if (_pods.addPod(msg.sender, pod, balanceOf(msg.sender)) > podsLimit) revert PodsLimitReachedForAccount();
     }
 
     function removePod(address pod) public virtual {
-        _removePod(msg.sender, pod);
+        _pods.removePod(msg.sender, pod, balanceOf(msg.sender));
     }
 
     function removeAllPods() public virtual {
-        _removeAllPods(msg.sender);
-    }
-
-    function _addPod(address account, address pod) internal virtual {
-        if (pod == address(0)) revert InvalidPodAddress();
-        if (!_pods[account].add(pod)) revert PodAlreadyAdded();
-        if (_pods[account].length() > podsLimit) revert PodsLimitReachedForAccount();
-
-        uint256 balance = balanceOf(account);
-        if (balance > 0) {
-            _updateBalances(pod, address(0), account, balance);
-        }
-    }
-
-    function _removePod(address account, address pod) internal virtual {
-        if (!_pods[account].remove(pod)) revert PodNotFound();
-
-        uint256 balance = balanceOf(account);
-        if (balance > 0) {
-            _updateBalances(pod, account, address(0), balance);
-        }
-    }
-
-    function _removeAllPods(address account) internal virtual {
-        address[] memory items = _pods[account].items.get();
-        uint256 balance = balanceOf(account);
-        unchecked {
-            for (uint256 i = items.length; i > 0; i--) {
-                if (balance > 0) {
-                    _updateBalances(items[i - 1], account, address(0), balance);
-                }
-                _pods[account].remove(items[i - 1]);
-            }
-        }
-    }
-
-    /// @notice Assembly implementation of the gas limited call to avoid return gas bomb,
-    // moreover call to a destructed pod would also revert even inside try-catch block in Solidity 0.8.17
-    /// @dev try IPod(pod).updateBalances{gas: _POD_CALL_GAS_LIMIT}(from, to, amount) {} catch {}
-    function _updateBalances(address pod, address from, address to, uint256 amount) private {
-        bytes4 selector = IPod.updateBalances.selector;
-        bytes4 exception = InsufficientGas.selector;
-        assembly {  // solhint-disable-line no-inline-assembly
-            let ptr := mload(0x40)
-            mstore(ptr, selector)
-            mstore(add(ptr, 0x04), from)
-            mstore(add(ptr, 0x24), to)
-            mstore(add(ptr, 0x44), amount)
-
-            if lt(div(mul(gas(), 63), 64), _POD_CALL_GAS_LIMIT) {
-                mstore(0, exception)
-                revert(0, 4)
-            }
-            pop(call(_POD_CALL_GAS_LIMIT, pod, 0, ptr, 0x64, 0, 0))
-        }
+        _pods.removeAllPods(msg.sender, balanceOf(msg.sender));
     }
 
     // ERC20 Overrides
 
     function _afterTokenTransfer(address from, address to, uint256 amount) internal nonReentrant(_guard) override virtual {
         super._afterTokenTransfer(from, to, amount);
-
-        unchecked {
-            if (amount > 0 && from != to) {
-                address[] memory a = _pods[from].items.get();
-                address[] memory b = _pods[to].items.get();
-                uint256 aLength = a.length;
-                uint256 bLength = b.length;
-
-                for (uint256 i = 0; i < aLength; i++) {
-                    address pod = a[i];
-
-                    uint256 j;
-                    for (j = 0; j < bLength; j++) {
-                        if (pod == b[j]) {
-                            // Both parties are participating of the same Pod
-                            _updateBalances(pod, from, to, amount);
-                            b[j] = address(0);
-                            break;
-                        }
-                    }
-
-                    if (j == bLength) {
-                        // Sender is participating in a Pod, but receiver is not
-                        _updateBalances(pod, from, address(0), amount);
-                    }
-                }
-
-                for (uint256 j = 0; j < bLength; j++) {
-                    address pod = b[j];
-                    if (pod != address(0)) {
-                        // Receiver is participating in a Pod, but sender is not
-                        _updateBalances(pod, address(0), to, amount);
-                    }
-                }
-            }
-        }
+        _pods.updateBalances(from, to, amount);
     }
 }

contracts/ERC721Pods.sol

@@ -0,0 +1,71 @@
+
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC721/ERC721.sol";
+import "@1inch/solidity-utils/contracts/libraries/AddressSet.sol";
+
+import "./interfaces/IERC721Pods.sol";
+import "./TokenPodsLib.sol";
+import "./libs/ReentrancyGuard.sol";
+
+abstract contract ERC721Pods is ERC721, IERC721Pods, ReentrancyGuardExt {
+    using TokenPodsLib for TokenPodsLib.Data;
+    using ReentrancyGuardLib for ReentrancyGuardLib.Data;
+
+    error PodsLimitReachedForAccount();
+
+    uint256 public immutable podsLimit;
+
+    ReentrancyGuardLib.Data private _guard;
+    TokenPodsLib.Data private _pods;
+
+    constructor(uint256 podsLimit_) {
+        podsLimit = podsLimit_;
+        _guard.init();
+    }
+
+    function hasPod(address account, address pod) public view virtual returns(bool) {
+        return _pods.hasPod(account, pod);
+    }
+
+    function podsCount(address account) public view virtual returns(uint256) {
+        return _pods.podsCount(account);
+    }
+
+    function podAt(address account, uint256 index) public view virtual returns(address) {
+        return _pods.podAt(account, index);
+    }
+
+    function pods(address account) public view virtual returns(address[] memory) {
+        return _pods.pods(account);
+    }
+
+    function balanceOf(address account) public nonReentrantView(_guard) view override(IERC721, ERC721) virtual returns(uint256) {
+        return super.balanceOf(account);
+    }
+
+    function podBalanceOf(address pod, address account) public nonReentrantView(_guard) view virtual returns(uint256) {
+        return _pods.podBalanceOf(account, pod, balanceOf(account));
+    }
+
+    function addPod(address pod) public virtual {
+        if (_pods.addPod(msg.sender, pod, balanceOf(msg.sender)) > podsLimit) revert PodsLimitReachedForAccount();
+    }
+
+    function removePod(address pod) public virtual {
+        _pods.removePod(msg.sender, pod, balanceOf(msg.sender));
+    }
+
+    function removeAllPods() public virtual {
+        _pods.removeAllPods(msg.sender, balanceOf(msg.sender));
+    }
+
+    // ERC721 Overrides
+
+    function _afterTokenTransfer(address from, address to, uint256 firstTokenId, uint256 batchSize) internal nonReentrant(_guard) override virtual {
+        super._afterTokenTransfer(from, to, firstTokenId, batchSize);
+        _pods.updateBalances(from, to, batchSize);
+    }
+}

contracts/Pod.sol

@@ -8,13 +8,30 @@ abstract contract Pod is IPod {
     error AccessDenied();
 
     address public immutable token;
+    uint256 public immutable tokenId;
 
     modifier onlyToken {
         if (msg.sender != token) revert AccessDenied();
         _;
     }
 
-    constructor(address token_) {
+    modifier onlyTokenId(uint256 id) {
+        if (id != tokenId) revert AccessDenied();
+        _;
+    }
+
+    constructor(address token_, uint256 tokenId_) {
         token = token_;
+        tokenId = tokenId_;
+    }
+
+    function updateBalancesWithTokenId(address from, address to, uint256 amount, uint256 id) external onlyToken onlyTokenId(id) {
+        _updateBalances(from, to, amount);
+    }
+
+    function updateBalances(address from, address to, uint256 amount) external onlyToken {
+        _updateBalances(from, to, amount);
     }
+
+    function _updateBalances(address from, address to, uint256 amount) internal virtual;
 }