Skip to content

Commit f652111

Browse files
mudgenmudgen
committed
Made memory improvements
1 parent d87f6bf commit f652111

File tree

1 file changed

+130
-32
lines changed

1 file changed

+130
-32
lines changed

‎src/diamond/DiamondUpgradeFacet.sol‎

Lines changed: 130 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -184,20 +184,82 @@ contract DiamondUpgradeFacet {
184184
*/
185185
error CannotReplaceFunctionFromNonReplacementFacet(bytes4 _selector);
186186

187-
function packedSelectors(address _facet) internal view returns (bytes memory) {
188-
if (_facet.code.length == 0) {
189-
revert NoBytecodeAtAddress(_facet);
190-
}
191-
(bool success, bytes memory selectors) =
192-
_facet.staticcall(abi.encodeWithSelector(bytes4(keccak256("packedSelectors()"))));
187+
function packedSelectors(address _facet) internal view returns (bytes memory selectors) {
188+
assembly ("memory-safe") {
189+
/**
190+
* 1. Point to the current Free Memory Pointer (0x40).
191+
* We use this as the start of our "Static Buffer" workspace.
192+
*/
193+
let ptr := mload(0x40)
194+
/**
195+
* 2. Prepare calldata for packedSelectors()
196+
* "0x3e62267c" is the function selector for packedSelectors()
197+
* "0x3e62267c" is bytes4(keccak256("packedSelectors()"))
198+
* We store it at ptr to reuse that memory immediately.
199+
*/
200+
mstore(ptr, 0x3e62267c00000000000000000000000000000000000000000000000000000000)
201+
mstore(add(ptr, 0x04), _facet)
202+
/**
203+
* 3. Perform the staticcall.
204+
* We pass 0 for out and outSize to keep the return data in the
205+
* "Free Waiting Room" (Return Data Buffer) until we verify it.
206+
*/
207+
let success :=
208+
staticcall(
209+
gas(), // pass all available gas
210+
_facet, // target address
211+
ptr, // pointer to start of input
212+
0x24, // input length (4 bytes for selector and 32 bytes for address)
213+
0, // output pointer, not used
214+
0 // output length, not used
215+
)
216+
/**
217+
* 4. Basic Safety Check.
218+
* Ensure the call succeeded and returned at least 68 bytes
219+
*/
220+
if iszero(success) {
221+
/**
222+
* error FunctionSelectorsCallFailed(address) selector: 0x30319baa
223+
*/
224+
mstore(0x00, 0x30319baa00000000000000000000000000000000000000000000000000000000)
225+
mstore(0x04, _facet)
226+
revert(0x00, 0x24)
227+
}
228+
/**
229+
* Minimum return data size is 68 bytes:
230+
* 32 bytes offset + 32 bytes array length + 4 bytes (at least one selector).
231+
* If facet address has no bytecode then return size will be 0.
232+
*/
233+
let size := returndatasize()
234+
if lt(size, 68) {
235+
/**
236+
* error NoSelectorsForFacet(address) selector: 0x9c23886b
237+
*/
238+
mstore(0x00, 0x9c23886b00000000000000000000000000000000000000000000000000000000)
239+
mstore(0x04, _facet)
240+
revert(0x00, 0x24)
241+
}
242+
/**
243+
* 5. Extraction.
244+
* Move the data from the hidden Return Data Buffer.
245+
* This overwrites the 4-byte selector and facet address we stored earlier.
246+
*
247+
* ABI encoding for 'bytes' includes a 32-byte offset word at the start.
248+
* We point 'selectors' to ptr + 0x20 to skip the offset and point
249+
* directly to the Length word, making it a valid Solidity bytes array.
250+
*/
251+
size := sub(size, 0x20) // Adjust size to account for the 32-byte offset word
252+
returndatacopy(ptr, 0x20, size)
193253

194-
if (success == false) {
195-
revert FunctionSelectorsCallFailed(_facet);
196-
}
197-
if (selectors.length < 4) {
198-
revert NoSelectorsForFacet(_facet);
254+
selectors := ptr
255+
256+
/**
257+
* 6. Update Free Memory Pointer
258+
* New Free Memory Pointer is after the copied selectors.
259+
* Notice that this is not 32-byte aligned.
260+
*/
261+
mstore(0x40, add(ptr, size))
199262
}
200-
return selectors;
201263
}
202264

203265
function at(bytes memory selectors, uint256 index) internal pure returns (bytes4 selector) {
@@ -223,6 +285,14 @@ contract DiamondUpgradeFacet {
223285
return;
224286
}
225287
FacetList memory facetList = s.facetList;
288+
/*
289+
* Store current Free Memory Pointer to restore later.
290+
* This is use to reuse memory in each loop iteration.
291+
*/
292+
uint256 freeMemoryPointerLocation;
293+
assembly ("memory-safe") {
294+
freeMemoryPointerLocation := mload(0x40)
295+
}
226296
bytes4 prevFacetNodeId = facetList.lastFacetNodeId;
227297
bytes memory selectors = packedSelectors(_facets[0]);
228298
bytes4 currentFacetNodeId = at(selectors, 0);
@@ -245,6 +315,12 @@ contract DiamondUpgradeFacet {
245315
nextSelectors = packedSelectors(_facets[nextI]);
246316
nextFacetNodeId = at(nextSelectors, 0);
247317
}
318+
/*
319+
* Restore Free Memory Pointer to reuse memory from packedSelectors() calls.
320+
*/
321+
assembly ("memory-safe") {
322+
mstore(0x40, freeMemoryPointerLocation)
323+
}
248324
address oldFacet = s.facetNodes[currentFacetNodeId].facet;
249325
if (oldFacet != address(0)) {
250326
revert CannotAddFunctionToDiamondThatAlreadyExists(currentFacetNodeId);
@@ -284,6 +360,14 @@ contract DiamondUpgradeFacet {
284360
function replaceFacets(FacetReplacement[] calldata _replaceFacets) internal {
285361
DiamondStorage storage s = getDiamondStorage();
286362
FacetList memory facetList = s.facetList;
363+
/*
364+
* Store current Free Memory Pointer to restore later.
365+
* This is use to reuse memory in each loop iteration.
366+
*/
367+
uint256 freeMemoryPointerLocation;
368+
assembly ("memory-safe") {
369+
freeMemoryPointerLocation := mload(0x40)
370+
}
287371
for (uint256 i; i < _replaceFacets.length; i++) {
288372
address oldFacet = _replaceFacets[i].oldFacet;
289373
address newFacet = _replaceFacets[i].newFacet;
@@ -298,15 +382,15 @@ contract DiamondUpgradeFacet {
298382
/**
299383
* Validate old facet exists.
300384
*/
301-
bytes4 prevFacetNodeId;
302-
bytes4 nextFacetNodeId;
303-
{
304-
FacetNode storage oldFacetNode = s.facetNodes[oldCurrentFacetNodeId];
305-
if (oldFacetNode.facet != oldFacet) {
306-
revert FacetToReplaceDoesNotExist(oldFacet);
307-
}
308-
prevFacetNodeId = oldFacetNode.prevFacetNodeId;
309-
nextFacetNodeId = oldFacetNode.nextFacetNodeId;
385+
FacetNode memory oldFacetNode = s.facetNodes[oldCurrentFacetNodeId];
386+
if (oldFacetNode.facet != oldFacet) {
387+
revert FacetToReplaceDoesNotExist(oldFacet);
388+
}
389+
/*
390+
* Restore Free Memory Pointer to reuse memory.
391+
*/
392+
assembly ("memory-safe") {
393+
mstore(0x40, freeMemoryPointerLocation)
310394
}
311395

312396
if (oldCurrentFacetNodeId != newCurrentFacetNodeId) {
@@ -330,20 +414,21 @@ contract DiamondUpgradeFacet {
330414
if (oldCurrentFacetNodeId == facetList.firstFacetNodeId) {
331415
facetList.firstFacetNodeId = newCurrentFacetNodeId;
332416
} else {
333-
s.facetNodes[prevFacetNodeId].nextFacetNodeId = newCurrentFacetNodeId;
417+
s.facetNodes[oldFacetNode.prevFacetNodeId].nextFacetNodeId = newCurrentFacetNodeId;
334418
}
335419
if (oldCurrentFacetNodeId == facetList.lastFacetNodeId) {
336420
facetList.lastFacetNodeId = newCurrentFacetNodeId;
337421
} else {
338-
s.facetNodes[nextFacetNodeId].prevFacetNodeId = newCurrentFacetNodeId;
422+
s.facetNodes[oldFacetNode.nextFacetNodeId].prevFacetNodeId = newCurrentFacetNodeId;
339423
}
340424
} else {
341425
if (keccak256(oldSelectors) == keccak256(newSelectors)) {
342426
/**
343427
* Same selectors, replace.
344428
*/
345429
emit DiamondFunctionReplaced(newCurrentFacetNodeId, oldFacet, newFacet);
346-
s.facetNodes[newCurrentFacetNodeId] = FacetNode(newFacet, prevFacetNodeId, nextFacetNodeId);
430+
s.facetNodes[newCurrentFacetNodeId] =
431+
FacetNode(newFacet, oldFacetNode.prevFacetNodeId, oldFacetNode.nextFacetNodeId);
347432
/**
348433
* Replace remaining selectors.
349434
*/
@@ -362,7 +447,8 @@ contract DiamondUpgradeFacet {
362447
*/
363448
emit DiamondFunctionReplaced(newCurrentFacetNodeId, oldFacet, newFacet);
364449
}
365-
s.facetNodes[newCurrentFacetNodeId] = FacetNode(newFacet, prevFacetNodeId, nextFacetNodeId);
450+
s.facetNodes[newCurrentFacetNodeId] =
451+
FacetNode(newFacet, oldFacetNode.prevFacetNodeId, oldFacetNode.nextFacetNodeId);
366452

367453
/**
368454
* Add or replace new selectors.
@@ -409,6 +495,14 @@ contract DiamondUpgradeFacet {
409495
function removeFacets(address[] calldata _facets) internal {
410496
DiamondStorage storage s = getDiamondStorage();
411497
FacetList memory facetList = s.facetList;
498+
/*
499+
* Store current Free Memory Pointer to restore later.
500+
* This is use to reuse memory in each loop iteration.
501+
*/
502+
uint256 freeMemoryPointerLocation;
503+
assembly ("memory-safe") {
504+
freeMemoryPointerLocation := mload(0x40)
505+
}
412506
for (uint256 i = 0; i < _facets.length; i++) {
413507
address facet = _facets[i];
414508
bytes memory selectors = packedSelectors(facet);
@@ -418,24 +512,28 @@ contract DiamondUpgradeFacet {
418512
facetList.selectorCount -= uint32(selectorsLength);
419513
}
420514
bytes4 currentFacetNodeId = at(selectors, 0);
421-
FacetNode storage facetNode = s.facetNodes[currentFacetNodeId];
515+
FacetNode memory facetNode = s.facetNodes[currentFacetNodeId];
422516
if (facetNode.facet != facet) {
423517
revert CannotRemoveFacetThatDoesNotExist(facet);
424518
}
519+
/*
520+
* Restore Free Memory Pointer to reuse memory.
521+
*/
522+
assembly ("memory-safe") {
523+
mstore(0x40, freeMemoryPointerLocation)
524+
}
425525
/**
426526
* Remove the facet from the linked list.
427527
*/
428-
bytes4 nextFacetNodeId = facetNode.nextFacetNodeId;
429-
bytes4 prevFacetNodeId = facetNode.prevFacetNodeId;
430528
if (currentFacetNodeId == facetList.firstFacetNodeId) {
431-
facetList.firstFacetNodeId = nextFacetNodeId;
529+
facetList.firstFacetNodeId = facetNode.nextFacetNodeId;
432530
} else {
433-
s.facetNodes[prevFacetNodeId].nextFacetNodeId = nextFacetNodeId;
531+
s.facetNodes[facetNode.prevFacetNodeId].nextFacetNodeId = facetNode.nextFacetNodeId;
434532
}
435533
if (currentFacetNodeId == facetList.lastFacetNodeId) {
436-
facetList.lastFacetNodeId = prevFacetNodeId;
534+
facetList.lastFacetNodeId = facetNode.prevFacetNodeId;
437535
} else {
438-
s.facetNodes[nextFacetNodeId].prevFacetNodeId = prevFacetNodeId;
536+
s.facetNodes[facetNode.nextFacetNodeId].prevFacetNodeId = facetNode.prevFacetNodeId;
439537
}
440538
/**
441539
* Remove facet selectors.

0 commit comments

Comments
 (0)