From a036899a599c7fd687e8396744c5fcbb42a68fb6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ernani=20S=C3=A3o=20Thiago?= Date: Tue, 7 Nov 2023 17:26:06 -0500 Subject: [PATCH] storing all TreeNodes in the same mapping --- contracts/RolesRegistry/SftRolesRegistry.sol | 72 +++++++++---------- .../RolesRegistry/interfaces/IERCXXXX.sol | 5 +- .../RolesRegistry/libraries/BinaryTrees.sol | 34 ++++----- 3 files changed, 53 insertions(+), 58 deletions(-) diff --git a/contracts/RolesRegistry/SftRolesRegistry.sol b/contracts/RolesRegistry/SftRolesRegistry.sol index 6bdd1e2..264dee3 100644 --- a/contracts/RolesRegistry/SftRolesRegistry.sol +++ b/contracts/RolesRegistry/SftRolesRegistry.sol @@ -10,21 +10,20 @@ import { IERC1155 } from "@openzeppelin/contracts/token/ERC1155/IERC1155.sol"; import { IERC1155Receiver } from "@openzeppelin/contracts/token/ERC1155/IERC1155Receiver.sol"; import { ERC1155Holder, ERC1155Receiver } from "@openzeppelin/contracts/token/ERC1155/utils/ERC1155Holder.sol"; import { ERC165Checker } from "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol"; -import { EIP712 } from "@openzeppelin/contracts/utils/cryptography/EIP712.sol"; // todo can revoke role withdraw when the role is expired? // todo can grant role of an NFT already deposited? // Semi-fungible token (SFT) roles registry -contract SftRolesRegistry is IERCXXXX, ERC1155Holder, EIP712("SftRolesRegistry", "1") { - using BinaryTrees for BinaryTrees.Tree; +contract SftRolesRegistry is IERCXXXX, ERC1155Holder { + using BinaryTrees for BinaryTrees.Trees; using BinaryTrees for BinaryTrees.TreeNode; - // grantee => role => tokenAddress => tokenId => Tree - mapping(address => mapping(bytes32 => mapping(address => mapping(uint256 => BinaryTrees.Tree)))) public trees; + // binary tree for each role balance (grantee + role + tokenAddress + tokenId) + BinaryTrees.Trees internal trees; // nonce => RoleData - mapping(uint256 => RoleData) public roleAssignments; +// mapping(uint256 => RoleData) public roleAssignments; // grantor => tokenAddress => operator => isApproved mapping(address => mapping(address => mapping(address => bool))) public tokenApprovals; @@ -65,10 +64,11 @@ contract SftRolesRegistry is IERCXXXX, ERC1155Holder, EIP712("SftRolesRegistry", _roleAssignment.grantor ); - BinaryTrees.Tree storage tree = trees[_roleAssignment.grantee][_roleAssignment.role][_roleAssignment.tokenAddress][_roleAssignment.tokenId]; - BinaryTrees.TreeNode storage node = tree.nodes[_roleAssignment.nonce]; + bytes32 rootKey = _getRootKey(_roleAssignment.grantee, _roleAssignment.role, _roleAssignment.tokenAddress, _roleAssignment.tokenId); + BinaryTrees.TreeNode storage node = trees.nodes[_roleAssignment.nonce]; if (node.data.expirationDate == 0) { // expirationDate is only zero when the node does not exist + // transfer tokens to roles registry _transferFrom( _roleAssignment.grantor, address(this), @@ -95,11 +95,11 @@ contract SftRolesRegistry is IERCXXXX, ERC1155Holder, EIP712("SftRolesRegistry", } // remove node from tree - tree.remove(_roleAssignment.nonce); + trees.remove(rootKey, _roleAssignment.nonce); } - RoleData memory roleData = RoleData( + RoleData memory data = RoleData( hash, _roleAssignment.tokenAmount, _roleAssignment.expirationDate, @@ -107,7 +107,7 @@ contract SftRolesRegistry is IERCXXXX, ERC1155Holder, EIP712("SftRolesRegistry", _roleAssignment.data ); - tree.insert(_roleAssignment.nonce, roleData); + trees.insert(rootKey, _roleAssignment.nonce, data); emit RoleGranted( _roleAssignment.nonce, @@ -158,34 +158,6 @@ contract SftRolesRegistry is IERCXXXX, ERC1155Holder, EIP712("SftRolesRegistry", IERC1155(_tokenAddress).safeTransferFrom(_from, _to, _tokenId, _tokenAmount, ""); } - function _hashRoleData( - uint256 _nonce, - bytes32 _role, - address _tokenAddress, - uint256 _tokenId, -// uint256 _tokenAmount, - address _grantor//, -// address _grantee - ) internal view returns (bytes32) { - return _hashTypedDataV4( - keccak256( - abi.encode( - keccak256( -// "RoleAssignment(uint256 nonce,bytes32 role,address tokenAddress,uint256 tokenId,uint256 tokenAmount,address grantor,address grantee)" - "RoleAssignment(uint256 nonce,bytes32 role,address tokenAddress,uint256 tokenId)" - ), - _nonce, - _role, - _tokenAddress, - _tokenId, -// _tokenAmount, - _grantor//, -// _grantee - ) - ) - ); - } - function _findCaller(RevokeRoleData calldata _revokeRoleData) internal view returns (address) { if (_revokeRoleData.revoker == msg.sender || isRoleApprovedForAll(_revokeRoleData.tokenAddress, _revokeRoleData.revoker, msg.sender) @@ -217,8 +189,30 @@ contract SftRolesRegistry is IERCXXXX, ERC1155Holder, EIP712("SftRolesRegistry", return tokenApprovals[_grantor][_tokenAddress][_operator]; } + function roleData(uint256 _recordId) external view returns (RoleData memory data_) { + return trees.nodes[_recordId].data; + } + + function roleExpirationDate(uint256 _recordId) external view returns (uint64 expirationDate_) { + return trees.nodes[_recordId].data.expirationDate; + } + function supportsInterface(bytes4 interfaceId) public view virtual override(ERC1155Receiver, IERC165) returns (bool) { return interfaceId == type(IERCXXXX).interfaceId || interfaceId == type(IERC1155Receiver).interfaceId; } + /** Helper Functions **/ + + function _hashRoleData( + uint256 _nonce, bytes32 _role, address _tokenAddress, uint256 _tokenId, address _grantor + ) internal pure returns (bytes32) { + return keccak256(abi.encode(_nonce, _role, _tokenAddress, _tokenId, _grantor)); + } + + function _getRootKey( + address _grantee, bytes32 _role, address _tokenAddress, uint256 _tokenId + ) internal pure returns (bytes32 rootKey_) { + return keccak256(abi.encodePacked(_grantee, _role, _tokenAddress, _tokenId)); + } + } diff --git a/contracts/RolesRegistry/interfaces/IERCXXXX.sol b/contracts/RolesRegistry/interfaces/IERCXXXX.sol index 4dd5a34..664067b 100644 --- a/contracts/RolesRegistry/interfaces/IERCXXXX.sol +++ b/contracts/RolesRegistry/interfaces/IERCXXXX.sol @@ -114,12 +114,11 @@ interface IERCXXXX is IERC165 { /// @notice Returns the custom data of a role assignment. /// @param _recordId The identifier of the record. -// function roleData(uint256 _recordId) external view returns (RoleData memory data_); -// function roleData(address _tokenAddress, uint256 _recordId) external view returns (RoleData memory data_); + function roleData(uint256 _recordId) external view returns (RoleData memory data_); /// @notice Returns the expiration date of a role assignment. /// @param _recordId The identifier of the record. -// function roleExpirationDate(uint256 _recordId) external view returns (uint64 expirationDate_); + function roleExpirationDate(uint256 _recordId) external view returns (uint64 expirationDate_); /// @notice Checks if the grantor approved the operator for all NFTs. /// @param _tokenAddress The token address. diff --git a/contracts/RolesRegistry/libraries/BinaryTrees.sol b/contracts/RolesRegistry/libraries/BinaryTrees.sol index d1d8cb8..64e4b90 100644 --- a/contracts/RolesRegistry/libraries/BinaryTrees.sol +++ b/contracts/RolesRegistry/libraries/BinaryTrees.sol @@ -4,15 +4,16 @@ pragma solidity 0.8.9; import { IERCXXXX } from "../interfaces/IERCXXXX.sol"; - library BinaryTrees { uint256 public constant EMPTY = 0; bool public constant RED = true; bool public constant BLACK = false; - struct Tree { - uint256 root; + struct Trees { + // grantee => role => tokenAddress => tokenId => treeRoot + mapping(bytes32 => uint256) roots; + // nonce => TreeNode mapping (uint256 => TreeNode) nodes; } @@ -26,31 +27,31 @@ library BinaryTrees { // Insert ================================================================ - function insert(Tree storage _self, uint256 _nonce, IERCXXXX.RoleData memory _data) internal { + function insert(Trees storage _self, bytes32 _rootKey, uint256 _nonce, IERCXXXX.RoleData memory _data) internal { - if (_self.root == EMPTY) { + if (_self.roots[_rootKey] == EMPTY) { // if the tree is empty // insert node here as black - _self.root = _nonce; + _self.roots[_rootKey] = _nonce; _self.nodes[_nonce] = TreeNode(_data, 0, 0, 0, BLACK); } else { // if root exists // start searching for the right place to insert it - (TreeNode storage parent, uint256 parentNonce) = _insertHelper(_self, _self.root, _nonce, _data.expirationDate); + (TreeNode storage parent, uint256 parentNonce) = _insertHelper(_self, _self.roots[_rootKey], _nonce, _data.expirationDate); // insert new node as red _self.nodes[_nonce] = TreeNode(_data, parentNonce, 0, 0, RED); // check for violations (only if parent is red) if (parent.color == RED) { // if parent is red // fix violations - _fixViolations(_self, _self.nodes[_nonce], _nonce, parent, parentNonce); + _fixViolations(_self, _self.roots[_rootKey], _self.nodes[_nonce], _nonce, parent, parentNonce); } } } function _insertHelper( - Tree storage _self, uint256 _parentNonce, uint256 _nonce, uint64 _expirationDate + Trees storage _self, uint256 _parentNonce, uint256 _nonce, uint64 _expirationDate ) private returns (TreeNode storage parent_, uint256 parentNonce_) { TreeNode storage parentNode = _self.nodes[_parentNonce]; @@ -80,7 +81,8 @@ library BinaryTrees { } function _fixViolations( - Tree storage _self, + Trees storage _self, + uint256 _root, TreeNode storage _node, uint256 _nonce, TreeNode storage _parent, @@ -97,7 +99,7 @@ library BinaryTrees { uncle.color = BLACK; _parent.color = BLACK; // only recolor grandparent if it's not the root - if (_parent.parent != _self.root) { + if (_parent.parent != _root) { grandParent.color = RED; } return; @@ -140,7 +142,7 @@ library BinaryTrees { } - function _leftRotateAndUpdateRelativesColor(Tree storage _self, TreeNode storage _node, uint256 _nonce) private { + function _leftRotateAndUpdateRelativesColor(Trees storage _self, TreeNode storage _node, uint256 _nonce) private { // left rotate grandparent TreeNode storage currentParent = _self.nodes[_node.parent]; _leftRotation(_self, _self.nodes[currentParent.parent], currentParent.parent); @@ -160,7 +162,7 @@ library BinaryTrees { // Helpers =============================================================== function _findUncle( - Tree storage _self, uint256 _parentNonce, TreeNode storage _parent + Trees storage _self, uint256 _parentNonce, TreeNode storage _parent ) private view returns (TreeNode storage uncle_) { TreeNode storage grandParent = _self.nodes[_parent.parent]; if (_parentNonce == grandParent.right) { @@ -171,7 +173,7 @@ library BinaryTrees { } function _leftRotation( - Tree storage _self, TreeNode storage _node, uint256 _nonce + Trees storage _self, TreeNode storage _node, uint256 _nonce ) private returns (TreeNode storage node_) { TreeNode storage rightChild = _self.nodes[_node.right]; uint256 rightChildLeftChildNonce = rightChild.left; @@ -195,7 +197,7 @@ library BinaryTrees { } function _rightRotation( - Tree storage _self, TreeNode storage _node, uint256 _nonce + Trees storage _self, TreeNode storage _node, uint256 _nonce ) private returns (TreeNode storage node_) { TreeNode storage oldParent = _self.nodes[_node.parent]; TreeNode storage rightChildNode = _self.nodes[_node.right]; @@ -226,7 +228,7 @@ library BinaryTrees { // ======================================================================= // todo verify more edge cases - function remove(Tree storage _self, uint256 _nonce) internal { + function remove(Trees storage _self, bytes32 _rootKey, uint256 _nonce) internal { TreeNode storage _nodeToRemove = _self.nodes[_nonce]; // // modify removed nonce parent