Skip to content

Commit

Permalink
storing all TreeNodes in the same mapping
Browse files Browse the repository at this point in the history
  • Loading branch information
@ernanirst
ernanirst committed Nov 7, 2023
1 parent 0e95d4e commit a036899
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 58 deletions.
72 changes: 33 additions & 39 deletions contracts/RolesRegistry/SftRolesRegistry.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<RoleData>
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;
Expand Down Expand Up @@ -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),
Expand All @@ -95,19 +95,19 @@ 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,
_roleAssignment.revocable,
_roleAssignment.data
);

tree.insert(_roleAssignment.nonce, roleData);
trees.insert(rootKey, _roleAssignment.nonce, data);

emit RoleGranted(
_roleAssignment.nonce,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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));
}

}
5 changes: 2 additions & 3 deletions contracts/RolesRegistry/interfaces/IERCXXXX.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
34 changes: 18 additions & 16 deletions contracts/RolesRegistry/libraries/BinaryTrees.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand All @@ -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];

Expand Down Expand Up @@ -80,7 +81,8 @@ library BinaryTrees {
}

function _fixViolations(
Tree storage _self,
Trees storage _self,
uint256 _root,
TreeNode storage _node,
uint256 _nonce,
TreeNode storage _parent,
Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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) {
Expand All @@ -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;
Expand All @@ -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];
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit a036899

Please sign in to comment.