Add _setRoles and bytecode sizers

Vectorized · Jul 18, 2023 · 6dc1fde · 6dc1fde
1 parent 1f81570
commit 6dc1fde
Show file tree

Hide file tree

Showing 6 changed files with 94 additions and 33 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -657,49 +657,52 @@ MulticallableTest:testMulticallableRevertWithMessage(string) (runs: 256, μ: 126
 MulticallableTest:testMulticallableRevertWithNothing() (gas: 10155)
 MulticallableTest:testMulticallableWithNoData() (gas: 6322)
 MulticallableTest:test__codesize() (gas: 9792)
-OwnableRolesTest:testGrantAndRemoveRolesDirect(address,uint256,uint256) (runs: 256, μ: 37537, ~: 40526)
-OwnableRolesTest:testGrantAndRevokeOrRenounceRoles(address,bool,bool,bool,uint256,uint256) (runs: 256, μ: 19727, ~: 12176)
-OwnableRolesTest:testGrantRoles() (gas: 36111)
+OwnableRolesTest:testBytecodeSize() (gas: 300756)
+OwnableRolesTest:testGrantAndRemoveRolesDirect(address,uint256,uint256) (runs: 256, μ: 37592, ~: 40582)
+OwnableRolesTest:testGrantAndRevokeOrRenounceRoles(address,bool,bool,bool,uint256,uint256) (runs: 256, μ: 19411, ~: 12198)
+OwnableRolesTest:testGrantRoles() (gas: 36133)
 OwnableRolesTest:testHandoverOwnership() (gas: 32342)
-OwnableRolesTest:testHandoverOwnership(address) (runs: 256, μ: 32382, ~: 32368)
+OwnableRolesTest:testHandoverOwnership(address) (runs: 256, μ: 32400, ~: 32386)
 OwnableRolesTest:testHandoverOwnershipAfterExpiration() (gas: 36930)
-OwnableRolesTest:testHandoverOwnershipBeforeExpiration() (gas: 28702)
+OwnableRolesTest:testHandoverOwnershipBeforeExpiration() (gas: 28667)
 OwnableRolesTest:testHandoverOwnershipRevertsIfCompleteIsNotOwner() (gas: 35585)
 OwnableRolesTest:testHandoverOwnershipWithCancellation() (gas: 30715)
-OwnableRolesTest:testHasAllRoles(address,uint256,uint256,uint256,bool) (runs: 256, μ: 33542, ~: 33540)
-OwnableRolesTest:testHasAnyRole(address,uint256,uint256) (runs: 256, μ: 32380, ~: 33313)
+OwnableRolesTest:testHasAllRoles(address,uint256,uint256,uint256,bool) (runs: 256, μ: 33564, ~: 33562)
+OwnableRolesTest:testHasAnyRole(address,uint256,uint256) (runs: 256, μ: 32424, ~: 33357)
 OwnableRolesTest:testInitializeOwnerDirect() (gas: 16724)
 OwnableRolesTest:testOnlyOwnerModifier(address,bool) (runs: 256, μ: 17535, ~: 11530)
-OwnableRolesTest:testOnlyOwnerOrRolesModifier() (gas: 36549)
+OwnableRolesTest:testOnlyOwnerOrRolesModifier() (gas: 36593)
 OwnableRolesTest:testOnlyOwnerOrRolesModifier(address,bool,uint256,uint256) (runs: 256, μ: 54616, ~: 58506)
-OwnableRolesTest:testOnlyRolesModifier(address,uint256,uint256) (runs: 256, μ: 52611, ~: 57878)
-OwnableRolesTest:testOnlyRolesOrOwnerModifier(address,bool,uint256,uint256) (runs: 256, μ: 55255, ~: 58349)
-OwnableRolesTest:testOrdinalsFromRoles() (gas: 3610371)
-OwnableRolesTest:testOrdinalsFromRoles(uint256) (runs: 256, μ: 70624, ~: 35533)
-OwnableRolesTest:testOwnershipHandoverValidForDefaultValue() (gas: 5803)
+OwnableRolesTest:testOnlyRolesModifier(address,uint256,uint256) (runs: 256, μ: 53166, ~: 57923)
+OwnableRolesTest:testOnlyRolesOrOwnerModifier(address,bool,uint256,uint256) (runs: 256, μ: 55052, ~: 58393)
+OwnableRolesTest:testOrdinalsFromRoles() (gas: 3610415)
+OwnableRolesTest:testOrdinalsFromRoles(uint256) (runs: 256, μ: 72818, ~: 38085)
+OwnableRolesTest:testOwnershipHandoverValidForDefaultValue() (gas: 5825)
 OwnableRolesTest:testRenounceOwnership() (gas: 12899)
-OwnableRolesTest:testRolesFromOrdinals() (gas: 651237)
+OwnableRolesTest:testRolesFromOrdinals() (gas: 651259)
 OwnableRolesTest:testRolesFromOrdinals(uint8[]) (runs: 256, μ: 78088, ~: 74221)
-OwnableRolesTest:testSetOwnerDirect() (gas: 17820)
-OwnableRolesTest:testSetOwnerDirect(address) (runs: 256, μ: 17930, ~: 17968)
-OwnableRolesTest:testTransferOwnership() (gas: 19722)
-OwnableRolesTest:testTransferOwnership(address,bool,bool) (runs: 256, μ: 13915, ~: 13117)
-OwnableRolesTest:test__codesize() (gas: 21072)
-OwnableTest:testHandoverOwnership() (gas: 32055)
-OwnableTest:testHandoverOwnership(address) (runs: 256, μ: 32077, ~: 32064)
-OwnableTest:testHandoverOwnershipAfterExpiration() (gas: 36620)
-OwnableTest:testHandoverOwnershipBeforeExpiration() (gas: 28498)
-OwnableTest:testHandoverOwnershipRevertsIfCompleteIsNotOwner() (gas: 35364)
-OwnableTest:testHandoverOwnershipWithCancellation() (gas: 30300)
+OwnableRolesTest:testSetOwnerDirect() (gas: 17776)
+OwnableRolesTest:testSetOwnerDirect(address) (runs: 256, μ: 17844, ~: 17901)
+OwnableRolesTest:testSetRolesDirect(uint256) (runs: 256, μ: 99847, ~: 99819)
+OwnableRolesTest:testTransferOwnership() (gas: 19657)
+OwnableRolesTest:testTransferOwnership(address,bool,bool) (runs: 256, μ: 13975, ~: 13095)
+OwnableRolesTest:test__codesize() (gas: 23554)
+OwnableTest:testBytecodeSize() (gas: 214898)
+OwnableTest:testHandoverOwnership() (gas: 32072)
+OwnableTest:testHandoverOwnership(address) (runs: 256, μ: 32094, ~: 32081)
+OwnableTest:testHandoverOwnershipAfterExpiration() (gas: 36642)
+OwnableTest:testHandoverOwnershipBeforeExpiration() (gas: 28516)
+OwnableTest:testHandoverOwnershipRevertsIfCompleteIsNotOwner() (gas: 35341)
+OwnableTest:testHandoverOwnershipWithCancellation() (gas: 30318)
 OwnableTest:testInitializeOwnerDirect() (gas: 16599)
-OwnableTest:testOnlyOwnerModifier(address,bool) (runs: 256, μ: 17850, ~: 11408)
-OwnableTest:testOwnershipHandoverValidForDefaultValue() (gas: 5780)
-OwnableTest:testRenounceOwnership() (gas: 12799)
-OwnableTest:testSetOwnerDirect() (gas: 17664)
-OwnableTest:testSetOwnerDirect(address) (runs: 256, μ: 17812, ~: 17812)
-OwnableTest:testTransferOwnership() (gas: 19445)
-OwnableTest:testTransferOwnership(address,bool,bool) (runs: 256, μ: 13708, ~: 12851)
-OwnableTest:test__codesize() (gas: 10752)
+OwnableTest:testOnlyOwnerModifier(address,bool) (runs: 256, μ: 17348, ~: 11408)
+OwnableTest:testOwnershipHandoverValidForDefaultValue() (gas: 5737)
+OwnableTest:testRenounceOwnership() (gas: 12755)
+OwnableTest:testSetOwnerDirect() (gas: 17686)
+OwnableTest:testSetOwnerDirect(address) (runs: 256, μ: 17793, ~: 17812)
+OwnableTest:testTransferOwnership() (gas: 19467)
+OwnableTest:testTransferOwnership(address,bool,bool) (runs: 256, μ: 13685, ~: 12828)
+OwnableTest:test__codesize() (gas: 12122)
 RedBlackTreeLibTest:testRedBlackTreeBenchUint160() (gas: 3438446)
 RedBlackTreeLibTest:testRedBlackTreeBenchUint256() (gas: 5850739)
 RedBlackTreeLibTest:testRedBlackTreeClear() (gas: 57543)

diff --git a/src/auth/OwnableRoles.sol b/src/auth/OwnableRoles.sol
@@ -40,6 +40,19 @@ abstract contract OwnableRoles is Ownable {
     /*                     INTERNAL FUNCTIONS                     */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
+    /// @dev Overwrite the roles directly without authorization guard.
+    function _setRoles(address user, uint256 roles) internal virtual {
+        // Compute the role slot.
+        assembly {
+            mstore(0x0c, _ROLE_SLOT_SEED)
+            mstore(0x00, user)
+            // Store the new value.
+            sstore(keccak256(0x0c, 0x20), roles)
+            // Emit the {RolesUpdated} event.
+            log3(0, 0, _ROLES_UPDATED_EVENT_SIGNATURE, shr(96, mload(0x0c)), roles)
+        }
+    }
+
     /// @dev Updates the roles directly without authorization guard.
     /// If `on` is true, each set bit of `roles` will be turned on,
     /// otherwise, each set bit of `roles` will be turned off.

diff --git a/test/Ownable.t.sol b/test/Ownable.t.sol
@@ -17,6 +17,12 @@ contract OwnableTest is SoladyTest {
         mockOwnable = new MockOwnable();
     }
 
+    function testBytecodeSize() public {
+        MockOwnableBytecodeSizer mock = new MockOwnableBytecodeSizer();
+        assertTrue(address(mock).code.length > 0);
+        assertEq(mock.owner(), address(this));
+    }
+
     function testInitializeOwnerDirect() public {
         vm.expectEmit(true, true, true, true);
         emit OwnershipTransferred(address(0), address(1));

diff --git a/test/OwnableRoles.t.sol b/test/OwnableRoles.t.sol
@@ -19,6 +19,12 @@ contract OwnableRolesTest is SoladyTest {
         mockOwnableRoles = new MockOwnableRoles();
     }
 
+    function testBytecodeSize() public {
+        MockOwnableRolesBytecodeSizer mock = new MockOwnableRolesBytecodeSizer();
+        assertTrue(address(mock).code.length > 0);
+        assertEq(mock.owner(), address(this));
+    }
+
     function testInitializeOwnerDirect() public {
         vm.expectEmit(true, true, true, true);
         emit OwnershipTransferred(address(0), address(1));
@@ -45,6 +51,23 @@ contract OwnableRolesTest is SoladyTest {
         assertEq(mockOwnableRoles.rolesOf(user), rolesToGrant ^ (rolesToGrant & rolesToRemove));
     }
 
+    function testSetRolesDirect(uint256) public {
+        address userA = _randomNonZeroAddress();
+        address userB = _randomNonZeroAddress();
+        while (userA == userB) userA = _randomNonZeroAddress();
+        for (uint256 t; t != 2; ++t) {
+            uint256 rolesA = _random();
+            uint256 rolesB = _random();
+            vm.expectEmit(true, true, true, true);
+            emit RolesUpdated(userA, rolesA);
+            mockOwnableRoles.setRolesDirect(userA, rolesA);
+            emit RolesUpdated(userB, rolesB);
+            mockOwnableRoles.setRolesDirect(userB, rolesB);
+            assertEq(mockOwnableRoles.rolesOf(userA), rolesA);
+            assertEq(mockOwnableRoles.rolesOf(userB), rolesB);
+        }
+    }
+
     function testSetOwnerDirect() public {
         testSetOwnerDirect(address(1));
     }

diff --git a/test/utils/mocks/MockOwnable.sol b/test/utils/mocks/MockOwnable.sol
@@ -81,3 +81,9 @@ contract MockOwnable is Ownable {
         }
     }
 }
+
+contract MockOwnableBytecodeSizer is Ownable {
+    constructor() payable {
+        _initializeOwner(msg.sender);
+    }
+}
diff --git a/test/utils/mocks/MockOwnableRoles.sol b/test/utils/mocks/MockOwnableRoles.sol
@@ -43,6 +43,10 @@ contract MockOwnableRoles is OwnableRoles {
         _setOwner(_brutalizedAddress(newOwner));
     }
 
+    function setRolesDirect(address user, uint256 roles) public payable {
+        _setRoles(_brutalizedAddress(user), roles);
+    }
+
     function grantRolesDirect(address user, uint256 roles) public payable {
         _grantRoles(_brutalizedAddress(user), roles);
     }
@@ -153,3 +157,9 @@ contract MockOwnableRoles is OwnableRoles {
         if (!resultIsOneOrZero) result = !result;
     }
 }
+
+contract MockOwnableRolesBytecodeSizer is OwnableRoles {
+    constructor() payable {
+        _initializeOwner(msg.sender);
+    }
+}