Skip to content

Commit 2bcca59

Browse files
mudgenmudgen
committed
Added implementations of facet-based diamonds
1 parent 93b1df1 commit 2bcca59

File tree

3 files changed

+351
-160
lines changed

3 files changed

+351
-160
lines changed

‎src/diamond/DiamondInspectFacet.sol‎

Lines changed: 78 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -13,24 +13,28 @@ contract DiamondInspectFacet {
1313
bytes32 constant DIAMOND_STORAGE_POSITION = keccak256("erc8109.diamond");
1414

1515
/**
16-
* @notice Data stored for each function selector.
17-
* @dev Facet address of function selector.
18-
* Position of selector in the 'bytes4[] selectors' array.
16+
* @notice Data stored for each function selector
17+
* @dev Facet address of function selector
18+
* Position of selector in the 'bytes4[] selectors' array
1919
*/
20-
struct FacetAndPosition {
20+
struct FacetNode {
2121
address facet;
22-
uint32 position;
22+
bytes4 prevFacetSelector;
23+
bytes4 nextFacetSelector;
24+
}
25+
26+
struct FacetList {
27+
uint32 facetCount;
28+
bytes4 firstFacetSelector;
29+
bytes4 lastFacetSelector;
2330
}
2431

2532
/**
2633
* @custom:storage-location erc8042:erc8109.diamond
2734
*/
2835
struct DiamondStorage {
29-
mapping(bytes4 functionSelector => FacetAndPosition) facetAndPosition;
30-
/**
31-
* Array of all function selectors that can be called in the diamond.
32-
*/
33-
bytes4[] selectors;
36+
mapping(bytes4 functionSelector => FacetNode) facetNodes;
37+
FacetList facetList;
3438
}
3539

3640
function getStorage() internal pure returns (DiamondStorage storage s) {
@@ -48,7 +52,7 @@ contract DiamondInspectFacet {
4852
*/
4953
function facetAddress(bytes4 _functionSelector) external view returns (address facet) {
5054
DiamondStorage storage s = getStorage();
51-
facet = s.facetAndPosition[_functionSelector].facet;
55+
facet = s.facetNodes[_functionSelector].facet;
5256
}
5357

5458
/**
@@ -60,7 +64,7 @@ contract DiamondInspectFacet {
6064
function facetFunctionSelectors(address _facet) external view returns (bytes4[] memory facetSelectors) {
6165
DiamondStorage storage s = getStorage();
6266
facetSelectors = IFacet(_facet).functionSelectors();
63-
if (facetSelectors.length == 0 || s.facetAndPosition[facetSelectors[0]].facet == address(0)) {
67+
if (facetSelectors.length == 0 || s.facetNodes[facetSelectors[0]].facet == address(0)) {
6468
facetSelectors = new bytes4[](0);
6569
}
6670
}
@@ -72,13 +76,14 @@ contract DiamondInspectFacet {
7276
*/
7377
function facetAddresses() external view returns (address[] memory allFacets) {
7478
DiamondStorage storage s = getStorage();
75-
bytes4[] memory selectors = s.selectors;
76-
uint256 facetCount = selectors.length;
79+
FacetList memory facetList = s.facetList;
80+
uint256 facetCount = facetList.facetCount;
7781
allFacets = new address[](facetCount);
78-
for (uint256 selectorIndex; selectorIndex < facetCount; selectorIndex++) {
79-
bytes4 selector = selectors[selectorIndex];
80-
address facet = s.facetAndPosition[selector].facet;
81-
allFacets[selectorIndex] = facet;
82+
bytes4 currentSelector = facetList.firstFacetSelector;
83+
for (uint256 i; i < facetCount; i++) {
84+
address facet = s.facetNodes[currentSelector].facet;
85+
allFacets[i] = facet;
86+
currentSelector = s.facetNodes[currentSelector].nextFacetSelector;
8287
}
8388
}
8489

@@ -95,15 +100,63 @@ contract DiamondInspectFacet {
95100
*/
96101
function facets() external view returns (Facet[] memory facetsAndSelectors) {
97102
DiamondStorage storage s = getStorage();
98-
bytes4[] memory selectors = s.selectors;
99-
uint256 facetCount = selectors.length;
103+
FacetList memory facetList = s.facetList;
104+
uint256 facetCount = facetList.facetCount;
105+
bytes4 currentSelector = facetList.firstFacetSelector;
100106
facetsAndSelectors = new Facet[](facetCount);
101-
for (uint256 selectorIndex; selectorIndex < facetCount; selectorIndex++) {
102-
bytes4 selector = selectors[selectorIndex];
103-
address facet = s.facetAndPosition[selector].facet;
107+
for (uint256 i; i < facetCount; i++) {
108+
address facet = s.facetNodes[currentSelector].facet;
104109
bytes4[] memory facetSelectors = IFacet(facet).functionSelectors();
105-
facetsAndSelectors[selectorIndex].facet = facet;
106-
facetsAndSelectors[selectorIndex].functionSelectors = facetSelectors;
110+
facetsAndSelectors[i].facet = facet;
111+
facetsAndSelectors[i].functionSelectors = facetSelectors;
112+
currentSelector = s.facetNodes[currentSelector].nextFacetSelector;
113+
}
114+
}
115+
116+
struct FunctionFacetPair {
117+
bytes4 selector;
118+
address facet;
119+
}
120+
121+
/**
122+
* @notice Returns an array of all function selectors and their
123+
* corresponding facet addresses.
124+
*
125+
* @dev Iterates through the diamond's stored selectors and pairs
126+
* each with its facet.
127+
* @return pairs An array of `FunctionFacetPair` structs, each containing
128+
* a selector and its facet address.
129+
*/
130+
function functionFacetPairs() external view returns (FunctionFacetPair[] memory pairs) {
131+
DiamondStorage storage s = getStorage();
132+
FacetList memory facetList = s.facetList;
133+
uint256 facetCount = facetList.facetCount;
134+
bytes4 currentSelector = facetList.firstFacetSelector;
135+
uint256 selectorCount;
136+
Facet[] memory facetsAndSelectors = new Facet[](facetCount);
137+
for (uint256 i; i < facetCount; i++) {
138+
address facet = s.facetNodes[currentSelector].facet;
139+
bytes4[] memory facetSelectors = IFacet(facet).functionSelectors();
140+
unchecked {
141+
selectorCount += facetSelectors.length;
142+
}
143+
facetsAndSelectors[i].facet = facet;
144+
facetsAndSelectors[i].functionSelectors = facetSelectors;
145+
currentSelector = s.facetNodes[currentSelector].nextFacetSelector;
146+
}
147+
pairs = new FunctionFacetPair[](selectorCount);
148+
selectorCount = 0;
149+
for (uint256 i; i < facetCount; i++) {
150+
Facet memory facetsAndSelector = facetsAndSelectors[i];
151+
address facet = facetsAndSelector.facet;
152+
uint256 selectorsLength = facetsAndSelector.functionSelectors.length;
153+
for (uint256 selectorIndex; selectorIndex < selectorsLength; selectorIndex++) {
154+
bytes4 selector = facetsAndSelector.functionSelectors[selectorIndex];
155+
pairs[selectorCount] = FunctionFacetPair(selector, facet);
156+
unchecked {
157+
selectorCount++;
158+
}
159+
}
107160
}
108161
}
109162

‎src/diamond/DiamondMod.sol‎

Lines changed: 65 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -14,25 +14,29 @@ pragma solidity >=0.8.30;
1414

1515
bytes32 constant DIAMOND_STORAGE_POSITION = keccak256("erc8109.diamond");
1616

17-
/*
18-
* @notice Data stored for each function selector.
19-
* @dev Facet address of function selector.
20-
* Position of selector in the `bytes4[] selectors` array.
21-
*/
22-
struct FacetAndPosition {
17+
/**
18+
* @notice Data stored for each function selector
19+
* @dev Facet address of function selector
20+
* Position of selector in the 'bytes4[] selectors' array
21+
*/
22+
struct FacetNode {
2323
address facet;
24-
uint32 position;
24+
bytes4 prevFacetSelector;
25+
bytes4 nextFacetSelector;
26+
}
27+
28+
struct FacetList {
29+
uint32 facetCount;
30+
bytes4 firstFacetSelector;
31+
bytes4 lastFacetSelector;
2532
}
2633

2734
/**
2835
* @custom:storage-location erc8042:erc8109.diamond
2936
*/
3037
struct DiamondStorage {
31-
mapping(bytes4 functionSelector => FacetAndPosition) facetAndPosition;
32-
/**
33-
* `selectors` contains all function selectors that can be called in the diamond.
34-
*/
35-
bytes4[] selectors;
38+
mapping(bytes4 functionSelector => FacetNode) facetNodes;
39+
FacetList facetList;
3640
}
3741

3842
function getStorage() pure returns (DiamondStorage storage s) {
@@ -42,6 +46,27 @@ function getStorage() pure returns (DiamondStorage storage s) {
4246
}
4347
}
4448

49+
error FunctionSelectorsCallFailed(address _facet);
50+
error NoSelectorsForFacet(address _facet);
51+
error NoBytecodeAtAddress(address _contractAddress);
52+
53+
bytes constant FUNCTION_SELECTORS_CALL = abi.encodeWithSignature("functionSelectors()");
54+
55+
function functionSelectors(address _facet) view returns (bytes4[] memory) {
56+
if (_facet.code.length == 0) {
57+
revert NoBytecodeAtAddress(_facet);
58+
}
59+
(bool success, bytes memory data) = _facet.staticcall(FUNCTION_SELECTORS_CALL);
60+
if (success == false) {
61+
revert FunctionSelectorsCallFailed(_facet);
62+
}
63+
bytes4[] memory selectors = abi.decode(data, (bytes4[]));
64+
if (selectors.length == 0) {
65+
revert NoSelectorsForFacet(_facet);
66+
}
67+
return selectors;
68+
}
69+
4570
/**
4671
* @notice Emitted when a function is added to a diamond.
4772
*
@@ -50,43 +75,46 @@ function getStorage() pure returns (DiamondStorage storage s) {
5075
*/
5176
event DiamondFunctionAdded(bytes4 indexed _selector, address indexed _facet);
5277

53-
struct FacetFunctions {
54-
address facet;
55-
bytes4[] selectors;
56-
}
57-
58-
interface IFacet {
59-
function functionSelectors() external view returns (bytes4[] memory);
60-
}
61-
6278
error CannotAddFunctionToDiamondThatAlreadyExists(bytes4 _selector);
63-
error NoSelectorsForFacet(address _facet);
79+
error NoFacetsToAdd();
6480

6581
/**
6682
* @notice Adds facets and their function selectors to the diamond.
67-
* @dev Only supports adding functions during diamond deployment.
6883
*/
69-
function addFacets(address[] memory _facets) {
84+
function addFacets(address[] calldata _facets) {
7085
DiamondStorage storage s = getStorage();
71-
uint32 selectorPosition = uint32(s.selectors.length);
72-
for (uint256 i; i < _facets.length; i++) {
73-
address facet = _facets[i];
74-
bytes4[] memory facetSelectors = IFacet(facet).functionSelectors();
75-
if (facetSelectors.length == 0) {
76-
revert NoSelectorsForFacet(facet);
86+
uint256 facetsLength = _facets.length;
87+
if (facetsLength == 0) {
88+
return;
89+
}
90+
FacetList memory facetList = s.facetList;
91+
bytes4 prevSelector = facetList.lastFacetSelector;
92+
bytes4 currentSelector;
93+
for (uint256 i; i < facetsLength; i++) {
94+
address facet = _facets[0];
95+
bytes4[] memory currentSelectors = functionSelectors(facet);
96+
currentSelector = currentSelectors[0];
97+
if (i == 0 && facetList.facetCount == 0) {
98+
facetList.firstFacetSelector = currentSelector;
99+
} else {
100+
s.facetNodes[prevSelector].nextFacetSelector = currentSelector;
77101
}
78-
s.selectors.push(facetSelectors[0]);
79-
for (uint256 selectorIndex; selectorIndex < facetSelectors.length; selectorIndex++) {
80-
bytes4 selector = facetSelectors[selectorIndex];
81-
address oldFacet = s.facetAndPosition[selector].facet;
102+
for (uint256 selectorIndex; selectorIndex < currentSelectors.length; selectorIndex++) {
103+
bytes4 selector = currentSelectors[selectorIndex];
104+
address oldFacet = s.facetNodes[selector].facet;
82105
if (oldFacet != address(0)) {
83106
revert CannotAddFunctionToDiamondThatAlreadyExists(selector);
84107
}
85-
s.facetAndPosition[selector] = FacetAndPosition(facet, selectorPosition);
108+
s.facetNodes[selector] = FacetNode(facet, prevSelector, bytes4(0));
86109
emit DiamondFunctionAdded(selector, facet);
87110
}
88-
selectorPosition++;
111+
prevSelector = currentSelector;
112+
}
113+
unchecked {
114+
facetList.facetCount += uint32(facetsLength);
89115
}
116+
facetList.lastFacetSelector = currentSelector;
117+
s.facetList = facetList;
90118
}
91119

92120
error FunctionNotFound(bytes4 _selector);
@@ -100,7 +128,7 @@ function diamondFallback() {
100128
/**
101129
* get facet from function selector
102130
*/
103-
address facet = s.facetAndPosition[msg.sig].facet;
131+
address facet = s.facetNodes[msg.sig].facet;
104132
if (facet == address(0)) {
105133
revert FunctionNotFound(msg.sig);
106134
}

0 commit comments

Comments
 (0)