Add ERC721 and ERC1155 support

README.md

@@ -1,8 +1,8 @@
 # ERC20Pods
 
-[![Build Status](https://github.com/1inch/erc20-pods/workflows/CI/badge.svg)](https://github.com/1inch/erc20-pods/actions)
-[![Coverage Status](https://codecov.io/gh/1inch/erc20-pods/branch/master/graph/badge.svg?token=Z3D5O3XUYV)](https://codecov.io/gh/1inch/erc20-pods)
-[![NPM Package](https://img.shields.io/npm/v/@1inch/erc20-pods.svg)](https://www.npmjs.org/package/@1inch/erc20-pods)
+[![Build Status](https://github.com/1inch/token-pods/workflows/CI/badge.svg)](https://github.com/1inch/token-pods/actions)
+[![Coverage Status](https://codecov.io/gh/1inch/token-pods/branch/master/graph/badge.svg?token=Z3D5O3XUYV)](https://codecov.io/gh/1inch/token-pods)
+[![NPM Package](https://img.shields.io/npm/v/@1inch/token-pods.svg)](https://www.npmjs.org/package/@1inch/token-pods)
 
 ERC20 extension enabling external smart contract based Pods to track balances of those users who opted-in to these Pods.
 

contracts/ERC1155Pods.sol

@@ -0,0 +1,85 @@
+// 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 ZeroPodsLimit();
+    error PodsLimitReachedForAccount();
+
+    uint256 public immutable podsLimit;
+    uint256 public immutable podCallGasLimit;
+
+    ReentrancyGuardLib.Data private _guard;
+    mapping(uint256 => TokenPodsLib.Data) private _pods;
+
+    constructor(uint256 podsLimit_, uint256 podCallGasLimit_) {
+        if (podsLimit_ == 0) revert ZeroPodsLimit();
+        podsLimit = podsLimit_;
+        podCallGasLimit = podCallGasLimit_;
+        _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, super.balanceOf(msg.sender, id));
+    }
+
+    function addPod(address pod, uint256 id) public virtual {
+        if (_pods[id].addPod(msg.sender, pod, balanceOf(msg.sender, id), podCallGasLimit) > podsLimit) revert PodsLimitReachedForAccount();
+    }
+
+    function removePod(address pod, uint256 id) public virtual {
+        _pods[id].removePod(msg.sender, pod, balanceOf(msg.sender, id), podCallGasLimit);
+    }
+
+    function removeAllPods(uint256 id) public virtual {
+        _pods[id].removeAllPods(msg.sender, balanceOf(msg.sender, id), podCallGasLimit);
+    }
+
+    // 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[i].updateBalancesWithTokenId(from, to, amounts[i], ids[i], podCallGasLimit);
+            }
+        }
+    }
+}

contracts/ERC20Pods.sol

@@ -6,26 +6,21 @@ 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();
     error ZeroPodsLimit();
+    error PodsLimitReachedForAccount();
 
     uint256 public immutable podsLimit;
     uint256 public immutable podCallGasLimit;
 
     ReentrancyGuardLib.Data private _guard;
-    mapping(address => AddressSet.Data) private _pods;
+    TokenPodsLib.Data private _pods;
 
     constructor(uint256 podsLimit_, uint256 podCallGasLimit_) {
         if (podsLimit_ == 0) revert ZeroPodsLimit();
@@ -35,141 +30,45 @@ abstract contract ERC20Pods is ERC20, IERC20Pods, ReentrancyGuardExt {
     }
 
     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) virtual returns(uint256) {
         return super.balanceOf(account);
     }
 
     function podBalanceOf(address pod, address account) public nonReentrantView(_guard) view virtual returns(uint256) {
-        if (hasPod(account, pod)) {
-            return super.balanceOf(account);
-        }
-        return 0;
+        return _pods.podBalanceOf(account, pod, super.balanceOf(account));
     }
 
     function addPod(address pod) public virtual {
-        _addPod(msg.sender, pod);
+        if (_pods.addPod(msg.sender, pod, balanceOf(msg.sender), podCallGasLimit) > podsLimit) revert PodsLimitReachedForAccount();
     }
 
     function removePod(address pod) public virtual {
-        _removePod(msg.sender, pod);
+        _pods.removePod(msg.sender, pod, balanceOf(msg.sender), podCallGasLimit);
     }
 
     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();
-
-        emit PodAdded(account, pod);
-        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();
-
-        emit PodRemoved(account, pod);
-        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--) {
-                _pods[account].remove(items[i - 1]);
-                emit PodRemoved(account, items[i - 1]);
-                if (balance > 0) {
-                    _updateBalances(items[i - 1], account, address(0), balance);
-                }
-            }
-        }
-    }
-
-    /// @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;
-        uint256 gasLimit = podCallGasLimit;
-        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), gasLimit) {
-                mstore(0, exception)
-                revert(0, 4)
-            }
-            pop(call(gasLimit, pod, 0, ptr, 0x64, 0, 0))
-        }
+        _pods.removeAllPods(msg.sender, balanceOf(msg.sender), podCallGasLimit);
     }
 
     // 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, podCallGasLimit);
     }
 }

contracts/ERC721Pods.sol

@@ -0,0 +1,75 @@
+
+// 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 ZeroPodsLimit();
+    error PodsLimitReachedForAccount();
+
+    uint256 public immutable podsLimit;
+    uint256 public immutable podCallGasLimit;
+
+    ReentrancyGuardLib.Data private _guard;
+    TokenPodsLib.Data private _pods;
+
+    constructor(uint256 podsLimit_, uint256 podCallGasLimit_) {
+        if (podsLimit_ == 0) revert ZeroPodsLimit();
+        podsLimit = podsLimit_;
+        podCallGasLimit = podCallGasLimit_;
+        _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, super.balanceOf(account));
+    }
+
+    function addPod(address pod) public virtual {
+        if (_pods.addPod(msg.sender, pod, balanceOf(msg.sender), podCallGasLimit) > podsLimit) revert PodsLimitReachedForAccount();
+    }
+
+    function removePod(address pod) public virtual {
+        _pods.removePod(msg.sender, pod, balanceOf(msg.sender), podCallGasLimit);
+    }
+
+    function removeAllPods() public virtual {
+        _pods.removeAllPods(msg.sender, balanceOf(msg.sender), podCallGasLimit);
+    }
+
+    // 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, podCallGasLimit);
+    }
+}

contracts/Pod.sol

@@ -3,25 +3,37 @@
 pragma solidity ^0.8.0;
 
 import "./interfaces/IPod.sol";
+import "./interfaces/IPodWithId.sol";
 import "./interfaces/IERC20Pods.sol";
 
-abstract contract Pod is IPod {
+abstract contract Pod is IPod, IPodWithId {
     error AccessDenied();
 
     IERC20Pods public immutable token;
+    uint256 public immutable tokenId;
 
     modifier onlyToken {
         if (msg.sender != address(token)) revert AccessDenied();
         _;
     }
 
-    constructor(IERC20Pods token_) {
+    modifier onlyTokenId(uint256 id) {
+        if (id != tokenId) revert AccessDenied();
+        _;
+    }
+
+    constructor(IERC20Pods token_, uint256 tokenId_) {
         token = token_;
+        tokenId = tokenId_;
     }
 
     function updateBalances(address from, address to, uint256 amount) external onlyToken {
         _updateBalances(from, to, amount);
     }
 
+    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) internal virtual;
 }