From 901123e7a7cbc1d26804a488d1c5905f9dbd7a01 Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:07:44 +0530 Subject: [PATCH 01/13] test(diamond): define BTT behavior tree --- test/trees/Diamond.tree | 147 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 test/trees/Diamond.tree diff --git a/test/trees/Diamond.tree b/test/trees/Diamond.tree new file mode 100644 index 00000000..c426edad --- /dev/null +++ b/test/trees/Diamond.tree @@ -0,0 +1,147 @@ +SelectorImport +├── when the facet address has no bytecode +│ └── it should revert with {NoBytecodeAtAddress} +├── when exportSelectors cannot be called successfully +│ └── it should revert with the surface-specific {FunctionSelectorsCallFailed} or {ExportSelectorsCallFailed} error +├── when exportSelectors return data is shorter than the ABI offset and length words +│ └── it should revert with {IncorrectSelectorsEncoding} +├── when the ABI offset is not 32 +│ └── it should revert with {IncorrectSelectorsEncoding} +├── when the encoded selector length exceeds the returned payload +│ └── it should revert with {IncorrectSelectorsEncoding} +├── when the packed selector bytes are empty +│ └── it should revert with {NoSelectorsForFacet} +├── when the packed selector length is not a multiple of four +│ └── it should revert with {IncorrectSelectorsEncoding} +└── when the facet returns valid packed selectors + └── it should import every selector in export order + +AddFacets +├── when the facet array is empty +│ └── it should leave the facet list and counts unchanged +├── when adding the first facet +│ ├── it should map every exported selector to the facet +│ ├── it should set head and tail to the first exported selector +│ ├── it should set the first node links to zero +│ ├── it should increment facetCount and selectorCount exactly +│ └── it should emit {FacetAdded} exactly once +├── when adding multiple facets +│ ├── it should preserve input order in the linked list +│ ├── it should link head, middle, and tail nodes in both directions +│ ├── it should keep non-head selector links zero +│ ├── it should increment facetCount and selectorCount exactly +│ └── it should emit one {FacetAdded} event per facet in input order +├── when appending facets to an existing list +│ ├── it should connect the previous tail to the appended head +│ └── it should preserve the existing head and set the new tail +├── when a selector already belongs to the diamond +│ └── it should revert with {CannotAddFunctionToDiamondThatAlreadyExists} +└── when a facet exports a selector more than once + └── it should revert with {CannotAddFunctionToDiamondThatAlreadyExists} + +ReplaceFacets +├── when the replacement array is empty +│ └── it should leave the facet list and counts unchanged +├── when old and new facet addresses are the same +│ └── it should revert with {CannotReplaceFacetWithSameFacet} +├── when the old facet is not registered +│ └── it should revert with {FacetToReplaceDoesNotExist} +├── when a new selector belongs to a different registered facet +│ └── it should revert with {CannotReplaceFunctionFromNonReplacementFacet} +├── when old and new facets export identical selectors +│ ├── it should route every selector to the new facet +│ ├── it should preserve head, tail, links, and counts +│ └── it should emit {FacetReplaced} exactly once +├── when the selector sets overlap but differ +│ ├── it should route shared selectors to the new facet +│ ├── it should add selectors unique to the new facet +│ ├── it should remove selectors unique to the old facet +│ └── it should update selectorCount to the exact new total +├── when replacing the first facet with a different node selector +│ └── it should update head and the next node's previous link +├── when replacing a middle facet with a different node selector +│ └── it should relink both neighboring nodes through the new node +└── when replacing the last facet with a different node selector + └── it should update tail and the previous node's next link + +RemoveFacets +├── when the removal array is empty +│ └── it should leave the facet list and counts unchanged +├── when the facet is not registered +│ └── it should revert with {CannotRemoveFacetThatDoesNotExist} +├── when removing the only facet +│ ├── it should clear every selector mapping +│ ├── it should clear head and tail +│ └── it should set facetCount and selectorCount to zero +├── when removing the first facet +│ └── it should move head and clear the new head's previous link +├── when removing a middle facet +│ └── it should link the previous and next nodes to each other +├── when removing the last facet +│ └── it should move tail and clear the new tail's next link +├── when removing multiple facets +│ ├── it should remove every exported selector +│ ├── it should decrement both counts exactly +│ └── it should emit one {FacetRemoved} event per facet in input order +└── when a removal fails + └── it should roll back all earlier removals in the transaction + +UpgradeDiamond +├── when the caller is not the owner on the upgrade facet +│ └── it should revert with {OwnerUnauthorizedAccount} +├── when additions replacements and removals are supplied together +│ ├── it should execute add then replace then remove +│ └── it should leave routing list links and counts consistent +├── when delegate is the zero address +│ └── it should skip delegatecall and omit {DiamondDelegateCall} +├── when delegate has no bytecode +│ └── it should revert with {NoBytecodeAtAddress} +├── when delegatecall succeeds +│ ├── it should mutate the caller diamond's storage +│ └── it should emit {DiamondDelegateCall} with exact calldata +├── when delegatecall reverts with data +│ └── it should bubble the exact revert data +├── when delegatecall reverts without data +│ └── it should revert with {DelegateCallReverted} +├── when tag is zero and metadata is empty +│ └── it should omit {DiamondMetadata} +├── when tag is nonzero or metadata is nonempty +│ └── it should emit {DiamondMetadata} with exact values +└── ExportSelectors + └── it should return only the upgradeDiamond selector + +Fallback +├── when no facet owns msg.sig +│ └── it should revert with {FunctionNotFound} and msg.sig +├── when a facet owns msg.sig +│ ├── it should forward exact calldata +│ ├── it should preserve msg.sender and msg.value +│ ├── it should write state in diamond storage not facet storage +│ └── it should return exact returndata +└── when the facet call reverts + └── it should bubble exact revert data + +Inspect +├── when no facets are registered +│ ├── facetAddress should return address zero +│ ├── facetAddresses should return an empty array +│ └── facets should return an empty array +├── when facets are registered +│ ├── facetAddress should return the owner of each selector +│ ├── facetFunctionSelectors should unpack selectors in export order +│ ├── facetAddresses should return linked-list order +│ └── facets should pair every facet with its exported selectors +├── when a queried valid facet is not registered +│ └── facetFunctionSelectors should return an empty array +├── when an old facet has been replaced and its first selector is now owned by the replacement +│ └── facetFunctionSelectors for the old facet should return an empty array +├── when a facet has been removed +│ └── every inspection view should exclude it and its selectors +├── ExportSelectors +│ └── it should return the four inspection selectors in declaration order +└── when add replace remove sequences are fuzzed + ├── facetCount should equal the number of traversable facet nodes + ├── selectorCount should equal the number of routed unique selectors + ├── head and tail should match the first and last traversable nodes + ├── every next link should have a matching previous link + └── inspection output should match the model after every step From f1f4166686bd3519c2044a263dc12c2cfca73d4b Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:14:17 +0530 Subject: [PATCH 02/13] test(diamond): add core fixtures and storage readers --- test/unit/diamond/DiamondModBase.t.sol | 73 +++++ test/unit/diamond/DiamondUpgradeBase.t.sol | 150 ++++++++++ .../harnesses/diamond/DiamondModHarness.sol | 24 ++ .../harnesses/diamond/DiamondProxyHarness.sol | 22 ++ .../diamond/DiamondUpgradeModHarness.sol | 24 ++ .../utils/mocks/diamond/DiamondFacetMocks.sol | 282 ++++++++++++++++++ test/utils/storage/DiamondStorageUtils.sol | 38 +++ 7 files changed, 613 insertions(+) create mode 100644 test/unit/diamond/DiamondModBase.t.sol create mode 100644 test/unit/diamond/DiamondUpgradeBase.t.sol create mode 100644 test/utils/harnesses/diamond/DiamondModHarness.sol create mode 100644 test/utils/harnesses/diamond/DiamondProxyHarness.sol create mode 100644 test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol create mode 100644 test/utils/mocks/diamond/DiamondFacetMocks.sol create mode 100644 test/utils/storage/DiamondStorageUtils.sol diff --git a/test/unit/diamond/DiamondModBase.t.sol b/test/unit/diamond/DiamondModBase.t.sol new file mode 100644 index 00000000..03c64ea8 --- /dev/null +++ b/test/unit/diamond/DiamondModBase.t.sol @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Base_Test} from "test/Base.t.sol"; +import {DiamondModHarness} from "test/utils/harnesses/diamond/DiamondModHarness.sol"; +import {DispatchFacet, FacetA, FacetB, FacetC} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +import {DiamondStorageUtils} from "test/utils/storage/DiamondStorageUtils.sol"; + +abstract contract DiamondMod_Base_Test is Base_Test { + event FacetAdded(address indexed _facet); + + DiamondModHarness internal harness; + FacetA internal facetA; + FacetB internal facetB; + FacetC internal facetC; + DispatchFacet internal dispatchFacet; + + function setUp() public virtual override { + Base_Test.setUp(); + harness = new DiamondModHarness(_emptyAddresses()); + facetA = new FacetA(); + facetB = new FacetB(); + facetC = new FacetC(); + dispatchFacet = new DispatchFacet(); + } + + function _emptyAddresses() internal pure returns (address[] memory values) { + values = new address[](0); + } + + function _singleAddress(address _value) internal pure returns (address[] memory values) { + values = new address[](1); + values[0] = _value; + } + + function _threeFacets(address _a, address _b, address _c) internal pure returns (address[] memory values) { + values = new address[](3); + values[0] = _a; + values[1] = _b; + values[2] = _c; + } + + function _assertFacetList( + address _target, + bytes4 _expectedHead, + bytes4 _expectedTail, + uint32 _expectedFacetCount, + uint32 _expectedSelectorCount + ) internal view { + (bytes4 head, bytes4 tail, uint32 facetCount, uint32 selectorCount) = DiamondStorageUtils.facetList(_target); + assertEq(head, _expectedHead, "head"); + assertEq(tail, _expectedTail, "tail"); + assertEq(facetCount, _expectedFacetCount, "facetCount"); + assertEq(selectorCount, _expectedSelectorCount, "selectorCount"); + } + + function _assertNode( + address _target, + bytes4 _selector, + address _expectedFacet, + bytes4 _expectedPrev, + bytes4 _expectedNext + ) internal view { + (address facet, bytes4 prev, bytes4 next) = DiamondStorageUtils.facetNode(_target, _selector); + assertEq(facet, _expectedFacet, "node facet"); + assertEq(prev, _expectedPrev, "node prev"); + assertEq(next, _expectedNext, "node next"); + } +} diff --git a/test/unit/diamond/DiamondUpgradeBase.t.sol b/test/unit/diamond/DiamondUpgradeBase.t.sol new file mode 100644 index 00000000..bf7194a1 --- /dev/null +++ b/test/unit/diamond/DiamondUpgradeBase.t.sol @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Base_Test} from "test/Base.t.sol"; +import {DiamondModHarness} from "test/utils/harnesses/diamond/DiamondModHarness.sol"; +import { + DelegateTarget, + FacetA, + FacetAChanged, + FacetAReplacement, + FacetB, + FacetBChanged, + FacetBReplacement, + FacetC, + FacetCChanged, + FacetCReplacement +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +import {DiamondStorageUtils} from "test/utils/storage/DiamondStorageUtils.sol"; + +struct Replacement { + address oldFacet; + address newFacet; +} + +abstract contract DiamondUpgrade_Base_Test is Base_Test { + using DiamondStorageUtils for address; + + event FacetAdded(address indexed _facet); + event FacetReplaced(address indexed _oldFacet, address indexed _newFacet); + event FacetRemoved(address indexed _facet); + event DiamondDelegateCall(address indexed _delegate, bytes _delegateCalldata); + event DiamondMetadata(bytes32 indexed _tag, bytes _data); + + address internal target; + FacetA internal facetA; + FacetAReplacement internal facetAReplacement; + FacetAChanged internal facetAChanged; + FacetB internal facetB; + FacetBReplacement internal facetBReplacement; + FacetBChanged internal facetBChanged; + FacetC internal facetC; + FacetCReplacement internal facetCReplacement; + FacetCChanged internal facetCChanged; + DelegateTarget internal delegateTarget; + + function setUp() public virtual override { + Base_Test.setUp(); + facetA = new FacetA(); + facetAReplacement = new FacetAReplacement(); + facetAChanged = new FacetAChanged(); + facetB = new FacetB(); + facetBReplacement = new FacetBReplacement(); + facetBChanged = new FacetBChanged(); + facetC = new FacetC(); + facetCReplacement = new FacetCReplacement(); + facetCChanged = new FacetCChanged(); + delegateTarget = new DelegateTarget(); + } + + function _upgrade( + address[] memory _adds, + Replacement[] memory _replacements, + address[] memory _removes, + address _delegate, + bytes memory _delegateCalldata, + bytes32 _tag, + bytes memory _metadata + ) internal virtual; + + function _assertCounts(uint32 _facets, uint32 _selectors) internal view { + (,, uint32 facetCount, uint32 selectorCount) = target.facetList(); + assertEq(facetCount, _facets, "facetCount"); + assertEq(selectorCount, _selectors, "selectorCount"); + } + + function _emptyAddresses() internal pure returns (address[] memory values) { + values = new address[](0); + } + + function _emptyReplacements() internal pure returns (Replacement[] memory values) { + values = new Replacement[](0); + } + + function _singleAddress(address _value) internal pure returns (address[] memory values) { + values = new address[](1); + values[0] = _value; + } + + function _singleReplacement(address _oldFacet, address _newFacet) + internal + pure + returns (Replacement[] memory values) + { + values = new Replacement[](1); + values[0] = Replacement(_oldFacet, _newFacet); + } + + function _threeFacets(address _a, address _b, address _c) internal pure returns (address[] memory values) { + values = new address[](3); + values[0] = _a; + values[1] = _b; + values[2] = _c; + } + + function _assertSelectorOwner(bytes4 _selector, address _expectedFacet) internal view { + (address actualFacet,,) = DiamondStorageUtils.facetNode(target, _selector); + assertEq(actualFacet, _expectedFacet, "selector owner"); + } + + function _assertNode(bytes4 _selector, address _facet, bytes4 _prev, bytes4 _next) internal view { + (address actualFacet, bytes4 actualPrev, bytes4 actualNext) = DiamondStorageUtils.facetNode(target, _selector); + assertEq(actualFacet, _facet, "node facet"); + assertEq(actualPrev, _prev, "node prev"); + assertEq(actualNext, _next, "node next"); + } + + function _assertThreeNodeList(address _a, address _b, address _c) internal view { + _assertNode(FacetA.a1.selector, _a, bytes4(0), FacetB.b1.selector); + _assertNode(FacetB.b1.selector, _b, FacetA.a1.selector, FacetC.c1.selector); + _assertNode(FacetC.c1.selector, _c, FacetB.b1.selector, bytes4(0)); + } +} + +contract DiamondStorageLayout_Unit_Test is Base_Test { + function test_ShouldProveStorageBitOffsets() external { + FacetA facetA = new FacetA(); + address[] memory facets = new address[](1); + facets[0] = address(facetA); + DiamondModHarness harness = new DiamondModHarness(facets); + + (bytes4 head, bytes4 tail, uint32 facetCount, uint32 selectorCount) = + DiamondStorageUtils.facetList(address(harness)); + assertEq(head, FacetA.a1.selector, "head"); + assertEq(tail, FacetA.a1.selector, "tail"); + assertEq(facetCount, 1, "facetCount"); + assertEq(selectorCount, 3, "selectorCount"); + + (address facet, bytes4 prev, bytes4 next) = DiamondStorageUtils.facetNode(address(harness), FacetA.a1.selector); + assertEq(facet, address(facetA), "node facet"); + assertEq(prev, bytes4(0), "node prev"); + assertEq(next, bytes4(0), "node next"); + + (facet,,) = DiamondStorageUtils.facetNode(address(harness), FacetA.a2.selector); + assertEq(facet, address(facetA), "non-head owner"); + } +} diff --git a/test/utils/harnesses/diamond/DiamondModHarness.sol b/test/utils/harnesses/diamond/DiamondModHarness.sol new file mode 100644 index 00000000..916db506 --- /dev/null +++ b/test/utils/harnesses/diamond/DiamondModHarness.sol @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import "src/diamond/DiamondMod.sol" as DiamondMod; + +contract DiamondModHarness { + constructor(address[] memory _facets) { + DiamondMod.addFacets(_facets); + } + + function addFacets(address[] memory _facets) external { + DiamondMod.addFacets(_facets); + } + + fallback() external payable { + DiamondMod.diamondFallback(); + } + + receive() external payable {} +} diff --git a/test/utils/harnesses/diamond/DiamondProxyHarness.sol b/test/utils/harnesses/diamond/DiamondProxyHarness.sol new file mode 100644 index 00000000..0effff46 --- /dev/null +++ b/test/utils/harnesses/diamond/DiamondProxyHarness.sol @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import "src/diamond/DiamondMod.sol" as DiamondMod; +import "src/access/Owner/Data/OwnerDataMod.sol" as OwnerDataMod; + +contract DiamondProxyHarness { + constructor(address[] memory _facets, address _diamondOwner) { + DiamondMod.addFacets(_facets); + OwnerDataMod.setContractOwner(_diamondOwner); + } + + fallback() external payable { + DiamondMod.diamondFallback(); + } + + receive() external payable {} +} diff --git a/test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol b/test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol new file mode 100644 index 00000000..321da03a --- /dev/null +++ b/test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import "src/diamond/DiamondUpgradeMod.sol" as DiamondUpgradeMod; + +contract DiamondUpgradeModHarness { + function upgradeDiamond( + address[] calldata _addFacets, + DiamondUpgradeMod.FacetReplacement[] calldata _replaceFacets, + address[] calldata _removeFacets, + address _delegate, + bytes calldata _delegateCalldata, + bytes32 _tag, + bytes calldata _metadata + ) external { + DiamondUpgradeMod.upgradeDiamond( + _addFacets, _replaceFacets, _removeFacets, _delegate, _delegateCalldata, _tag, _metadata + ); + } +} diff --git a/test/utils/mocks/diamond/DiamondFacetMocks.sol b/test/utils/mocks/diamond/DiamondFacetMocks.sol new file mode 100644 index 00000000..e2501d1a --- /dev/null +++ b/test/utils/mocks/diamond/DiamondFacetMocks.sol @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +bytes32 constant DIAMOND_TEST_STORAGE_POSITION = keccak256("compose.test.diamond"); + +contract FacetA { + function a1() external pure returns (uint256) { + return 1; + } + + function a2(uint256 _value) external pure returns (uint256) { + return _value; + } + + function a3() external pure returns (bytes32) { + return keccak256("a3"); + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.a1.selector, this.a2.selector, this.a3.selector); + } +} + +contract FacetAReplacement { + function a1() external pure returns (uint256) { + return 2; + } + + function a2(uint256 _value) external pure returns (uint256) { + return _value + 1; + } + + function a3() external pure returns (bytes32) { + return keccak256("a3-v2"); + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.a1.selector, this.a2.selector, this.a3.selector); + } +} + +contract FacetAChanged { + function a4() external pure returns (uint256) { + return 4; + } + + function a2(uint256 _value) external pure returns (uint256) { + return _value + 2; + } + + function a5() external pure returns (uint256) { + return 5; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.a4.selector, this.a2.selector, this.a5.selector); + } +} + +contract FacetB { + function b1() external pure returns (uint256) { + return 11; + } + + function b2(bytes32 _value) external pure returns (bytes32) { + return _value; + } + + function b3(address _value) external pure returns (address) { + return _value; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.b1.selector, this.b2.selector, this.b3.selector); + } +} + +contract FacetBReplacement { + function b1() external pure returns (uint256) { + return 12; + } + + function b2(bytes32 _value) external pure returns (bytes32) { + return keccak256(abi.encode(_value)); + } + + function b3(address _value) external pure returns (address) { + return address(uint160(_value) ^ 1); + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.b1.selector, this.b2.selector, this.b3.selector); + } +} + +contract FacetBChanged { + function b4() external pure returns (uint256) { + return 14; + } + + function b2(bytes32 _value) external pure returns (bytes32) { + return _value; + } + + function b5() external pure returns (uint256) { + return 15; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.b4.selector, this.b2.selector, this.b5.selector); + } +} + +contract FacetC { + function c1() external pure returns (uint256) { + return 21; + } + + function c2(bool _value) external pure returns (bool) { + return _value; + } + + function c3(bytes calldata _value) external pure returns (bytes32) { + return keccak256(_value); + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.c1.selector, this.c2.selector, this.c3.selector); + } +} + +contract FacetCReplacement { + function c1() external pure returns (uint256) { + return 22; + } + + function c2(bool _value) external pure returns (bool) { + return !_value; + } + + function c3(bytes calldata _value) external pure returns (bytes32) { + return keccak256(abi.encode(_value)); + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.c1.selector, this.c2.selector, this.c3.selector); + } +} + +contract FacetCChanged { + function c4() external pure returns (uint256) { + return 24; + } + + function c2(bool _value) external pure returns (bool) { + return _value; + } + + function c5() external pure returns (uint256) { + return 25; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.c4.selector, this.c2.selector, this.c5.selector); + } +} + +contract EmptySelectorsFacet { + function exportSelectors() external pure returns (bytes memory) { + return bytes(""); + } +} + +contract MisalignedSelectorsFacet { + function exportSelectors() external pure returns (bytes memory) { + return hex"0102030405"; + } +} + +contract RevertingSelectorsFacet { + error SelectorExportReverted(); + + function exportSelectors() external pure returns (bytes memory) { + revert SelectorExportReverted(); + } +} + +contract MissingSelectorsFacet {} + +contract ShortReturnFacet { + fallback() external { + assembly ("memory-safe") { + mstore(0, 0x20) + return(0, 0x20) + } + } +} + +contract BadOffsetFacet { + fallback() external { + assembly ("memory-safe") { + mstore(0, 0) + mstore(0x20, 4) + return(0, 0x40) + } + } +} + +contract OversizedLengthFacet { + fallback() external { + assembly ("memory-safe") { + mstore(0, 0x20) + mstore(0x20, 0x20) + return(0, 0x40) + } + } +} + +contract SelectorConflictFacet { + function conflictHead() external pure returns (uint256) { + return 1; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.conflictHead.selector, FacetB.b2.selector); + } +} + +contract DuplicateSelectorFacet { + function duplicate() external pure returns (uint256) { + return 1; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.duplicate.selector, this.duplicate.selector); + } +} + +contract DispatchFacet { + error DispatchFailure(uint256 _value); + + function context(uint256 _value) external payable returns (address sender, uint256 value, uint256 argument) { + bytes32 position = DIAMOND_TEST_STORAGE_POSITION; + assembly ("memory-safe") { + sstore(position, _value) + } + return (msg.sender, msg.value, _value); + } + + function fail(uint256 _value) external pure { + revert DispatchFailure(_value); + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.context.selector, this.fail.selector); + } +} + +contract DelegateTarget { + error DelegateFailure(uint256 _value); + + function initialize(uint256 _value) external { + bytes32 position = DIAMOND_TEST_STORAGE_POSITION; + assembly ("memory-safe") { + sstore(position, _value) + } + } + + function failWithData(uint256 _value) external pure { + revert DelegateFailure(_value); + } + + function failWithoutData() external pure { + assembly ("memory-safe") { + revert(0, 0) + } + } +} diff --git a/test/utils/storage/DiamondStorageUtils.sol b/test/utils/storage/DiamondStorageUtils.sol new file mode 100644 index 00000000..26572d4f --- /dev/null +++ b/test/utils/storage/DiamondStorageUtils.sol @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Vm} from "forge-std/Vm.sol"; + +library DiamondStorageUtils { + Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code"))))); + bytes32 internal constant DIAMOND_STORAGE_POSITION = keccak256("erc8153.diamond"); + + function facetNode(address _target, bytes4 _selector) + internal + view + returns (address facet, bytes4 prevFacetNodeId, bytes4 nextFacetNodeId) + { + bytes32 slot = keccak256(abi.encode(_selector, DIAMOND_STORAGE_POSITION)); + uint256 word = uint256(vm.load(_target, slot)); + facet = address(uint160(word)); + prevFacetNodeId = bytes4(uint32(word >> 160)); + nextFacetNodeId = bytes4(uint32(word >> 192)); + } + + function facetList(address _target) + internal + view + returns (bytes4 headFacetNodeId, bytes4 tailFacetNodeId, uint32 facetCount, uint32 selectorCount) + { + bytes32 slot = bytes32(uint256(DIAMOND_STORAGE_POSITION) + 1); + uint256 word = uint256(vm.load(_target, slot)); + headFacetNodeId = bytes4(uint32(word)); + tailFacetNodeId = bytes4(uint32(word >> 32)); + facetCount = uint32(word >> 64); + selectorCount = uint32(word >> 96); + } +} From ebf7cc2c868a70ecf41df441c50244112be08ecc Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:23:05 +0530 Subject: [PATCH 03/13] test(diamond): cover selector import add and fallback --- test/unit/diamond/mod/fuzz/addFacets.t.sol | 127 ++++++++++++++++++ test/unit/diamond/mod/fuzz/fallback.t.sol | 54 ++++++++ .../diamond/mod/fuzz/importSelectors.t.sol | 81 +++++++++++ 3 files changed, 262 insertions(+) create mode 100644 test/unit/diamond/mod/fuzz/addFacets.t.sol create mode 100644 test/unit/diamond/mod/fuzz/fallback.t.sol create mode 100644 test/unit/diamond/mod/fuzz/importSelectors.t.sol diff --git a/test/unit/diamond/mod/fuzz/addFacets.t.sol b/test/unit/diamond/mod/fuzz/addFacets.t.sol new file mode 100644 index 00000000..3ccf744a --- /dev/null +++ b/test/unit/diamond/mod/fuzz/addFacets.t.sol @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Vm} from "forge-std/Vm.sol"; + +import "src/diamond/DiamondMod.sol" as DiamondMod; +import {DiamondMod_Base_Test} from "test/unit/diamond/DiamondModBase.t.sol"; +import { + DuplicateSelectorFacet, + FacetA, + FacetB, + FacetC, + SelectorConflictFacet +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract AddFacets_DiamondMod_Fuzz_Unit_Test is DiamondMod_Base_Test { + bytes32 private constant FACET_ADDED_TOPIC = keccak256("FacetAdded(address)"); + + function test_ShouldLeaveFacetListUnchanged_WhenFacetArrayIsEmpty() external { + harness.addFacets(_singleAddress(address(facetA))); + vm.recordLogs(); + + harness.addFacets(_emptyAddresses()); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + _assertFacetList(address(harness), FacetA.a1.selector, FacetA.a1.selector, 1, 3); + _assertNode(address(harness), FacetA.a1.selector, address(facetA), bytes4(0), bytes4(0)); + assertEq(logs.length, 0, "unexpected logs"); + } + + function test_ShouldRegisterFirstFacetSelectorsAndList() external { + harness.addFacets(_singleAddress(address(facetA))); + + _assertFacetList(address(harness), FacetA.a1.selector, FacetA.a1.selector, 1, 3); + _assertNode(address(harness), FacetA.a1.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(address(harness), FacetA.a2.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(address(harness), FacetA.a3.selector, address(facetA), bytes4(0), bytes4(0)); + } + + function test_ShouldEmitFacetAddedExactlyOnce_WhenAddingFirstFacet() external { + vm.recordLogs(); + + harness.addFacets(_singleAddress(address(facetA))); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + assertEq(logs.length, 1, "log count"); + _assertFacetAddedLog(logs[0], address(facetA)); + } + + function test_ShouldLinkMultipleFacetsInInputOrder() external { + address[] memory facets = new address[](3); + facets[0] = address(facetA); + facets[1] = address(facetB); + facets[2] = address(facetC); + + harness.addFacets(facets); + + _assertFacetList(address(harness), FacetA.a1.selector, FacetC.c1.selector, 3, 9); + _assertNode(address(harness), FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); + _assertNode(address(harness), FacetB.b1.selector, address(facetB), FacetA.a1.selector, FacetC.c1.selector); + _assertNode(address(harness), FacetC.c1.selector, address(facetC), FacetB.b1.selector, bytes4(0)); + _assertNode(address(harness), FacetB.b2.selector, address(facetB), bytes4(0), bytes4(0)); + } + + function test_ShouldEmitOneFacetAddedPerFacetInInputOrder() external { + vm.recordLogs(); + + harness.addFacets(_threeFacets(address(facetA), address(facetB), address(facetC))); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + assertEq(logs.length, 3, "log count"); + _assertFacetAddedLog(logs[0], address(facetA)); + _assertFacetAddedLog(logs[1], address(facetB)); + _assertFacetAddedLog(logs[2], address(facetC)); + } + + function test_ShouldConnectPreviousTailToAppendedFacets() external { + harness.addFacets(_singleAddress(address(facetA))); + + address[] memory appendedFacets = new address[](2); + appendedFacets[0] = address(facetB); + appendedFacets[1] = address(facetC); + harness.addFacets(appendedFacets); + + _assertFacetList(address(harness), FacetA.a1.selector, FacetC.c1.selector, 3, 9); + _assertNode(address(harness), FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); + _assertNode(address(harness), FacetB.b1.selector, address(facetB), FacetA.a1.selector, FacetC.c1.selector); + _assertNode(address(harness), FacetC.c1.selector, address(facetC), FacetB.b1.selector, bytes4(0)); + } + + function test_RevertWhen_SelectorBelongsToAnotherFacet() external { + SelectorConflictFacet conflictFacet = new SelectorConflictFacet(); + harness.addFacets(_singleAddress(address(facetB))); + + vm.expectRevert( + abi.encodeWithSelector(DiamondMod.CannotAddFunctionToDiamondThatAlreadyExists.selector, FacetB.b2.selector) + ); + harness.addFacets(_singleAddress(address(conflictFacet))); + } + + function test_RevertWhen_FacetExportsDuplicateSelector() external { + DuplicateSelectorFacet duplicateFacet = new DuplicateSelectorFacet(); + + vm.expectRevert( + abi.encodeWithSelector( + DiamondMod.CannotAddFunctionToDiamondThatAlreadyExists.selector, + DuplicateSelectorFacet.duplicate.selector + ) + ); + harness.addFacets(_singleAddress(address(duplicateFacet))); + } + + function _assertFacetAddedLog(Vm.Log memory _log, address _facet) private view { + assertEq(_log.emitter, address(harness), "log emitter"); + assertEq(_log.topics.length, 2, "topic count"); + assertEq(_log.topics[0], FACET_ADDED_TOPIC, "event signature"); + assertEq(_log.topics[1], bytes32(uint256(uint160(_facet))), "facet topic"); + assertEq(_log.data.length, 0, "event data"); + } +} diff --git a/test/unit/diamond/mod/fuzz/fallback.t.sol b/test/unit/diamond/mod/fuzz/fallback.t.sol new file mode 100644 index 00000000..1bc9a7f8 --- /dev/null +++ b/test/unit/diamond/mod/fuzz/fallback.t.sol @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import "src/diamond/DiamondMod.sol" as DiamondMod; +import {DiamondMod_Base_Test} from "test/unit/diamond/DiamondModBase.t.sol"; +import {DIAMOND_TEST_STORAGE_POSITION, DispatchFacet} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract Fallback_DiamondMod_Fuzz_Unit_Test is DiamondMod_Base_Test { + function test_ShouldRevertWithFunctionNotFoundForUnknownSelector() external { + bytes4 unknownSelector = bytes4(keccak256("unknownFunction(uint256)")); + bytes memory callData = abi.encodeWithSelector(unknownSelector, uint256(123)); + + (bool success, bytes memory data) = address(harness).call(callData); + + assertEq(success, false, "fallback call"); + assertEq(data, abi.encodeWithSelector(DiamondMod.FunctionNotFound.selector, unknownSelector), "revert data"); + } + + function testFuzz_ShouldPreserveContextAndReturnData(uint96 _value, uint256 _argument, address _caller) external { + vm.assume(_caller != address(0)); + harness.addFacets(_singleAddress(address(dispatchFacet))); + vm.deal(_caller, _value); + + vm.stopPrank(); + vm.prank(_caller); + (bool success, bytes memory data) = + address(harness).call{value: _value}(abi.encodeCall(DispatchFacet.context, (_argument))); + + assertEq(success, true, "fallback call"); + assertEq(data, abi.encode(_caller, uint256(_value), _argument), "returndata"); + (address sender, uint256 value, uint256 argument) = abi.decode(data, (address, uint256, uint256)); + assertEq(sender, _caller, "msg.sender"); + assertEq(value, _value, "msg.value"); + assertEq(argument, _argument, "argument"); + assertEq(uint256(vm.load(address(harness), DIAMOND_TEST_STORAGE_POSITION)), _argument, "diamond storage"); + assertEq(uint256(vm.load(address(dispatchFacet), DIAMOND_TEST_STORAGE_POSITION)), 0, "facet storage"); + } + + function testFuzz_ShouldBubbleExactRevertData(uint256 _value) external { + harness.addFacets(_singleAddress(address(dispatchFacet))); + + (bool success, bytes memory data) = address(harness).call(abi.encodeCall(DispatchFacet.fail, (_value))); + + assertEq(success, false, "fallback call"); + assertEq(data, abi.encodeWithSelector(DispatchFacet.DispatchFailure.selector, _value), "revert data"); + } +} diff --git a/test/unit/diamond/mod/fuzz/importSelectors.t.sol b/test/unit/diamond/mod/fuzz/importSelectors.t.sol new file mode 100644 index 00000000..fcd7737b --- /dev/null +++ b/test/unit/diamond/mod/fuzz/importSelectors.t.sol @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import "src/diamond/DiamondMod.sol" as DiamondMod; +import {DiamondMod_Base_Test} from "test/unit/diamond/DiamondModBase.t.sol"; +import { + BadOffsetFacet, + EmptySelectorsFacet, + FacetA, + MisalignedSelectorsFacet, + OversizedLengthFacet, + RevertingSelectorsFacet, + ShortReturnFacet +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract ImportSelectors_DiamondMod_Fuzz_Unit_Test is DiamondMod_Base_Test { + function test_RevertWhen_FacetHasNoBytecode() external { + address noCode = makeAddr("no-code"); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.NoBytecodeAtAddress.selector, noCode)); + harness.addFacets(_singleAddress(noCode)); + } + + function test_RevertWhen_ExportSelectorsCallFails() external { + address badFacet = address(new RevertingSelectorsFacet()); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.FunctionSelectorsCallFailed.selector, badFacet)); + harness.addFacets(_singleAddress(badFacet)); + } + + function test_RevertWhen_NoSelectorsAreReturned() external { + address badFacet = address(new EmptySelectorsFacet()); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.NoSelectorsForFacet.selector, badFacet)); + harness.addFacets(_singleAddress(badFacet)); + } + + function test_RevertWhen_ReturnDataIsTooShort() external { + address badFacet = address(new ShortReturnFacet()); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.IncorrectSelectorsEncoding.selector, badFacet)); + harness.addFacets(_singleAddress(badFacet)); + } + + function test_RevertWhen_ReturnDataHasBadOffset() external { + address badFacet = address(new BadOffsetFacet()); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.IncorrectSelectorsEncoding.selector, badFacet)); + harness.addFacets(_singleAddress(badFacet)); + } + + function test_RevertWhen_SelectorLengthExceedsPayload() external { + address badFacet = address(new OversizedLengthFacet()); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.IncorrectSelectorsEncoding.selector, badFacet)); + harness.addFacets(_singleAddress(badFacet)); + } + + function test_RevertWhen_SelectorLengthIsMisaligned() external { + address badFacet = address(new MisalignedSelectorsFacet()); + + vm.expectRevert(abi.encodeWithSelector(DiamondMod.IncorrectSelectorsEncoding.selector, badFacet)); + harness.addFacets(_singleAddress(badFacet)); + } + + function test_ShouldImportEverySelectorInExportOrder() external { + harness.addFacets(_singleAddress(address(facetA))); + + _assertFacetList(address(harness), FacetA.a1.selector, FacetA.a1.selector, 1, 3); + _assertNode(address(harness), FacetA.a1.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(address(harness), FacetA.a2.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(address(harness), FacetA.a3.selector, address(facetA), bytes4(0), bytes4(0)); + } +} From c4dd2b7533a30e8326a5ab92f6f0863304888e3a Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:33:02 +0530 Subject: [PATCH 04/13] test(diamond): harden fallback forwarding coverage --- test/unit/diamond/mod/fuzz/fallback.t.sol | 29 +++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/test/unit/diamond/mod/fuzz/fallback.t.sol b/test/unit/diamond/mod/fuzz/fallback.t.sol index 1bc9a7f8..724c2eb2 100644 --- a/test/unit/diamond/mod/fuzz/fallback.t.sol +++ b/test/unit/diamond/mod/fuzz/fallback.t.sol @@ -9,6 +9,19 @@ import "src/diamond/DiamondMod.sol" as DiamondMod; import {DiamondMod_Base_Test} from "test/unit/diamond/DiamondModBase.t.sol"; import {DIAMOND_TEST_STORAGE_POSITION, DispatchFacet} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract RawCalldataFacet { + function rawCalldata() external payable returns (bytes memory) { + return msg.data; + } + + function exportSelectors() external pure returns (bytes memory) { + return bytes.concat(this.rawCalldata.selector); + } +} + /** * @dev BTT spec: test/trees/Diamond.tree */ @@ -23,8 +36,23 @@ contract Fallback_DiamondMod_Fuzz_Unit_Test is DiamondMod_Base_Test { assertEq(data, abi.encodeWithSelector(DiamondMod.FunctionNotFound.selector, unknownSelector), "revert data"); } + function test_ShouldForwardNoncanonicalTrailingCalldataByteForByte() external { + RawCalldataFacet rawCalldataFacet = new RawCalldataFacet(); + harness.addFacets(_singleAddress(address(rawCalldataFacet))); + bytes memory callData = bytes.concat(RawCalldataFacet.rawCalldata.selector, hex"decafbad0102030405060708"); + + (bool success, bytes memory data) = address(harness).call(callData); + + assertEq(success, true, "fallback call"); + assertEq(data, abi.encode(callData), "raw returndata"); + bytes memory observedCalldata = abi.decode(data, (bytes)); + assertEq(observedCalldata, callData, "forwarded calldata"); + } + function testFuzz_ShouldPreserveContextAndReturnData(uint96 _value, uint256 _argument, address _caller) external { vm.assume(_caller != address(0)); + _value = uint96(bound(_value, 0, 100 ether)); + _argument = bound(_argument, 0, type(uint128).max); harness.addFacets(_singleAddress(address(dispatchFacet))); vm.deal(_caller, _value); @@ -44,6 +72,7 @@ contract Fallback_DiamondMod_Fuzz_Unit_Test is DiamondMod_Base_Test { } function testFuzz_ShouldBubbleExactRevertData(uint256 _value) external { + _value = bound(_value, 0, type(uint128).max); harness.addFacets(_singleAddress(address(dispatchFacet))); (bool success, bytes memory data) = address(harness).call(abi.encodeCall(DispatchFacet.fail, (_value))); From c1c753fb33642a4abcd66ec10e2e2ac7d40ba924 Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:41:09 +0530 Subject: [PATCH 05/13] test(diamond): cover upgrade facet and module additions --- .../upgrade/facet/fuzz/upgradeDiamond.t.sol | 64 +++++++++++ .../upgrade/mod/fuzz/upgradeDiamond.t.sol | 61 ++++++++++ .../upgrade/shared/AddFacetsBehavior.t.sol | 108 ++++++++++++++++++ 3 files changed, 233 insertions(+) create mode 100644 test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol create mode 100644 test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol create mode 100644 test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol diff --git a/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol b/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol new file mode 100644 index 00000000..d9503fa4 --- /dev/null +++ b/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {DiamondUpgradeFacet} from "src/diamond/DiamondUpgradeFacet.sol"; +import {DiamondUpgrade_Base_Test, Replacement} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; +import {AddFacetsBehavior} from "test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol"; +import {OwnerStorageUtils} from "test/utils/storage/OwnerStorageUtils.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is AddFacetsBehavior { + DiamondUpgradeFacet internal upgradeFacet; + + function setUp() public override(DiamondUpgrade_Base_Test) { + super.setUp(); + upgradeFacet = new DiamondUpgradeFacet(); + target = address(upgradeFacet); + vm.label(target, "DiamondUpgradeFacet"); + OwnerStorageUtils.setOwner(target, users.alice); + vm.stopPrank(); + } + + function _upgrade( + address[] memory _adds, + Replacement[] memory _replacements, + address[] memory _removes, + address _delegate, + bytes memory _delegateCalldata, + bytes32 _tag, + bytes memory _metadata + ) internal override { + DiamondUpgradeFacet.FacetReplacement[] memory replacements = + new DiamondUpgradeFacet.FacetReplacement[](_replacements.length); + for (uint256 i; i < _replacements.length; i++) { + replacements[i] = DiamondUpgradeFacet.FacetReplacement({ + oldFacet: _replacements[i].oldFacet, newFacet: _replacements[i].newFacet + }); + } + + vm.prank(users.alice); + upgradeFacet.upgradeDiamond(_adds, replacements, _removes, _delegate, _delegateCalldata, _tag, _metadata); + } + + function _noBytecodeAtAddressError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.NoBytecodeAtAddress.selector; + } + + function _exportSelectorsCallFailedError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.ExportSelectorsCallFailed.selector; + } + + function _noSelectorsForFacetError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.NoSelectorsForFacet.selector; + } + + function _incorrectSelectorsEncodingError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.IncorrectSelectorsEncoding.selector; + } +} diff --git a/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol b/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol new file mode 100644 index 00000000..82329375 --- /dev/null +++ b/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import "src/diamond/DiamondUpgradeMod.sol" as DiamondUpgradeMod; +import {DiamondUpgrade_Base_Test, Replacement} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; +import {AddFacetsBehavior} from "test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol"; +import {DiamondUpgradeModHarness} from "test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is AddFacetsBehavior { + DiamondUpgradeModHarness internal harness; + + function setUp() public override(DiamondUpgrade_Base_Test) { + super.setUp(); + harness = new DiamondUpgradeModHarness(); + target = address(harness); + vm.label(target, "DiamondUpgradeModHarness"); + } + + function _upgrade( + address[] memory _adds, + Replacement[] memory _replacements, + address[] memory _removes, + address _delegate, + bytes memory _delegateCalldata, + bytes32 _tag, + bytes memory _metadata + ) internal override { + DiamondUpgradeMod.FacetReplacement[] memory replacements = + new DiamondUpgradeMod.FacetReplacement[](_replacements.length); + for (uint256 i; i < _replacements.length; i++) { + replacements[i] = DiamondUpgradeMod.FacetReplacement({ + oldFacet: _replacements[i].oldFacet, newFacet: _replacements[i].newFacet + }); + } + + harness.upgradeDiamond(_adds, replacements, _removes, _delegate, _delegateCalldata, _tag, _metadata); + } + + function _noBytecodeAtAddressError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.NoBytecodeAtAddress.selector; + } + + function _exportSelectorsCallFailedError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.ExportSelectorsCallFailed.selector; + } + + function _noSelectorsForFacetError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.NoSelectorsForFacet.selector; + } + + function _incorrectSelectorsEncodingError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.IncorrectSelectorsEncoding.selector; + } +} diff --git a/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol b/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol new file mode 100644 index 00000000..1eb97521 --- /dev/null +++ b/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {DiamondUpgrade_Base_Test} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; +import { + BadOffsetFacet, + EmptySelectorsFacet, + MisalignedSelectorsFacet, + OversizedLengthFacet, + RevertingSelectorsFacet, + ShortReturnFacet +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; + +abstract contract AddFacetsBehavior is DiamondUpgrade_Base_Test { + function _noBytecodeAtAddressError() internal pure virtual returns (bytes4); + + function _exportSelectorsCallFailedError() internal pure virtual returns (bytes4); + + function _noSelectorsForFacetError() internal pure virtual returns (bytes4); + + function _incorrectSelectorsEncodingError() internal pure virtual returns (bytes4); + + function test_ShouldNotChangeState_WhenAddArrayIsEmpty() external { + _upgrade( + _emptyAddresses(), _emptyReplacements(), _emptyAddresses(), address(0), bytes(""), bytes32(0), bytes("") + ); + _assertCounts(0, 0); + } + + function test_ShouldAddAndLinkThreeFacets() external { + _upgrade( + _threeFacets(address(facetA), address(facetB), address(facetC)), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + _assertThreeNodeList(address(facetA), address(facetB), address(facetC)); + _assertCounts(3, 9); + } + + function test_RevertWhen_AddFacetHasNoBytecode() external { + address badFacet = makeAddr("no-code"); + + vm.expectRevert(abi.encodeWithSelector(_noBytecodeAtAddressError(), badFacet)); + _addFacet(badFacet); + } + + function test_RevertWhen_AddFacetExportSelectorsCallFails() external { + address badFacet = address(new RevertingSelectorsFacet()); + + vm.expectRevert(abi.encodeWithSelector(_exportSelectorsCallFailedError(), badFacet)); + _addFacet(badFacet); + } + + function test_RevertWhen_AddFacetReturnsNoSelectors() external { + address badFacet = address(new EmptySelectorsFacet()); + + vm.expectRevert(abi.encodeWithSelector(_noSelectorsForFacetError(), badFacet)); + _addFacet(badFacet); + } + + function test_RevertWhen_AddFacetSelectorReturnDataIsTooShort() external { + address badFacet = address(new ShortReturnFacet()); + + vm.expectRevert(abi.encodeWithSelector(_incorrectSelectorsEncodingError(), badFacet)); + _addFacet(badFacet); + } + + function test_RevertWhen_AddFacetSelectorReturnDataHasBadOffset() external { + address badFacet = address(new BadOffsetFacet()); + + vm.expectRevert(abi.encodeWithSelector(_incorrectSelectorsEncodingError(), badFacet)); + _addFacet(badFacet); + } + + function test_RevertWhen_AddFacetSelectorLengthExceedsPayload() external { + address badFacet = address(new OversizedLengthFacet()); + + vm.expectRevert(abi.encodeWithSelector(_incorrectSelectorsEncodingError(), badFacet)); + _addFacet(badFacet); + } + + function test_RevertWhen_AddFacetSelectorLengthIsMisaligned() external { + address badFacet = address(new MisalignedSelectorsFacet()); + + vm.expectRevert(abi.encodeWithSelector(_incorrectSelectorsEncodingError(), badFacet)); + _addFacet(badFacet); + } + + function _addFacet(address _facet) private { + _upgrade( + _singleAddress(_facet), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } +} From c6c0a8e0beb4a298ba2616ccac3370ea8bf7049d Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:52:26 +0530 Subject: [PATCH 06/13] test(diamond): cover replace remove and linked-list integrity --- .../upgrade/facet/fuzz/upgradeDiamond.t.sol | 24 +- .../upgrade/mod/fuzz/upgradeDiamond.t.sol | 24 +- .../upgrade/shared/RemoveFacetsBehavior.t.sol | 207 ++++++++++++++++ .../shared/ReplaceFacetsBehavior.t.sol | 224 ++++++++++++++++++ 4 files changed, 477 insertions(+), 2 deletions(-) create mode 100644 test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol create mode 100644 test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol diff --git a/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol b/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol index d9503fa4..fef02e5b 100644 --- a/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol +++ b/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol @@ -8,12 +8,18 @@ pragma solidity >=0.8.30 <0.9.0; import {DiamondUpgradeFacet} from "src/diamond/DiamondUpgradeFacet.sol"; import {DiamondUpgrade_Base_Test, Replacement} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; import {AddFacetsBehavior} from "test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol"; +import {RemoveFacetsBehavior} from "test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol"; +import {ReplaceFacetsBehavior} from "test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol"; import {OwnerStorageUtils} from "test/utils/storage/OwnerStorageUtils.sol"; /** * @dev BTT spec: test/trees/Diamond.tree */ -contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is AddFacetsBehavior { +contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is + AddFacetsBehavior, + ReplaceFacetsBehavior, + RemoveFacetsBehavior +{ DiamondUpgradeFacet internal upgradeFacet; function setUp() public override(DiamondUpgrade_Base_Test) { @@ -61,4 +67,20 @@ contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is AddFacetsBehavior function _incorrectSelectorsEncodingError() internal pure override returns (bytes4) { return DiamondUpgradeFacet.IncorrectSelectorsEncoding.selector; } + + function _cannotReplaceFacetWithSameFacetError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.CannotReplaceFacetWithSameFacet.selector; + } + + function _facetToReplaceDoesNotExistError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.FacetToReplaceDoesNotExist.selector; + } + + function _cannotReplaceFunctionFromNonReplacementFacetError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.CannotReplaceFunctionFromNonReplacementFacet.selector; + } + + function _cannotRemoveFacetThatDoesNotExistError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.CannotRemoveFacetThatDoesNotExist.selector; + } } diff --git a/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol b/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol index 82329375..7ea67e8a 100644 --- a/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol +++ b/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol @@ -8,12 +8,18 @@ pragma solidity >=0.8.30 <0.9.0; import "src/diamond/DiamondUpgradeMod.sol" as DiamondUpgradeMod; import {DiamondUpgrade_Base_Test, Replacement} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; import {AddFacetsBehavior} from "test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol"; +import {RemoveFacetsBehavior} from "test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol"; +import {ReplaceFacetsBehavior} from "test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol"; import {DiamondUpgradeModHarness} from "test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol"; /** * @dev BTT spec: test/trees/Diamond.tree */ -contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is AddFacetsBehavior { +contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is + AddFacetsBehavior, + ReplaceFacetsBehavior, + RemoveFacetsBehavior +{ DiamondUpgradeModHarness internal harness; function setUp() public override(DiamondUpgrade_Base_Test) { @@ -58,4 +64,20 @@ contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is AddFacetsBehavior { function _incorrectSelectorsEncodingError() internal pure override returns (bytes4) { return DiamondUpgradeMod.IncorrectSelectorsEncoding.selector; } + + function _cannotReplaceFacetWithSameFacetError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.CannotReplaceFacetWithSameFacet.selector; + } + + function _facetToReplaceDoesNotExistError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.FacetToReplaceDoesNotExist.selector; + } + + function _cannotReplaceFunctionFromNonReplacementFacetError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.CannotReplaceFunctionFromNonReplacementFacet.selector; + } + + function _cannotRemoveFacetThatDoesNotExistError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.CannotRemoveFacetThatDoesNotExist.selector; + } } diff --git a/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol b/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol new file mode 100644 index 00000000..ec1f6899 --- /dev/null +++ b/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Vm} from "forge-std/Vm.sol"; +import {DiamondUpgrade_Base_Test} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; +import {FacetA, FacetB, FacetC} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +import {DiamondStorageUtils} from "test/utils/storage/DiamondStorageUtils.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { + using DiamondStorageUtils for address; + + bytes32 private constant FACET_REMOVED_TOPIC = keccak256("FacetRemoved(address)"); + + function _cannotRemoveFacetThatDoesNotExistError() internal pure virtual returns (bytes4); + + function test_ShouldNotChangeState_WhenRemoveArrayIsEmpty() external { + _remove(_emptyAddresses()); + + _assertRemovalList(bytes4(0), bytes4(0), 0, 0); + } + + function test_RevertWhen_RemovingFacetThatDoesNotExist() external { + _addRemovalFacet(address(facetA)); + + vm.expectRevert(abi.encodeWithSelector(_cannotRemoveFacetThatDoesNotExistError(), address(facetAReplacement))); + _remove(_singleAddress(address(facetAReplacement))); + + _assertOnlyFacetA(); + } + + function test_ShouldClearList_WhenRemovingOnlyFacet() external { + _addRemovalFacet(address(facetA)); + + _remove(_singleAddress(address(facetA))); + + _assertFacetACleared(); + _assertRemovalList(bytes4(0), bytes4(0), 0, 0); + } + + function test_ShouldRelink_WhenRemovingHeadFacet() external { + _seedRemovalABC(); + + _remove(_singleAddress(address(facetA))); + + _assertFacetACleared(); + _assertNode(FacetB.b1.selector, address(facetB), bytes4(0), FacetC.c1.selector); + _assertNode(FacetC.c1.selector, address(facetC), FacetB.b1.selector, bytes4(0)); + _assertFacetBSelectors(); + _assertFacetCSelectors(); + _assertRemovalList(FacetB.b1.selector, FacetC.c1.selector, 2, 6); + } + + function test_ShouldRelink_WhenRemovingMiddleFacet() external { + _seedRemovalABC(); + + _remove(_singleAddress(address(facetB))); + + _assertFacetBCleared(); + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetC.c1.selector); + _assertNode(FacetC.c1.selector, address(facetC), FacetA.a1.selector, bytes4(0)); + _assertFacetASelectors(); + _assertFacetCSelectors(); + _assertRemovalList(FacetA.a1.selector, FacetC.c1.selector, 2, 6); + } + + function test_ShouldRelink_WhenRemovingTailFacet() external { + _seedRemovalABC(); + + _remove(_singleAddress(address(facetC))); + + _assertFacetCCleared(); + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); + _assertNode(FacetB.b1.selector, address(facetB), FacetA.a1.selector, bytes4(0)); + _assertFacetASelectors(); + _assertFacetBSelectors(); + _assertRemovalList(FacetA.a1.selector, FacetB.b1.selector, 2, 6); + } + + function test_ShouldRemoveMultipleFacetsAndEmitInInputOrder() external { + _seedRemovalABC(); + address[] memory removes = new address[](2); + removes[0] = address(facetC); + removes[1] = address(facetA); + + vm.recordLogs(); + _remove(removes); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 2, "removal log count"); + _assertRemovalLog(logs[0], address(facetC)); + _assertRemovalLog(logs[1], address(facetA)); + _assertFacetACleared(); + _assertFacetCCleared(); + _assertNode(FacetB.b1.selector, address(facetB), bytes4(0), bytes4(0)); + _assertFacetBSelectors(); + _assertRemovalList(FacetB.b1.selector, FacetB.b1.selector, 1, 3); + } + + function test_ShouldRollBackEarlierRemoval_WhenLaterRemovalFails() external { + _seedRemovalABC(); + address[] memory removes = new address[](2); + removes[0] = address(facetA); + removes[1] = address(facetAReplacement); + + vm.expectRevert(abi.encodeWithSelector(_cannotRemoveFacetThatDoesNotExistError(), address(facetAReplacement))); + _remove(removes); + + _assertThreeNodeList(address(facetA), address(facetB), address(facetC)); + _assertFacetASelectors(); + _assertFacetBSelectors(); + _assertFacetCSelectors(); + _assertRemovalList(FacetA.a1.selector, FacetC.c1.selector, 3, 9); + } + + function _remove(address[] memory _facets) private { + _upgrade(_emptyAddresses(), _emptyReplacements(), _facets, address(0), bytes(""), bytes32(0), bytes("")); + } + + function _addRemovalFacet(address _facet) private { + _upgrade( + _singleAddress(_facet), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _seedRemovalABC() private { + _upgrade( + _threeFacets(address(facetA), address(facetB), address(facetC)), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _assertOnlyFacetA() private view { + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), bytes4(0)); + _assertFacetASelectors(); + _assertRemovalList(FacetA.a1.selector, FacetA.a1.selector, 1, 3); + } + + function _assertFacetASelectors() private view { + _assertSelectorOwner(FacetA.a1.selector, address(facetA)); + _assertSelectorOwner(FacetA.a2.selector, address(facetA)); + _assertSelectorOwner(FacetA.a3.selector, address(facetA)); + } + + function _assertFacetBSelectors() private view { + _assertSelectorOwner(FacetB.b1.selector, address(facetB)); + _assertSelectorOwner(FacetB.b2.selector, address(facetB)); + _assertSelectorOwner(FacetB.b3.selector, address(facetB)); + } + + function _assertFacetCSelectors() private view { + _assertSelectorOwner(FacetC.c1.selector, address(facetC)); + _assertSelectorOwner(FacetC.c2.selector, address(facetC)); + _assertSelectorOwner(FacetC.c3.selector, address(facetC)); + } + + function _assertFacetACleared() private view { + _assertNode(FacetA.a1.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetA.a2.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetA.a3.selector, address(0), bytes4(0), bytes4(0)); + } + + function _assertFacetBCleared() private view { + _assertNode(FacetB.b1.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetB.b2.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetB.b3.selector, address(0), bytes4(0), bytes4(0)); + } + + function _assertFacetCCleared() private view { + _assertNode(FacetC.c1.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetC.c2.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetC.c3.selector, address(0), bytes4(0), bytes4(0)); + } + + function _assertRemovalList(bytes4 _head, bytes4 _tail, uint32 _facetCount, uint32 _selectorCount) private view { + (bytes4 head, bytes4 tail, uint32 facetCount, uint32 selectorCount) = target.facetList(); + assertEq(head, _head, "head"); + assertEq(tail, _tail, "tail"); + assertEq(facetCount, _facetCount, "facetCount"); + assertEq(selectorCount, _selectorCount, "selectorCount"); + } + + function _assertRemovalLog(Vm.Log memory _log, address _facet) private view { + assertEq(_log.emitter, target, "removal emitter"); + assertEq(_log.topics.length, 2, "removal topic count"); + assertEq(_log.topics[0], FACET_REMOVED_TOPIC, "removal topic"); + assertEq(address(uint160(uint256(_log.topics[1]))), _facet, "removed facet"); + assertEq(_log.data.length, 0, "removal data length"); + } +} diff --git a/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol b/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol new file mode 100644 index 00000000..22dbbdd4 --- /dev/null +++ b/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol @@ -0,0 +1,224 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Vm} from "forge-std/Vm.sol"; +import {DiamondUpgrade_Base_Test} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; +import { + FacetA, + FacetAChanged, + FacetB, + FacetBChanged, + FacetC, + FacetCChanged +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +import {DiamondStorageUtils} from "test/utils/storage/DiamondStorageUtils.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { + using DiamondStorageUtils for address; + + bytes32 private constant FACET_REPLACED_TOPIC = keccak256("FacetReplaced(address,address)"); + + function _cannotReplaceFacetWithSameFacetError() internal pure virtual returns (bytes4); + + function _facetToReplaceDoesNotExistError() internal pure virtual returns (bytes4); + + function _cannotReplaceFunctionFromNonReplacementFacetError() internal pure virtual returns (bytes4); + + function test_RevertWhen_ReplacingFacetWithSameFacet() external { + vm.expectRevert(abi.encodeWithSelector(_cannotReplaceFacetWithSameFacetError(), address(facetA))); + _replace(address(facetA), address(facetA)); + } + + function test_RevertWhen_FacetToReplaceDoesNotExist() external { + vm.expectRevert(abi.encodeWithSelector(_facetToReplaceDoesNotExistError(), address(facetA))); + _replace(address(facetA), address(facetAReplacement)); + } + + function test_RevertWhen_NewFirstSelectorBelongsToDifferentFacet() external { + _seedReplacementABC(); + + vm.expectRevert( + abi.encodeWithSelector(_cannotReplaceFunctionFromNonReplacementFacetError(), FacetB.b1.selector) + ); + _replace(address(facetA), address(facetBReplacement)); + } + + function test_RevertWhen_NewNonFirstSelectorBelongsToDifferentFacet() external { + _seedReplacementABC(); + + vm.expectRevert( + abi.encodeWithSelector(_cannotReplaceFunctionFromNonReplacementFacetError(), FacetB.b2.selector) + ); + _replace(address(facetA), address(facetBChanged)); + } + + function test_ShouldReplaceInPlace_WhenSelectorSetsAreIdentical() external { + _addReplacementFacet(address(facetA)); + + vm.recordLogs(); + _replace(address(facetA), address(facetAReplacement)); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 1, "replacement log count"); + _assertReplacementLog(logs[0], address(facetA), address(facetAReplacement)); + _assertNode(FacetA.a1.selector, address(facetAReplacement), bytes4(0), bytes4(0)); + _assertNode(FacetA.a2.selector, address(facetAReplacement), bytes4(0), bytes4(0)); + _assertNode(FacetA.a3.selector, address(facetAReplacement), bytes4(0), bytes4(0)); + _assertReplacementList(FacetA.a1.selector, FacetA.a1.selector, 1, 3); + } + + function test_ShouldAddSharedAndRemoveStaleSelectors_WhenSelectorSetsDiffer() external { + _addReplacementFacet(address(facetA)); + + _replace(address(facetA), address(facetAChanged)); + + _assertSelectorOwner(FacetA.a1.selector, address(0)); + _assertSelectorOwner(FacetA.a2.selector, address(facetAChanged)); + _assertSelectorOwner(FacetA.a3.selector, address(0)); + _assertSelectorOwner(FacetAChanged.a4.selector, address(facetAChanged)); + _assertSelectorOwner(FacetAChanged.a5.selector, address(facetAChanged)); + _assertNode(FacetAChanged.a4.selector, address(facetAChanged), bytes4(0), bytes4(0)); + _assertNode(FacetA.a2.selector, address(facetAChanged), bytes4(0), bytes4(0)); + _assertNode(FacetAChanged.a5.selector, address(facetAChanged), bytes4(0), bytes4(0)); + _assertCounts(1, 3); + _assertReplacementList(FacetAChanged.a4.selector, FacetAChanged.a4.selector, 1, 3); + } + + function test_ShouldRelink_WhenReplacingHeadFacet() external { + _seedReplacementABC(); + + _replace(address(facetA), address(facetAChanged)); + + _assertSelectorOwner(FacetA.a1.selector, address(0)); + _assertSelectorOwner(FacetA.a3.selector, address(0)); + _assertNode(FacetAChanged.a4.selector, address(facetAChanged), bytes4(0), FacetB.b1.selector); + _assertNode(FacetA.a2.selector, address(facetAChanged), bytes4(0), bytes4(0)); + _assertNode(FacetAChanged.a5.selector, address(facetAChanged), bytes4(0), bytes4(0)); + _assertNode(FacetB.b1.selector, address(facetB), FacetAChanged.a4.selector, FacetC.c1.selector); + _assertNode(FacetC.c1.selector, address(facetC), FacetB.b1.selector, bytes4(0)); + _assertUnchangedBAndCSelectors(); + _assertReplacementList(FacetAChanged.a4.selector, FacetC.c1.selector, 3, 9); + } + + function test_ShouldRelink_WhenReplacingMiddleFacet() external { + _seedReplacementABC(); + + _replace(address(facetB), address(facetBChanged)); + + _assertSelectorOwner(FacetB.b1.selector, address(0)); + _assertSelectorOwner(FacetB.b3.selector, address(0)); + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetBChanged.b4.selector); + _assertNode(FacetBChanged.b4.selector, address(facetBChanged), FacetA.a1.selector, FacetC.c1.selector); + _assertNode(FacetB.b2.selector, address(facetBChanged), bytes4(0), bytes4(0)); + _assertNode(FacetBChanged.b5.selector, address(facetBChanged), bytes4(0), bytes4(0)); + _assertNode(FacetC.c1.selector, address(facetC), FacetBChanged.b4.selector, bytes4(0)); + _assertUnchangedAAndCSelectors(); + _assertReplacementList(FacetA.a1.selector, FacetC.c1.selector, 3, 9); + } + + function test_ShouldRelink_WhenReplacingTailFacet() external { + _seedReplacementABC(); + + _replace(address(facetC), address(facetCChanged)); + + _assertSelectorOwner(FacetC.c1.selector, address(0)); + _assertSelectorOwner(FacetC.c3.selector, address(0)); + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); + _assertNode(FacetB.b1.selector, address(facetB), FacetA.a1.selector, FacetCChanged.c4.selector); + _assertNode(FacetCChanged.c4.selector, address(facetCChanged), FacetB.b1.selector, bytes4(0)); + _assertNode(FacetC.c2.selector, address(facetCChanged), bytes4(0), bytes4(0)); + _assertNode(FacetCChanged.c5.selector, address(facetCChanged), bytes4(0), bytes4(0)); + _assertUnchangedAAndBSelectors(); + _assertReplacementList(FacetA.a1.selector, FacetCChanged.c4.selector, 3, 9); + } + + function _replace(address _oldFacet, address _newFacet) private { + _upgrade( + _emptyAddresses(), + _singleReplacement(_oldFacet, _newFacet), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _addReplacementFacet(address _facet) private { + _upgrade( + _singleAddress(_facet), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _seedReplacementABC() private { + _upgrade( + _threeFacets(address(facetA), address(facetB), address(facetC)), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _assertReplacementList(bytes4 _head, bytes4 _tail, uint32 _facetCount, uint32 _selectorCount) + private + view + { + (bytes4 head, bytes4 tail, uint32 facetCount, uint32 selectorCount) = target.facetList(); + assertEq(head, _head, "head"); + assertEq(tail, _tail, "tail"); + assertEq(facetCount, _facetCount, "facetCount"); + assertEq(selectorCount, _selectorCount, "selectorCount"); + } + + function _assertReplacementLog(Vm.Log memory _log, address _oldFacet, address _newFacet) private view { + assertEq(_log.emitter, target, "replacement emitter"); + assertEq(_log.topics.length, 3, "replacement topic count"); + assertEq(_log.topics[0], FACET_REPLACED_TOPIC, "replacement topic"); + assertEq(address(uint160(uint256(_log.topics[1]))), _oldFacet, "replacement old facet"); + assertEq(address(uint160(uint256(_log.topics[2]))), _newFacet, "replacement new facet"); + assertEq(_log.data.length, 0, "replacement data length"); + } + + function _assertUnchangedAAndBSelectors() private view { + _assertSelectorOwner(FacetA.a1.selector, address(facetA)); + _assertSelectorOwner(FacetA.a2.selector, address(facetA)); + _assertSelectorOwner(FacetA.a3.selector, address(facetA)); + _assertSelectorOwner(FacetB.b1.selector, address(facetB)); + _assertSelectorOwner(FacetB.b2.selector, address(facetB)); + _assertSelectorOwner(FacetB.b3.selector, address(facetB)); + } + + function _assertUnchangedAAndCSelectors() private view { + _assertSelectorOwner(FacetA.a1.selector, address(facetA)); + _assertSelectorOwner(FacetA.a2.selector, address(facetA)); + _assertSelectorOwner(FacetA.a3.selector, address(facetA)); + _assertSelectorOwner(FacetC.c1.selector, address(facetC)); + _assertSelectorOwner(FacetC.c2.selector, address(facetC)); + _assertSelectorOwner(FacetC.c3.selector, address(facetC)); + } + + function _assertUnchangedBAndCSelectors() private view { + _assertSelectorOwner(FacetB.b1.selector, address(facetB)); + _assertSelectorOwner(FacetB.b2.selector, address(facetB)); + _assertSelectorOwner(FacetB.b3.selector, address(facetB)); + _assertSelectorOwner(FacetC.c1.selector, address(facetC)); + _assertSelectorOwner(FacetC.c2.selector, address(facetC)); + _assertSelectorOwner(FacetC.c3.selector, address(facetC)); + } +} From 588ab5dd4bdcff6a7121e6f1087efa83f20f66c1 Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:00:35 +0530 Subject: [PATCH 07/13] test(diamond): harden replace remove lifecycle coverage --- .../upgrade/shared/RemoveFacetsBehavior.t.sol | 37 ++++++++++-- .../shared/ReplaceFacetsBehavior.t.sol | 57 ++++++++++++++----- 2 files changed, 76 insertions(+), 18 deletions(-) diff --git a/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol b/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol index ec1f6899..d2721b4d 100644 --- a/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol +++ b/test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol @@ -21,9 +21,14 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { function _cannotRemoveFacetThatDoesNotExistError() internal pure virtual returns (bytes4); function test_ShouldNotChangeState_WhenRemoveArrayIsEmpty() external { + _seedRemovalABC(); + + vm.recordLogs(); _remove(_emptyAddresses()); + Vm.Log[] memory logs = vm.getRecordedLogs(); - _assertRemovalList(bytes4(0), bytes4(0), 0, 0); + assertEq(logs.length, 0, "empty removal log count"); + _assertRemovalOriginalABC(); } function test_RevertWhen_RemovingFacetThatDoesNotExist() external { @@ -38,7 +43,7 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldClearList_WhenRemovingOnlyFacet() external { _addRemovalFacet(address(facetA)); - _remove(_singleAddress(address(facetA))); + _removeAndAssertSingleEvent(address(facetA)); _assertFacetACleared(); _assertRemovalList(bytes4(0), bytes4(0), 0, 0); @@ -47,7 +52,7 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldRelink_WhenRemovingHeadFacet() external { _seedRemovalABC(); - _remove(_singleAddress(address(facetA))); + _removeAndAssertSingleEvent(address(facetA)); _assertFacetACleared(); _assertNode(FacetB.b1.selector, address(facetB), bytes4(0), FacetC.c1.selector); @@ -60,7 +65,7 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldRelink_WhenRemovingMiddleFacet() external { _seedRemovalABC(); - _remove(_singleAddress(address(facetB))); + _removeAndAssertSingleEvent(address(facetB)); _assertFacetBCleared(); _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetC.c1.selector); @@ -73,7 +78,7 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldRelink_WhenRemovingTailFacet() external { _seedRemovalABC(); - _remove(_singleAddress(address(facetC))); + _removeAndAssertSingleEvent(address(facetC)); _assertFacetCCleared(); _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); @@ -123,6 +128,15 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { _upgrade(_emptyAddresses(), _emptyReplacements(), _facets, address(0), bytes(""), bytes32(0), bytes("")); } + function _removeAndAssertSingleEvent(address _facet) private { + vm.recordLogs(); + _remove(_singleAddress(_facet)); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 1, "removal log count"); + _assertRemovalLog(logs[0], _facet); + } + function _addRemovalFacet(address _facet) private { _upgrade( _singleAddress(_facet), @@ -204,4 +218,17 @@ abstract contract RemoveFacetsBehavior is DiamondUpgrade_Base_Test { assertEq(address(uint160(uint256(_log.topics[1]))), _facet, "removed facet"); assertEq(_log.data.length, 0, "removal data length"); } + + function _assertRemovalOriginalABC() private view { + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); + _assertNode(FacetA.a2.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(FacetA.a3.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(FacetB.b1.selector, address(facetB), FacetA.a1.selector, FacetC.c1.selector); + _assertNode(FacetB.b2.selector, address(facetB), bytes4(0), bytes4(0)); + _assertNode(FacetB.b3.selector, address(facetB), bytes4(0), bytes4(0)); + _assertNode(FacetC.c1.selector, address(facetC), FacetB.b1.selector, bytes4(0)); + _assertNode(FacetC.c2.selector, address(facetC), bytes4(0), bytes4(0)); + _assertNode(FacetC.c3.selector, address(facetC), bytes4(0), bytes4(0)); + _assertRemovalList(FacetA.a1.selector, FacetC.c1.selector, 3, 9); + } } diff --git a/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol b/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol index 22dbbdd4..fbc00e9a 100644 --- a/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol +++ b/test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol @@ -59,15 +59,24 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { _replace(address(facetA), address(facetBChanged)); } - function test_ShouldReplaceInPlace_WhenSelectorSetsAreIdentical() external { - _addReplacementFacet(address(facetA)); + function test_ShouldNotChangeState_WhenReplacementArrayIsEmpty() external { + _seedReplacementABC(); vm.recordLogs(); - _replace(address(facetA), address(facetAReplacement)); + _upgrade( + _emptyAddresses(), _emptyReplacements(), _emptyAddresses(), address(0), bytes(""), bytes32(0), bytes("") + ); Vm.Log[] memory logs = vm.getRecordedLogs(); - assertEq(logs.length, 1, "replacement log count"); - _assertReplacementLog(logs[0], address(facetA), address(facetAReplacement)); + assertEq(logs.length, 0, "empty replacement log count"); + _assertReplacementOriginalABC(); + } + + function test_ShouldReplaceInPlace_WhenSelectorSetsAreIdentical() external { + _addReplacementFacet(address(facetA)); + + _replaceAndAssertSingleEvent(address(facetA), address(facetAReplacement)); + _assertNode(FacetA.a1.selector, address(facetAReplacement), bytes4(0), bytes4(0)); _assertNode(FacetA.a2.selector, address(facetAReplacement), bytes4(0), bytes4(0)); _assertNode(FacetA.a3.selector, address(facetAReplacement), bytes4(0), bytes4(0)); @@ -77,9 +86,9 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldAddSharedAndRemoveStaleSelectors_WhenSelectorSetsDiffer() external { _addReplacementFacet(address(facetA)); - _replace(address(facetA), address(facetAChanged)); + _replaceAndAssertSingleEvent(address(facetA), address(facetAChanged)); - _assertSelectorOwner(FacetA.a1.selector, address(0)); + _assertNode(FacetA.a1.selector, address(0), bytes4(0), bytes4(0)); _assertSelectorOwner(FacetA.a2.selector, address(facetAChanged)); _assertSelectorOwner(FacetA.a3.selector, address(0)); _assertSelectorOwner(FacetAChanged.a4.selector, address(facetAChanged)); @@ -94,9 +103,9 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldRelink_WhenReplacingHeadFacet() external { _seedReplacementABC(); - _replace(address(facetA), address(facetAChanged)); + _replaceAndAssertSingleEvent(address(facetA), address(facetAChanged)); - _assertSelectorOwner(FacetA.a1.selector, address(0)); + _assertNode(FacetA.a1.selector, address(0), bytes4(0), bytes4(0)); _assertSelectorOwner(FacetA.a3.selector, address(0)); _assertNode(FacetAChanged.a4.selector, address(facetAChanged), bytes4(0), FacetB.b1.selector); _assertNode(FacetA.a2.selector, address(facetAChanged), bytes4(0), bytes4(0)); @@ -110,9 +119,9 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldRelink_WhenReplacingMiddleFacet() external { _seedReplacementABC(); - _replace(address(facetB), address(facetBChanged)); + _replaceAndAssertSingleEvent(address(facetB), address(facetBChanged)); - _assertSelectorOwner(FacetB.b1.selector, address(0)); + _assertNode(FacetB.b1.selector, address(0), bytes4(0), bytes4(0)); _assertSelectorOwner(FacetB.b3.selector, address(0)); _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetBChanged.b4.selector); _assertNode(FacetBChanged.b4.selector, address(facetBChanged), FacetA.a1.selector, FacetC.c1.selector); @@ -126,9 +135,9 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { function test_ShouldRelink_WhenReplacingTailFacet() external { _seedReplacementABC(); - _replace(address(facetC), address(facetCChanged)); + _replaceAndAssertSingleEvent(address(facetC), address(facetCChanged)); - _assertSelectorOwner(FacetC.c1.selector, address(0)); + _assertNode(FacetC.c1.selector, address(0), bytes4(0), bytes4(0)); _assertSelectorOwner(FacetC.c3.selector, address(0)); _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); _assertNode(FacetB.b1.selector, address(facetB), FacetA.a1.selector, FacetCChanged.c4.selector); @@ -151,6 +160,15 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { ); } + function _replaceAndAssertSingleEvent(address _oldFacet, address _newFacet) private { + vm.recordLogs(); + _replace(_oldFacet, _newFacet); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 1, "replacement log count"); + _assertReplacementLog(logs[0], _oldFacet, _newFacet); + } + function _addReplacementFacet(address _facet) private { _upgrade( _singleAddress(_facet), @@ -195,6 +213,19 @@ abstract contract ReplaceFacetsBehavior is DiamondUpgrade_Base_Test { assertEq(_log.data.length, 0, "replacement data length"); } + function _assertReplacementOriginalABC() private view { + _assertNode(FacetA.a1.selector, address(facetA), bytes4(0), FacetB.b1.selector); + _assertNode(FacetA.a2.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(FacetA.a3.selector, address(facetA), bytes4(0), bytes4(0)); + _assertNode(FacetB.b1.selector, address(facetB), FacetA.a1.selector, FacetC.c1.selector); + _assertNode(FacetB.b2.selector, address(facetB), bytes4(0), bytes4(0)); + _assertNode(FacetB.b3.selector, address(facetB), bytes4(0), bytes4(0)); + _assertNode(FacetC.c1.selector, address(facetC), FacetB.b1.selector, bytes4(0)); + _assertNode(FacetC.c2.selector, address(facetC), bytes4(0), bytes4(0)); + _assertNode(FacetC.c3.selector, address(facetC), bytes4(0), bytes4(0)); + _assertReplacementList(FacetA.a1.selector, FacetC.c1.selector, 3, 9); + } + function _assertUnchangedAAndBSelectors() private view { _assertSelectorOwner(FacetA.a1.selector, address(facetA)); _assertSelectorOwner(FacetA.a2.selector, address(facetA)); From ea0d7078adedfcefc2cc268c9c2280fc4302e003 Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:08:10 +0530 Subject: [PATCH 08/13] test(diamond): cover upgrade effects events and authorization --- .../upgrade/facet/fuzz/upgradeDiamond.t.sol | 33 ++- .../upgrade/mod/fuzz/upgradeDiamond.t.sol | 15 +- .../shared/UpgradeEffectsBehavior.t.sol | 236 ++++++++++++++++++ 3 files changed, 280 insertions(+), 4 deletions(-) create mode 100644 test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol diff --git a/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol b/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol index fef02e5b..e62eb60e 100644 --- a/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol +++ b/test/unit/diamond/upgrade/facet/fuzz/upgradeDiamond.t.sol @@ -10,6 +10,7 @@ import {DiamondUpgrade_Base_Test, Replacement} from "test/unit/diamond/DiamondUp import {AddFacetsBehavior} from "test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol"; import {RemoveFacetsBehavior} from "test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol"; import {ReplaceFacetsBehavior} from "test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol"; +import {UpgradeEffectsBehavior} from "test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol"; import {OwnerStorageUtils} from "test/utils/storage/OwnerStorageUtils.sol"; /** @@ -18,7 +19,8 @@ import {OwnerStorageUtils} from "test/utils/storage/OwnerStorageUtils.sol"; contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is AddFacetsBehavior, ReplaceFacetsBehavior, - RemoveFacetsBehavior + RemoveFacetsBehavior, + UpgradeEffectsBehavior { DiamondUpgradeFacet internal upgradeFacet; @@ -52,7 +54,12 @@ contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is upgradeFacet.upgradeDiamond(_adds, replacements, _removes, _delegate, _delegateCalldata, _tag, _metadata); } - function _noBytecodeAtAddressError() internal pure override returns (bytes4) { + function _noBytecodeAtAddressError() + internal + pure + override(AddFacetsBehavior, UpgradeEffectsBehavior) + returns (bytes4) + { return DiamondUpgradeFacet.NoBytecodeAtAddress.selector; } @@ -83,4 +90,26 @@ contract UpgradeDiamond_DiamondUpgradeFacet_Fuzz_Unit_Test is function _cannotRemoveFacetThatDoesNotExistError() internal pure override returns (bytes4) { return DiamondUpgradeFacet.CannotRemoveFacetThatDoesNotExist.selector; } + + function _delegateCallRevertedError() internal pure override returns (bytes4) { + return DiamondUpgradeFacet.DelegateCallReverted.selector; + } + + function test_RevertWhen_CallerIsNotOwner() external { + setMsgSender(users.bob); + vm.expectRevert(DiamondUpgradeFacet.OwnerUnauthorizedAccount.selector); + upgradeFacet.upgradeDiamond( + new address[](0), + new DiamondUpgradeFacet.FacetReplacement[](0), + new address[](0), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function test_ShouldExportOnlyUpgradeDiamondSelector() external view { + assertEq(upgradeFacet.exportSelectors(), bytes.concat(DiamondUpgradeFacet.upgradeDiamond.selector)); + } } diff --git a/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol b/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol index 7ea67e8a..ea22566e 100644 --- a/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol +++ b/test/unit/diamond/upgrade/mod/fuzz/upgradeDiamond.t.sol @@ -10,6 +10,7 @@ import {DiamondUpgrade_Base_Test, Replacement} from "test/unit/diamond/DiamondUp import {AddFacetsBehavior} from "test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol"; import {RemoveFacetsBehavior} from "test/unit/diamond/upgrade/shared/RemoveFacetsBehavior.t.sol"; import {ReplaceFacetsBehavior} from "test/unit/diamond/upgrade/shared/ReplaceFacetsBehavior.t.sol"; +import {UpgradeEffectsBehavior} from "test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol"; import {DiamondUpgradeModHarness} from "test/utils/harnesses/diamond/DiamondUpgradeModHarness.sol"; /** @@ -18,7 +19,8 @@ import {DiamondUpgradeModHarness} from "test/utils/harnesses/diamond/DiamondUpgr contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is AddFacetsBehavior, ReplaceFacetsBehavior, - RemoveFacetsBehavior + RemoveFacetsBehavior, + UpgradeEffectsBehavior { DiamondUpgradeModHarness internal harness; @@ -49,7 +51,12 @@ contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is harness.upgradeDiamond(_adds, replacements, _removes, _delegate, _delegateCalldata, _tag, _metadata); } - function _noBytecodeAtAddressError() internal pure override returns (bytes4) { + function _noBytecodeAtAddressError() + internal + pure + override(AddFacetsBehavior, UpgradeEffectsBehavior) + returns (bytes4) + { return DiamondUpgradeMod.NoBytecodeAtAddress.selector; } @@ -80,4 +87,8 @@ contract UpgradeDiamond_DiamondUpgradeMod_Fuzz_Unit_Test is function _cannotRemoveFacetThatDoesNotExistError() internal pure override returns (bytes4) { return DiamondUpgradeMod.CannotRemoveFacetThatDoesNotExist.selector; } + + function _delegateCallRevertedError() internal pure override returns (bytes4) { + return DiamondUpgradeMod.DelegateCallReverted.selector; + } } diff --git a/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol b/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol new file mode 100644 index 00000000..4847746e --- /dev/null +++ b/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {Vm} from "forge-std/Vm.sol"; +import {DiamondUpgrade_Base_Test} from "test/unit/diamond/DiamondUpgradeBase.t.sol"; +import { + DIAMOND_TEST_STORAGE_POSITION, + DelegateTarget, + FacetA, + FacetB, + FacetC +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +import {DiamondStorageUtils} from "test/utils/storage/DiamondStorageUtils.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +abstract contract UpgradeEffectsBehavior is DiamondUpgrade_Base_Test { + bytes32 private constant FACET_ADDED_TOPIC = keccak256("FacetAdded(address)"); + bytes32 private constant FACET_REPLACED_TOPIC = keccak256("FacetReplaced(address,address)"); + bytes32 private constant FACET_REMOVED_TOPIC = keccak256("FacetRemoved(address)"); + bytes32 private constant DIAMOND_DELEGATE_CALL_TOPIC = keccak256("DiamondDelegateCall(address,bytes)"); + bytes32 private constant DIAMOND_METADATA_TOPIC = keccak256("DiamondMetadata(bytes32,bytes)"); + + function _noBytecodeAtAddressError() internal pure virtual returns (bytes4); + + function _delegateCallRevertedError() internal pure virtual returns (bytes4); + + function testFuzz_ShouldDelegatecallAndEmit(uint256 _value) external { + bytes memory callData = abi.encodeCall(DelegateTarget.initialize, (_value)); + vm.expectEmit(target); + emit DiamondDelegateCall(address(delegateTarget), callData); + + _upgrade( + _emptyAddresses(), + _emptyReplacements(), + _emptyAddresses(), + address(delegateTarget), + callData, + bytes32(0), + bytes("") + ); + + assertEq(uint256(vm.load(target, DIAMOND_TEST_STORAGE_POSITION)), _value, "target state"); + assertEq(uint256(vm.load(address(delegateTarget), DIAMOND_TEST_STORAGE_POSITION)), 0, "delegate state"); + } + + function test_RevertWhen_DelegateHasNoBytecode() external { + address noCodeDelegate = makeAddr("no-code delegate"); + + vm.expectRevert(abi.encodeWithSelector(_noBytecodeAtAddressError(), noCodeDelegate)); + _upgrade( + _emptyAddresses(), _emptyReplacements(), _emptyAddresses(), noCodeDelegate, bytes(""), bytes32(0), bytes("") + ); + } + + function testFuzz_RevertWhen_DelegateBubblesError(uint256 _value) external { + bytes memory callData = abi.encodeCall(DelegateTarget.failWithData, (_value)); + + vm.expectRevert(abi.encodeWithSelector(DelegateTarget.DelegateFailure.selector, _value)); + _upgrade( + _emptyAddresses(), + _emptyReplacements(), + _emptyAddresses(), + address(delegateTarget), + callData, + bytes32(0), + bytes("") + ); + } + + function test_RevertWhen_DelegateRevertsWithoutData() external { + bytes memory callData = abi.encodeCall(DelegateTarget.failWithoutData, ()); + + vm.expectRevert(abi.encodeWithSelector(_delegateCallRevertedError(), address(delegateTarget), callData)); + _upgrade( + _emptyAddresses(), + _emptyReplacements(), + _emptyAddresses(), + address(delegateTarget), + callData, + bytes32(0), + bytes("") + ); + } + + function testFuzz_ShouldNotDelegatecallOrEmit_WhenDelegateIsZero(bytes calldata _delegateCalldata) external { + bytes32 tag = keccak256("zero-delegate"); + + vm.recordLogs(); + _upgrade( + _emptyAddresses(), _emptyReplacements(), _emptyAddresses(), address(0), _delegateCalldata, tag, bytes("") + ); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 1, "zero delegate log count"); + _assertMetadataLog(logs[0], tag, bytes("")); + assertNotEq(logs[0].topics[0], DIAMOND_DELEGATE_CALL_TOPIC, "delegate event omitted"); + } + + function test_ShouldNotEmitMetadata_WhenTagAndMetadataAreEmpty() external { + vm.recordLogs(); + _upgrade( + _emptyAddresses(), _emptyReplacements(), _emptyAddresses(), address(0), bytes(""), bytes32(0), bytes("") + ); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 0, "zero tag empty metadata log count"); + } + + function testFuzz_ShouldEmitMetadata_WhenTagIsNonzeroAndMetadataIsEmpty(bytes32 _tag) external { + vm.assume(_tag != bytes32(0)); + + _upgradeAndAssertSingleMetadata(_tag, bytes("")); + } + + function testFuzz_ShouldEmitMetadata_WhenTagIsZeroAndMetadataIsNonempty(bytes calldata _metadata) external { + vm.assume(_metadata.length > 0); + + _upgradeAndAssertSingleMetadata(bytes32(0), _metadata); + } + + function testFuzz_ShouldEmitMetadata_WhenTagAndMetadataAreNonempty(bytes32 _tag, bytes calldata _metadata) + external + { + vm.assume(_tag != bytes32(0)); + vm.assume(_metadata.length > 0); + + _upgradeAndAssertSingleMetadata(_tag, _metadata); + } + + function test_ShouldApplyUpgradeEffectsInStrictOrder() external { + address[] memory seedFacets = new address[](2); + seedFacets[0] = address(facetA); + seedFacets[1] = address(facetB); + _upgrade(seedFacets, _emptyReplacements(), _emptyAddresses(), address(0), bytes(""), bytes32(0), bytes("")); + + bytes memory callData = abi.encodeCall(DelegateTarget.initialize, (42)); + bytes32 tag = keccak256("v2"); + bytes memory metadata = abi.encode("issue-339"); + + vm.recordLogs(); + _upgrade( + _singleAddress(address(facetC)), + _singleReplacement(address(facetA), address(facetAReplacement)), + _singleAddress(address(facetB)), + address(delegateTarget), + callData, + tag, + metadata + ); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 5, "compound upgrade log count"); + _assertFacetAddedLog(logs[0], address(facetC)); + _assertFacetReplacedLog(logs[1], address(facetA), address(facetAReplacement)); + _assertFacetRemovedLog(logs[2], address(facetB)); + _assertDelegateCallLog(logs[3], address(delegateTarget), callData); + _assertMetadataLog(logs[4], tag, metadata); + _assertCompoundUpgradeState(); + assertEq(uint256(vm.load(target, DIAMOND_TEST_STORAGE_POSITION)), 42, "compound delegate state"); + assertEq(uint256(vm.load(address(delegateTarget), DIAMOND_TEST_STORAGE_POSITION)), 0, "delegate target state"); + } + + function _upgradeAndAssertSingleMetadata(bytes32 _tag, bytes memory _metadata) private { + vm.recordLogs(); + _upgrade(_emptyAddresses(), _emptyReplacements(), _emptyAddresses(), address(0), bytes(""), _tag, _metadata); + Vm.Log[] memory logs = vm.getRecordedLogs(); + + assertEq(logs.length, 1, "metadata log count"); + _assertMetadataLog(logs[0], _tag, _metadata); + } + + function _assertCompoundUpgradeState() private view { + (bytes4 head, bytes4 tail, uint32 facetCount, uint32 selectorCount) = DiamondStorageUtils.facetList(target); + assertEq(head, FacetA.a1.selector, "compound head"); + assertEq(tail, FacetC.c1.selector, "compound tail"); + assertEq(facetCount, 2, "compound facetCount"); + assertEq(selectorCount, 6, "compound selectorCount"); + + _assertNode(FacetA.a1.selector, address(facetAReplacement), bytes4(0), FacetC.c1.selector); + _assertNode(FacetC.c1.selector, address(facetC), FacetA.a1.selector, bytes4(0)); + _assertSelectorOwner(FacetA.a2.selector, address(facetAReplacement)); + _assertSelectorOwner(FacetA.a3.selector, address(facetAReplacement)); + _assertNode(FacetB.b1.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetB.b2.selector, address(0), bytes4(0), bytes4(0)); + _assertNode(FacetB.b3.selector, address(0), bytes4(0), bytes4(0)); + _assertSelectorOwner(FacetC.c2.selector, address(facetC)); + _assertSelectorOwner(FacetC.c3.selector, address(facetC)); + } + + function _assertFacetAddedLog(Vm.Log memory _log, address _facet) private view { + assertEq(_log.emitter, target, "added emitter"); + assertEq(_log.topics.length, 2, "added topic count"); + assertEq(_log.topics[0], FACET_ADDED_TOPIC, "added topic"); + assertEq(address(uint160(uint256(_log.topics[1]))), _facet, "added facet"); + assertEq(_log.data.length, 0, "added data length"); + } + + function _assertFacetReplacedLog(Vm.Log memory _log, address _oldFacet, address _newFacet) private view { + assertEq(_log.emitter, target, "replaced emitter"); + assertEq(_log.topics.length, 3, "replaced topic count"); + assertEq(_log.topics[0], FACET_REPLACED_TOPIC, "replaced topic"); + assertEq(address(uint160(uint256(_log.topics[1]))), _oldFacet, "replaced old facet"); + assertEq(address(uint160(uint256(_log.topics[2]))), _newFacet, "replaced new facet"); + assertEq(_log.data.length, 0, "replaced data length"); + } + + function _assertFacetRemovedLog(Vm.Log memory _log, address _facet) private view { + assertEq(_log.emitter, target, "removed emitter"); + assertEq(_log.topics.length, 2, "removed topic count"); + assertEq(_log.topics[0], FACET_REMOVED_TOPIC, "removed topic"); + assertEq(address(uint160(uint256(_log.topics[1]))), _facet, "removed facet"); + assertEq(_log.data.length, 0, "removed data length"); + } + + function _assertDelegateCallLog(Vm.Log memory _log, address _delegate, bytes memory _callData) private view { + assertEq(_log.emitter, target, "delegate emitter"); + assertEq(_log.topics.length, 2, "delegate topic count"); + assertEq(_log.topics[0], DIAMOND_DELEGATE_CALL_TOPIC, "delegate topic"); + assertEq(address(uint160(uint256(_log.topics[1]))), _delegate, "delegate address"); + assertEq(_log.data, abi.encode(_callData), "delegate calldata"); + } + + function _assertMetadataLog(Vm.Log memory _log, bytes32 _tag, bytes memory _metadata) private view { + assertEq(_log.emitter, target, "metadata emitter"); + assertEq(_log.topics.length, 2, "metadata topic count"); + assertEq(_log.topics[0], DIAMOND_METADATA_TOPIC, "metadata topic"); + assertEq(_log.topics[1], _tag, "metadata tag"); + assertEq(_log.data, abi.encode(_metadata), "metadata data"); + } +} From abb7537df859f959ccaf723bb562515eaee387fe Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:25:32 +0530 Subject: [PATCH 09/13] test(diamond): require exact delegate event count --- .../diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol b/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol index 4847746e..0a8febe9 100644 --- a/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol +++ b/test/unit/diamond/upgrade/shared/UpgradeEffectsBehavior.t.sol @@ -32,9 +32,8 @@ abstract contract UpgradeEffectsBehavior is DiamondUpgrade_Base_Test { function testFuzz_ShouldDelegatecallAndEmit(uint256 _value) external { bytes memory callData = abi.encodeCall(DelegateTarget.initialize, (_value)); - vm.expectEmit(target); - emit DiamondDelegateCall(address(delegateTarget), callData); + vm.recordLogs(); _upgrade( _emptyAddresses(), _emptyReplacements(), @@ -44,7 +43,10 @@ abstract contract UpgradeEffectsBehavior is DiamondUpgrade_Base_Test { bytes32(0), bytes("") ); + Vm.Log[] memory logs = vm.getRecordedLogs(); + assertEq(logs.length, 1, "delegate log count"); + _assertDelegateCallLog(logs[0], address(delegateTarget), callData); assertEq(uint256(vm.load(target, DIAMOND_TEST_STORAGE_POSITION)), _value, "target state"); assertEq(uint256(vm.load(address(delegateTarget), DIAMOND_TEST_STORAGE_POSITION)), 0, "delegate state"); } From d1ef584b99370d6ba8def6b109cfa59558689f92 Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:33:23 +0530 Subject: [PATCH 10/13] fix(diamond): validate facet ownership during inspection --- src/diamond/DiamondInspectFacet.sol | 2 +- test/integration/diamond/DiamondBase.t.sol | 165 ++++++++++++++++++++ test/integration/diamond/fuzz/inspect.t.sol | 141 +++++++++++++++++ 3 files changed, 307 insertions(+), 1 deletion(-) create mode 100644 test/integration/diamond/DiamondBase.t.sol create mode 100644 test/integration/diamond/fuzz/inspect.t.sol diff --git a/src/diamond/DiamondInspectFacet.sol b/src/diamond/DiamondInspectFacet.sol index 4f1c4490..8069a9df 100644 --- a/src/diamond/DiamondInspectFacet.sol +++ b/src/diamond/DiamondInspectFacet.sol @@ -114,7 +114,7 @@ contract DiamondInspectFacet { function facetFunctionSelectors(address _facet) external view returns (bytes4[] memory facetSelectors) { DiamondStorage storage s = getStorage(); facetSelectors = unpackSelectors(IFacet(_facet).exportSelectors()); - if (facetSelectors.length == 0 || s.facetNodes[facetSelectors[0]].facet == address(0)) { + if (facetSelectors.length == 0 || s.facetNodes[facetSelectors[0]].facet != _facet) { facetSelectors = new bytes4[](0); } } diff --git a/test/integration/diamond/DiamondBase.t.sol b/test/integration/diamond/DiamondBase.t.sol new file mode 100644 index 00000000..fa330c1f --- /dev/null +++ b/test/integration/diamond/DiamondBase.t.sol @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {DiamondInspectFacet} from "src/diamond/DiamondInspectFacet.sol"; +import {DiamondUpgradeFacet} from "src/diamond/DiamondUpgradeFacet.sol"; +import {Base_Test} from "test/Base.t.sol"; +import {DiamondProxyHarness} from "test/utils/harnesses/diamond/DiamondProxyHarness.sol"; +import {FacetA, FacetAReplacement, FacetB, FacetC} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; + +interface IDiamondInspect { + struct Facet { + address facet; + bytes4[] functionSelectors; + } + + function facetAddress(bytes4 _selector) external view returns (address); + + function facetFunctionSelectors(address _facet) external view returns (bytes4[] memory); + + function facetAddresses() external view returns (address[] memory); + + function facets() external view returns (Facet[] memory); +} + +interface IDiamondUpgrade { + struct FacetReplacement { + address oldFacet; + address newFacet; + } + + function upgradeDiamond( + address[] calldata _adds, + FacetReplacement[] calldata _replacements, + address[] calldata _removes, + address _delegate, + bytes calldata _delegateCalldata, + bytes32 _tag, + bytes calldata _metadata + ) external; +} + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +abstract contract Diamond_Base_Test is Base_Test { + DiamondProxyHarness internal proxy; + DiamondUpgradeFacet internal upgradeFacet; + DiamondInspectFacet internal inspectFacet; + DiamondInspectFacet internal isolatedInspectFacet; + + IDiamondUpgrade internal upgrade; + IDiamondInspect internal inspect; + + FacetA internal facetA; + FacetAReplacement internal facetAReplacement; + FacetB internal facetB; + FacetC internal facetC; + + function setUp() public virtual override { + Base_Test.setUp(); + + upgradeFacet = new DiamondUpgradeFacet(); + inspectFacet = new DiamondInspectFacet(); + isolatedInspectFacet = new DiamondInspectFacet(); + facetA = new FacetA(); + facetAReplacement = new FacetAReplacement(); + facetB = new FacetB(); + facetC = new FacetC(); + + address[] memory baselineFacets = new address[](2); + baselineFacets[0] = address(upgradeFacet); + baselineFacets[1] = address(inspectFacet); + proxy = new DiamondProxyHarness(baselineFacets, users.alice); + + upgrade = IDiamondUpgrade(address(proxy)); + inspect = IDiamondInspect(address(proxy)); + } + + function _addFacet(address _facet) internal { + upgrade.upgradeDiamond( + _singleAddress(_facet), + _emptyReplacements(), + _emptyAddresses(), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _addFacets(address _a, address _b, address _c) internal { + address[] memory adds = new address[](3); + adds[0] = _a; + adds[1] = _b; + adds[2] = _c; + upgrade.upgradeDiamond( + adds, _emptyReplacements(), _emptyAddresses(), address(0), bytes(""), bytes32(0), bytes("") + ); + } + + function _replaceFacet(address _oldFacet, address _newFacet) internal { + IDiamondUpgrade.FacetReplacement[] memory replacements = new IDiamondUpgrade.FacetReplacement[](1); + replacements[0] = IDiamondUpgrade.FacetReplacement({oldFacet: _oldFacet, newFacet: _newFacet}); + upgrade.upgradeDiamond( + _emptyAddresses(), replacements, _emptyAddresses(), address(0), bytes(""), bytes32(0), bytes("") + ); + } + + function _removeFacet(address _facet) internal { + upgrade.upgradeDiamond( + _emptyAddresses(), + _emptyReplacements(), + _singleAddress(_facet), + address(0), + bytes(""), + bytes32(0), + bytes("") + ); + } + + function _emptyAddresses() internal pure returns (address[] memory values) { + values = new address[](0); + } + + function _emptyReplacements() internal pure returns (IDiamondUpgrade.FacetReplacement[] memory values) { + values = new IDiamondUpgrade.FacetReplacement[](0); + } + + function _singleAddress(address _value) internal pure returns (address[] memory values) { + values = new address[](1); + values[0] = _value; + } + + function _assertSelectors(bytes4[] memory _actual, bytes memory _expectedPacked, string memory _message) + internal + pure + { + uint256 expectedLength = _expectedPacked.length / 4; + assertEq(_actual.length, expectedLength, string.concat(_message, " length")); + for (uint256 i; i < expectedLength; i++) { + bytes4 expected; + assembly ("memory-safe") { + expected := mload(add(add(_expectedPacked, 0x20), mul(i, 4))) + } + assertEq(_actual[i], expected, string.concat(_message, " selector")); + } + } + + function _upgradeSelectors() internal pure returns (bytes memory) { + return bytes.concat(DiamondUpgradeFacet.upgradeDiamond.selector); + } + + function _inspectSelectors() internal pure returns (bytes memory) { + return bytes.concat( + DiamondInspectFacet.facetAddress.selector, + DiamondInspectFacet.facetFunctionSelectors.selector, + DiamondInspectFacet.facetAddresses.selector, + DiamondInspectFacet.facets.selector + ); + } +} diff --git a/test/integration/diamond/fuzz/inspect.t.sol b/test/integration/diamond/fuzz/inspect.t.sol new file mode 100644 index 00000000..9eff0a97 --- /dev/null +++ b/test/integration/diamond/fuzz/inspect.t.sol @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {DiamondInspectFacet} from "src/diamond/DiamondInspectFacet.sol"; +import {Diamond_Base_Test, IDiamondInspect} from "test/integration/diamond/DiamondBase.t.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract Inspect_DiamondInspectFacet_Fuzz_Integration_Test is Diamond_Base_Test { + function test_ShouldReturnEmptyViews_WhenNoFacetsAreRegistered() external view { + assertEq(isolatedInspectFacet.facetAddress(bytes4(keccak256("unknown()"))), address(0), "unknown selector"); + assertEq(isolatedInspectFacet.facetAddresses().length, 0, "facet addresses"); + assertEq(isolatedInspectFacet.facets().length, 0, "facets"); + } + + function test_ShouldReturnEmptySelectors_ForValidUnregisteredFacet() external view { + bytes4[] memory selectors = inspect.facetFunctionSelectors(address(facetA)); + + assertEq(selectors.length, 0, "unregistered facet selectors"); + } + + function test_ShouldReturnFacetAddress_ForEveryRegisteredSelector() external { + _addFacets(address(facetA), address(facetB), address(facetC)); + + _assertFacetOwnsPackedSelectors(address(upgradeFacet), _upgradeSelectors()); + _assertFacetOwnsPackedSelectors(address(inspectFacet), _inspectSelectors()); + _assertFacetOwnsPackedSelectors(address(facetA), facetA.exportSelectors()); + _assertFacetOwnsPackedSelectors(address(facetB), facetB.exportSelectors()); + _assertFacetOwnsPackedSelectors(address(facetC), facetC.exportSelectors()); + assertEq(inspect.facetAddress(bytes4(keccak256("unknown()"))), address(0), "unknown selector"); + } + + function test_ShouldReturnFacetSelectors_InExportOrder() external { + _addFacets(address(facetA), address(facetB), address(facetC)); + + _assertSelectors(inspect.facetFunctionSelectors(address(upgradeFacet)), _upgradeSelectors(), "upgrade"); + _assertSelectors(inspect.facetFunctionSelectors(address(inspectFacet)), _inspectSelectors(), "inspect"); + _assertSelectors(inspect.facetFunctionSelectors(address(facetA)), facetA.exportSelectors(), "facet A"); + _assertSelectors(inspect.facetFunctionSelectors(address(facetB)), facetB.exportSelectors(), "facet B"); + _assertSelectors(inspect.facetFunctionSelectors(address(facetC)), facetC.exportSelectors(), "facet C"); + } + + function test_ShouldReturnFacetAddresses_InLinkedListOrder() external { + _addFacets(address(facetA), address(facetB), address(facetC)); + + address[] memory expected = new address[](5); + expected[0] = address(upgradeFacet); + expected[1] = address(inspectFacet); + expected[2] = address(facetA); + expected[3] = address(facetB); + expected[4] = address(facetC); + assertEq(inspect.facetAddresses(), expected, "facet address order"); + } + + function test_ShouldReturnFacets_PairedWithExactSelectors() external { + _addFacets(address(facetA), address(facetB), address(facetC)); + + IDiamondInspect.Facet[] memory inspectedFacets = inspect.facets(); + assertEq(inspectedFacets.length, 5, "facet count"); + _assertInspectedFacet(inspectedFacets[0], address(upgradeFacet), _upgradeSelectors(), "upgrade"); + _assertInspectedFacet(inspectedFacets[1], address(inspectFacet), _inspectSelectors(), "inspect"); + _assertInspectedFacet(inspectedFacets[2], address(facetA), facetA.exportSelectors(), "facet A"); + _assertInspectedFacet(inspectedFacets[3], address(facetB), facetB.exportSelectors(), "facet B"); + _assertInspectedFacet(inspectedFacets[4], address(facetC), facetC.exportSelectors(), "facet C"); + } + + function test_ShouldExcludeRemovedFacetAndSelectors_FromEveryInspectionView() external { + _addFacets(address(facetA), address(facetB), address(facetC)); + _removeFacet(address(facetB)); + + bytes memory removedSelectors = facetB.exportSelectors(); + for (uint256 i; i < removedSelectors.length / 4; i++) { + assertEq(inspect.facetAddress(_selectorAt(removedSelectors, i)), address(0), "removed selector owner"); + } + assertEq(inspect.facetFunctionSelectors(address(facetB)).length, 0, "removed facet selectors"); + + address[] memory expectedAddresses = new address[](4); + expectedAddresses[0] = address(upgradeFacet); + expectedAddresses[1] = address(inspectFacet); + expectedAddresses[2] = address(facetA); + expectedAddresses[3] = address(facetC); + assertEq(inspect.facetAddresses(), expectedAddresses, "addresses after removal"); + + IDiamondInspect.Facet[] memory inspectedFacets = inspect.facets(); + assertEq(inspectedFacets.length, 4, "facets after removal"); + _assertInspectedFacet(inspectedFacets[0], address(upgradeFacet), _upgradeSelectors(), "upgrade"); + _assertInspectedFacet(inspectedFacets[1], address(inspectFacet), _inspectSelectors(), "inspect"); + _assertInspectedFacet(inspectedFacets[2], address(facetA), facetA.exportSelectors(), "facet A"); + _assertInspectedFacet(inspectedFacets[3], address(facetC), facetC.exportSelectors(), "facet C"); + } + + function test_ShouldReturnEmptySelectors_ForFacetReplacedWithSameSelectorSet() external { + _addFacet(address(facetA)); + _replaceFacet(address(facetA), address(facetAReplacement)); + + bytes4[] memory oldSelectors = inspect.facetFunctionSelectors(address(facetA)); + bytes4[] memory newSelectors = inspect.facetFunctionSelectors(address(facetAReplacement)); + + assertEq(oldSelectors.length, 0, "old facet selectors"); + assertEq(newSelectors.length, 3, "replacement selectors"); + } + + function test_ShouldExportInspectionSelectors_InDeclarationOrder() external view { + assertEq( + inspectFacet.exportSelectors(), + bytes.concat( + DiamondInspectFacet.facetAddress.selector, + DiamondInspectFacet.facetFunctionSelectors.selector, + DiamondInspectFacet.facetAddresses.selector, + DiamondInspectFacet.facets.selector + ) + ); + } + + function _assertFacetOwnsPackedSelectors(address _facet, bytes memory _packedSelectors) private view { + for (uint256 i; i < _packedSelectors.length / 4; i++) { + assertEq(inspect.facetAddress(_selectorAt(_packedSelectors, i)), _facet, "selector owner"); + } + } + + function _assertInspectedFacet( + IDiamondInspect.Facet memory _actual, + address _expectedAddress, + bytes memory _expectedSelectors, + string memory _message + ) private pure { + assertEq(_actual.facet, _expectedAddress, string.concat(_message, " address")); + _assertSelectors(_actual.functionSelectors, _expectedSelectors, _message); + } + + function _selectorAt(bytes memory _packedSelectors, uint256 _index) private pure returns (bytes4 selector) { + assembly ("memory-safe") { + selector := mload(add(add(_packedSelectors, 0x20), mul(_index, 4))) + } + } +} From 9833f16714e61420b9af35da09386e990c6b615f Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:42:38 +0530 Subject: [PATCH 11/13] test(diamond): fuzz upgrade sequences and invariants --- test/integration/diamond/fuzz/sequences.t.sol | 326 ++++++++++++++++++ 1 file changed, 326 insertions(+) create mode 100644 test/integration/diamond/fuzz/sequences.t.sol diff --git a/test/integration/diamond/fuzz/sequences.t.sol b/test/integration/diamond/fuzz/sequences.t.sol new file mode 100644 index 00000000..8f9ab3f2 --- /dev/null +++ b/test/integration/diamond/fuzz/sequences.t.sol @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.30 <0.9.0; + +/* Compose + * https://compose.diamonds + */ + +import {DiamondInspectFacet} from "src/diamond/DiamondInspectFacet.sol"; +import {DiamondUpgradeFacet} from "src/diamond/DiamondUpgradeFacet.sol"; +import {Diamond_Base_Test, IDiamondInspect} from "test/integration/diamond/DiamondBase.t.sol"; +import { + FacetA, + FacetAReplacement, + FacetB, + FacetBReplacement, + FacetC, + FacetCReplacement +} from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +import {DiamondStorageUtils} from "test/utils/storage/DiamondStorageUtils.sol"; + +/** + * @dev BTT spec: test/trees/Diamond.tree + */ +contract Sequences_Diamond_Fuzz_Integration_Test is Diamond_Base_Test { + uint256 private constant MAX_ACTIONS = 24; + uint256 private constant BASELINE_FACET_COUNT = 2; + uint256 private constant BASELINE_SELECTOR_COUNT = 5; + + struct FamilyModel { + address original; + address replacement; + address currentFacet; + bytes selectors; + bool registered; + } + + struct SequenceModel { + FamilyModel[3] families; + uint8[3] order; + uint8 activeCount; + } + + FacetBReplacement internal facetBReplacement; + FacetCReplacement internal facetCReplacement; + + function setUp() public override { + super.setUp(); + + facetBReplacement = new FacetBReplacement(); + facetCReplacement = new FacetCReplacement(); + } + + function test_AddReplaceRemove_AllUserFacets() external { + _runDeterministic(hex"000306010407020508"); + } + + function test_RemoveHeadThenAppend_ReusesValidTailLinks() external { + _runDeterministic(hex"0003060200"); + } + + function test_RemoveTailThenAppend_UsesPreviousTail() external { + _runDeterministic(hex"0003060806"); + } + + function test_ReplaceMiddleThenRemoveNeighbors_PreservesSingleton() external { + _runDeterministic(hex"000306040208"); + } + + function test_RemoveAllThenReAddInDifferentOrder_RebuildsList() external { + _runDeterministic(hex"000306020508060003"); + } + + function testFuzz_AddReplaceRemove_SequenceMatchesModel(bytes calldata _actions) external { + SequenceModel memory model = _newModel(); + _assertModel(model); + + uint256 steps = bound(_actions.length, 0, MAX_ACTIONS); + for (uint256 i; i < steps; i++) { + uint8 action = uint8(_actions[i]) % 9; + _applyAction(model, action); + _assertModel(model); + } + } + + function _runDeterministic(bytes memory _actions) private { + assertLe(_actions.length, MAX_ACTIONS, "deterministic action bound"); + + SequenceModel memory model = _newModel(); + _assertModel(model); + for (uint256 i; i < _actions.length; i++) { + _applyAction(model, uint8(_actions[i]) % 9); + _assertModel(model); + } + } + + function _newModel() private view returns (SequenceModel memory model) { + model.families[0] = FamilyModel({ + original: address(facetA), + replacement: address(facetAReplacement), + currentFacet: address(facetA), + selectors: bytes.concat(FacetA.a1.selector, FacetA.a2.selector, FacetA.a3.selector), + registered: false + }); + model.families[1] = FamilyModel({ + original: address(facetB), + replacement: address(facetBReplacement), + currentFacet: address(facetB), + selectors: bytes.concat(FacetB.b1.selector, FacetB.b2.selector, FacetB.b3.selector), + registered: false + }); + model.families[2] = FamilyModel({ + original: address(facetC), + replacement: address(facetCReplacement), + currentFacet: address(facetC), + selectors: bytes.concat(FacetC.c1.selector, FacetC.c2.selector, FacetC.c3.selector), + registered: false + }); + } + + function _applyAction(SequenceModel memory _model, uint8 _action) private { + uint8 family = _action / 3; + uint8 operation = _action % 3; + FamilyModel memory familyModel = _model.families[family]; + + if (operation == 0 && !familyModel.registered) { + _addFacet(familyModel.currentFacet); + _model.families[family].registered = true; + _model.order[_model.activeCount] = family; + _model.activeCount++; + } else if (operation == 1 && familyModel.registered) { + address replacement; + if (familyModel.currentFacet == familyModel.original) { + replacement = familyModel.replacement; + } else { + replacement = familyModel.original; + } + _replaceFacet(familyModel.currentFacet, replacement); + _model.families[family].currentFacet = replacement; + } else if (operation == 2 && familyModel.registered) { + _removeFacet(familyModel.currentFacet); + _model.families[family].registered = false; + _removeFromOrder(_model, family); + } + } + + function _removeFromOrder(SequenceModel memory _model, uint8 _family) private pure { + uint256 removeIndex; + for (uint256 i; i < _model.activeCount; i++) { + if (_model.order[i] == _family) { + removeIndex = i; + break; + } + } + for (uint256 i = removeIndex; i + 1 < _model.activeCount; i++) { + _model.order[i] = _model.order[i + 1]; + } + _model.activeCount--; + } + + function _assertModel(SequenceModel memory _model) private view { + uint256 expectedFacetCount = BASELINE_FACET_COUNT + _model.activeCount; + uint256 expectedSelectorCount = BASELINE_SELECTOR_COUNT + uint256(_model.activeCount) * 3; + + _assertRoutesAndFacetSelectors(_model); + _assertCountsAndTraversal(_model, expectedFacetCount, expectedSelectorCount); + _assertInspectionViews(_model, expectedFacetCount); + } + + function _assertRoutesAndFacetSelectors(SequenceModel memory _model) private view { + _assertFacetRoutes(address(upgradeFacet), _upgradeSelectors()); + _assertFacetRoutes(address(inspectFacet), _inspectSelectors()); + + for (uint256 family; family < _model.families.length; family++) { + FamilyModel memory familyModel = _model.families[family]; + address expectedOwner = familyModel.registered ? familyModel.currentFacet : address(0); + for (uint256 selectorIndex; selectorIndex < 3; selectorIndex++) { + assertEq( + inspect.facetAddress(_selectorAt(familyModel.selectors, selectorIndex)), + expectedOwner, + "model selector owner" + ); + } + + if (familyModel.registered && familyModel.currentFacet == familyModel.original) { + _assertSelectors( + inspect.facetFunctionSelectors(familyModel.original), familyModel.selectors, "active original" + ); + assertEq( + inspect.facetFunctionSelectors(familyModel.replacement).length, 0, "inactive replacement selectors" + ); + } else if (familyModel.registered) { + assertEq(inspect.facetFunctionSelectors(familyModel.original).length, 0, "replaced old facet selectors"); + _assertSelectors( + inspect.facetFunctionSelectors(familyModel.replacement), familyModel.selectors, "active replacement" + ); + } else { + assertEq(inspect.facetFunctionSelectors(familyModel.original).length, 0, "removed original selectors"); + assertEq( + inspect.facetFunctionSelectors(familyModel.replacement).length, 0, "removed replacement selectors" + ); + } + } + } + + function _assertFacetRoutes(address _expectedFacet, bytes memory _selectors) private view { + for (uint256 i; i < _selectors.length / 4; i++) { + assertEq(inspect.facetAddress(_selectorAt(_selectors, i)), _expectedFacet, "baseline selector owner"); + } + } + + function _assertCountsAndTraversal( + SequenceModel memory _model, + uint256 _expectedFacetCount, + uint256 _expectedSelectorCount + ) private view { + (bytes4 head, bytes4 tail, uint32 facetCount, uint32 selectorCount) = + DiamondStorageUtils.facetList(address(proxy)); + assertEq(uint256(facetCount), _expectedFacetCount, "model facetCount"); + assertEq(uint256(selectorCount), _expectedSelectorCount, "model selectorCount"); + assertEq( + uint256(selectorCount) - BASELINE_SELECTOR_COUNT, + uint256(_model.activeCount) * 3, + "active family selectorCount" + ); + + bytes4[] memory visited = new bytes4[](_expectedFacetCount); + bytes4 current = head; + bytes4 previous; + for (uint256 i; i < _expectedFacetCount; i++) { + for (uint256 visitedIndex; visitedIndex < i; visitedIndex++) { + assertTrue(visited[visitedIndex] != current, "traversal node uniqueness"); + } + visited[i] = current; + current = _assertTraversalNode(_model, i, current, previous, _expectedFacetCount); + previous = visited[i]; + } + + assertEq(head, visited[0], "model head"); + assertEq(tail, visited[_expectedFacetCount - 1], "model tail"); + assertEq(current, bytes4(0), "traversal terminates after facetCount"); + } + + function _assertTraversalNode( + SequenceModel memory _model, + uint256 _index, + bytes4 _current, + bytes4 _previous, + uint256 _expectedFacetCount + ) private view returns (bytes4 actualNext) { + assertEq(_current, _expectedNodeAt(_model, _index), "model traversal node"); + + (address actualFacet, bytes4 actualPrev, bytes4 next) = DiamondStorageUtils.facetNode(address(proxy), _current); + assertEq(actualFacet, _expectedFacetAt(_model, _index), "model traversal facet"); + assertEq(actualPrev, _previous, "model previous link"); + + bytes4 expectedNext = _index + 1 < _expectedFacetCount ? _expectedNodeAt(_model, _index + 1) : bytes4(0); + assertEq(next, expectedNext, "model next link"); + actualNext = next; + } + + function _assertInspectionViews(SequenceModel memory _model, uint256 _expectedFacetCount) private view { + address[] memory expectedAddresses = new address[](_expectedFacetCount); + for (uint256 i; i < _expectedFacetCount; i++) { + expectedAddresses[i] = _expectedFacetAt(_model, i); + } + assertEq(inspect.facetAddresses(), expectedAddresses, "model facet addresses"); + + IDiamondInspect.Facet[] memory actualFacets = inspect.facets(); + assertEq(actualFacets.length, _expectedFacetCount, "model facets length"); + for (uint256 i; i < _expectedFacetCount; i++) { + assertEq(actualFacets[i].facet, expectedAddresses[i], "model facets address"); + _assertSelectors( + actualFacets[i].functionSelectors, _expectedSelectorsAt(_model, i), "model facets selectors" + ); + } + } + + function _expectedFacetAt(SequenceModel memory _model, uint256 _index) private view returns (address) { + if (_index == 0) { + return address(upgradeFacet); + } + if (_index == 1) { + return address(inspectFacet); + } + return _model.families[_model.order[_index - BASELINE_FACET_COUNT]].currentFacet; + } + + function _expectedNodeAt(SequenceModel memory _model, uint256 _index) private pure returns (bytes4) { + if (_index == 0) { + return DiamondUpgradeFacet.upgradeDiamond.selector; + } + if (_index == 1) { + return DiamondInspectFacet.facetAddress.selector; + } + + uint8 family = _model.order[_index - BASELINE_FACET_COUNT]; + if (family == 0) { + return FacetA.a1.selector; + } + if (family == 1) { + return FacetB.b1.selector; + } + return FacetC.c1.selector; + } + + function _expectedSelectorsAt(SequenceModel memory _model, uint256 _index) private pure returns (bytes memory) { + if (_index == 0) { + return bytes.concat(DiamondUpgradeFacet.upgradeDiamond.selector); + } + if (_index == 1) { + return bytes.concat( + DiamondInspectFacet.facetAddress.selector, + DiamondInspectFacet.facetFunctionSelectors.selector, + DiamondInspectFacet.facetAddresses.selector, + DiamondInspectFacet.facets.selector + ); + } + return _model.families[_model.order[_index - BASELINE_FACET_COUNT]].selectors; + } + + function _selectorAt(bytes memory _selectors, uint256 _index) private pure returns (bytes4 selector) { + assembly ("memory-safe") { + selector := mload(add(add(_selectors, 0x20), mul(_index, 4))) + } + } +} From c572dcb1d1b2be5bdcd51e1c3e7ebe9dca6bbe90 Mon Sep 17 00:00:00 2001 From: megabyte0x <79016290+megabyte0x@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:01:57 +0530 Subject: [PATCH 12/13] test(diamond): complete issue 339 coverage gates --- test/unit/diamond/DiamondUpgradeBase.t.sol | 3 +++ test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol | 3 +++ 2 files changed, 6 insertions(+) diff --git a/test/unit/diamond/DiamondUpgradeBase.t.sol b/test/unit/diamond/DiamondUpgradeBase.t.sol index bf7194a1..177c7bf2 100644 --- a/test/unit/diamond/DiamondUpgradeBase.t.sol +++ b/test/unit/diamond/DiamondUpgradeBase.t.sol @@ -125,6 +125,9 @@ abstract contract DiamondUpgrade_Base_Test is Base_Test { } } +/** + * @dev BTT spec: test/trees/Diamond.tree + */ contract DiamondStorageLayout_Unit_Test is Base_Test { function test_ShouldProveStorageBitOffsets() external { FacetA facetA = new FacetA(); diff --git a/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol b/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol index 1eb97521..f12a1d3f 100644 --- a/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol +++ b/test/unit/diamond/upgrade/shared/AddFacetsBehavior.t.sol @@ -15,6 +15,9 @@ import { ShortReturnFacet } from "test/utils/mocks/diamond/DiamondFacetMocks.sol"; +/** + * @dev BTT spec: test/trees/Diamond.tree + */ abstract contract AddFacetsBehavior is DiamondUpgrade_Base_Test { function _noBytecodeAtAddressError() internal pure virtual returns (bytes4); From c28e9b35a4c82926c3e370799312cdec71233fa1 Mon Sep 17 00:00:00 2001 From: MN Date: Tue, 4 Aug 2026 17:06:35 -0400 Subject: [PATCH 13/13] fix coverage: extract inner add selector loop. remove stack too deep error with -ir-min --- src/diamond/DiamondMod.sol | 61 ++++++++++++++++------------- src/diamond/DiamondUpgradeFacet.sol | 61 ++++++++++++++++------------- src/diamond/DiamondUpgradeMod.sol | 61 ++++++++++++++++------------- 3 files changed, 102 insertions(+), 81 deletions(-) diff --git a/src/diamond/DiamondMod.sol b/src/diamond/DiamondMod.sol index 4cb090fb..236a7676 100644 --- a/src/diamond/DiamondMod.sol +++ b/src/diamond/DiamondMod.sol @@ -183,6 +183,33 @@ function at(bytes memory selectors, uint256 index) pure returns (bytes4 selector } } +/* + * Add all selectors of a facet to the diamond, except the first. + * + * The first selector (index 0) is the facet's linked-list node identifier. + * It carries prev/next pointers and must be stored by the caller after + * resolving the linked-list position. This function stores the remaining + * selectors as leaf nodes with zero prev/next links. + * + * Reverts if any selector already belongs to a different facet. + * Returns the number of selectors processed (excluding the first). + */ +function addFacetSelectors(address _facet, bytes memory _selectors) returns (uint256 selectorsLength) { + DiamondStorage storage s = getDiamondStorage(); + /* + * Shift right by 2 is the same as dividing by 4, but cheaper. + * We do this to get the number of selectors. + */ + selectorsLength = _selectors.length >> 2; + for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { + bytes4 selector = at(_selectors, selectorIndex); + if (s.facetNodes[selector].facet != address(0)) { + revert CannotAddFunctionToDiamondThatAlreadyExists(selector); + } + s.facetNodes[selector] = FacetNode(_facet, bytes4(0), bytes4(0)); + } +} + function addFacets(address[] memory _facets) { DiamondStorage storage s = getDiamondStorage(); uint256 facetLength = _facets.length; @@ -235,23 +262,13 @@ function addFacets(address[] memory _facets) { */ s.facetNodes[prevFacetNodeId].nextFacetNodeId = currentFacetNodeId; } - /* - * Shift right by 2 is the same as dividing by 4, but cheaper. - * We do this to get the number of selectors - */ - uint256 selectorsLength = selectors.length >> 2; - unchecked { - facetList.selectorCount += uint32(selectorsLength); - } /* * Add all selectors, except the first, to the diamond. + * The first selector was already extracted as currentFacetNodeId above. + * The returned count is used to update selectorCount. */ - for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { - bytes4 selector = at(selectors, selectorIndex); - if (s.facetNodes[selector].facet != address(0)) { - revert CannotAddFunctionToDiamondThatAlreadyExists(selector); - } - s.facetNodes[selector] = FacetNode(facet, bytes4(0), bytes4(0)); + unchecked { + facetList.selectorCount += uint32(addFacetSelectors(facet, selectors)); } /* * Reset memory for the main loop. @@ -296,23 +313,13 @@ function addFacets(address[] memory _facets) { facet = nextFacet; prevFacetNodeId = currentFacetNodeId; currentFacetNodeId = nextFacetNodeId; - /* - * Shift right by 2 is the same as dividing by 4, but cheaper. - * We do this to get the number of selectors. - */ - selectorsLength = selectors.length >> 2; /* * Add all the selectors of the facet to the diamond, except the first selector. + * The first selector (currentFacetNodeId) is the pending linked-list node + * that will be stored in the next iteration once nextFacetNodeId is known. */ - for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { - bytes4 selector = at(selectors, selectorIndex); - if (s.facetNodes[selector].facet != address(0)) { - revert CannotAddFunctionToDiamondThatAlreadyExists(selector); - } - s.facetNodes[selector] = FacetNode(facet, bytes4(0), bytes4(0)); - } unchecked { - facetList.selectorCount += uint32(selectorsLength); + facetList.selectorCount += uint32(addFacetSelectors(facet, selectors)); } /* * Restore Free Memory Pointer to reuse memory from packedSelectors() calls. diff --git a/src/diamond/DiamondUpgradeFacet.sol b/src/diamond/DiamondUpgradeFacet.sol index dfde194e..94cef2d7 100644 --- a/src/diamond/DiamondUpgradeFacet.sol +++ b/src/diamond/DiamondUpgradeFacet.sol @@ -283,6 +283,33 @@ contract DiamondUpgradeFacet { } } + /* + * Add all selectors of a facet to the diamond, except the first. + * + * The first selector (index 0) is the facet's linked-list node identifier. + * It carries prev/next pointers and must be stored by the caller after + * resolving the linked-list position. This function stores the remaining + * selectors as leaf nodes with zero prev/next links. + * + * Reverts if any selector already belongs to a different facet. + * Returns the number of selectors processed (excluding the first). + */ + function addFacetSelectors(address _facet, bytes memory _selectors) internal returns (uint256 selectorsLength) { + DiamondStorage storage s = getDiamondStorage(); + /* + * Shift right by 2 is the same as dividing by 4, but cheaper. + * We do this to get the number of selectors. + */ + selectorsLength = _selectors.length >> 2; + for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { + bytes4 selector = at(_selectors, selectorIndex); + if (s.facetNodes[selector].facet != address(0)) { + revert CannotAddFunctionToDiamondThatAlreadyExists(selector); + } + s.facetNodes[selector] = FacetNode(_facet, bytes4(0), bytes4(0)); + } + } + function addFacets(address[] calldata _facets) internal { DiamondStorage storage s = getDiamondStorage(); uint256 facetLength = _facets.length; @@ -335,23 +362,13 @@ contract DiamondUpgradeFacet { */ s.facetNodes[prevFacetNodeId].nextFacetNodeId = currentFacetNodeId; } - /* - * Shift right by 2 is the same as dividing by 4, but cheaper. - * We do this to get the number of selectors - */ - uint256 selectorsLength = selectors.length >> 2; - unchecked { - facetList.selectorCount += uint32(selectorsLength); - } /* * Add all selectors, except the first, to the diamond. + * The first selector was already extracted as currentFacetNodeId above. + * The returned count is used to update selectorCount. */ - for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { - bytes4 selector = at(selectors, selectorIndex); - if (s.facetNodes[selector].facet != address(0)) { - revert CannotAddFunctionToDiamondThatAlreadyExists(selector); - } - s.facetNodes[selector] = FacetNode(facet, bytes4(0), bytes4(0)); + unchecked { + facetList.selectorCount += uint32(addFacetSelectors(facet, selectors)); } /* * Reset memory for the main loop. @@ -396,23 +413,13 @@ contract DiamondUpgradeFacet { facet = nextFacet; prevFacetNodeId = currentFacetNodeId; currentFacetNodeId = nextFacetNodeId; - /* - * Shift right by 2 is the same as dividing by 4, but cheaper. - * We do this to get the number of selectors. - */ - selectorsLength = selectors.length >> 2; /* * Add all the selectors of the facet to the diamond, except the first selector. + * The first selector (currentFacetNodeId) is the pending linked-list node + * that will be stored in the next iteration once nextFacetNodeId is known. */ - for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { - bytes4 selector = at(selectors, selectorIndex); - if (s.facetNodes[selector].facet != address(0)) { - revert CannotAddFunctionToDiamondThatAlreadyExists(selector); - } - s.facetNodes[selector] = FacetNode(facet, bytes4(0), bytes4(0)); - } unchecked { - facetList.selectorCount += uint32(selectorsLength); + facetList.selectorCount += uint32(addFacetSelectors(facet, selectors)); } /* * Restore Free Memory Pointer to reuse memory from packedSelectors() calls. diff --git a/src/diamond/DiamondUpgradeMod.sol b/src/diamond/DiamondUpgradeMod.sol index a1892d93..8920a980 100644 --- a/src/diamond/DiamondUpgradeMod.sol +++ b/src/diamond/DiamondUpgradeMod.sol @@ -256,6 +256,33 @@ function at(bytes memory selectors, uint256 index) pure returns (bytes4 selector } } +/* + * Add all selectors of a facet to the diamond, except the first. + * + * The first selector (index 0) is the facet's linked-list node identifier. + * It carries prev/next pointers and must be stored by the caller after + * resolving the linked-list position. This function stores the remaining + * selectors as leaf nodes with zero prev/next links. + * + * Reverts if any selector already belongs to a different facet. + * Returns the number of selectors processed (excluding the first). + */ +function addFacetSelectors(address _facet, bytes memory _selectors) returns (uint256 selectorsLength) { + DiamondStorage storage s = getDiamondStorage(); + /* + * Shift right by 2 is the same as dividing by 4, but cheaper. + * We do this to get the number of selectors. + */ + selectorsLength = _selectors.length >> 2; + for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { + bytes4 selector = at(_selectors, selectorIndex); + if (s.facetNodes[selector].facet != address(0)) { + revert CannotAddFunctionToDiamondThatAlreadyExists(selector); + } + s.facetNodes[selector] = FacetNode(_facet, bytes4(0), bytes4(0)); + } +} + function addFacets(address[] calldata _facets) { DiamondStorage storage s = getDiamondStorage(); uint256 facetLength = _facets.length; @@ -308,23 +335,13 @@ function addFacets(address[] calldata _facets) { */ s.facetNodes[prevFacetNodeId].nextFacetNodeId = currentFacetNodeId; } - /* - * Shift right by 2 is the same as dividing by 4, but cheaper. - * We do this to get the number of selectors - */ - uint256 selectorsLength = selectors.length >> 2; - unchecked { - facetList.selectorCount += uint32(selectorsLength); - } /* * Add all selectors, except the first, to the diamond. + * The first selector was already extracted as currentFacetNodeId above. + * The returned count is used to update selectorCount. */ - for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { - bytes4 selector = at(selectors, selectorIndex); - if (s.facetNodes[selector].facet != address(0)) { - revert CannotAddFunctionToDiamondThatAlreadyExists(selector); - } - s.facetNodes[selector] = FacetNode(facet, bytes4(0), bytes4(0)); + unchecked { + facetList.selectorCount += uint32(addFacetSelectors(facet, selectors)); } /* * Reset memory for the main loop. @@ -369,23 +386,13 @@ function addFacets(address[] calldata _facets) { facet = nextFacet; prevFacetNodeId = currentFacetNodeId; currentFacetNodeId = nextFacetNodeId; - /* - * Shift right by 2 is the same as dividing by 4, but cheaper. - * We do this to get the number of selectors. - */ - selectorsLength = selectors.length >> 2; /* * Add all the selectors of the facet to the diamond, except the first selector. + * The first selector (currentFacetNodeId) is the pending linked-list node + * that will be stored in the next iteration once nextFacetNodeId is known. */ - for (uint256 selectorIndex = 1; selectorIndex < selectorsLength; selectorIndex++) { - bytes4 selector = at(selectors, selectorIndex); - if (s.facetNodes[selector].facet != address(0)) { - revert CannotAddFunctionToDiamondThatAlreadyExists(selector); - } - s.facetNodes[selector] = FacetNode(facet, bytes4(0), bytes4(0)); - } unchecked { - facetList.selectorCount += uint32(selectorsLength); + facetList.selectorCount += uint32(addFacetSelectors(facet, selectors)); } /* * Restore Free Memory Pointer to reuse memory from packedSelectors() calls.