✨ Add opt-in storage layout compatibility guard for UUPS

src/utils/UUPSUpgradeable.sol

@@ -20,6 +20,8 @@ abstract contract UUPSUpgradeable is CallContextChecker {
 
     /// @dev The upgrade failed.
     error UpgradeFailed();
+    /// @dev The storagelayout mismatch.
+    error StorageLayoutMismatch();
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
     /*                           EVENTS                           */
@@ -52,6 +54,10 @@ abstract contract UUPSUpgradeable is CallContextChecker {
     /// ```
     function _authorizeUpgrade(address newImplementation) internal virtual;
 
+    /// @dev Optional hook to guard against storage layout incompatibility.
+    /// Override in implementation if desired (e.g. compare a constant layout hash).
+    function _checkStorageLayout(address newImplementation) internal virtual {}
+
     /// @dev Returns the storage slot used by the implementation,
     /// as specified in [ERC1822](https://eips.ethereum.org/EIPS/eip-1822).
     ///
@@ -73,6 +79,7 @@ abstract contract UUPSUpgradeable is CallContextChecker {
         onlyProxy
     {
         _authorizeUpgrade(newImplementation);
+        _checkStorageLayout(newImplementation);
         /// @solidity memory-safe-assembly
         assembly {
             newImplementation := shr(96, shl(96, newImplementation)) // Clears upper 96 bits.

test/UUPSUpgradeable.t.sol

@@ -5,6 +5,7 @@ import "./utils/SoladyTest.sol";
 import {CallContextChecker, UUPSUpgradeable} from "../src/utils/UUPSUpgradeable.sol";
 import {LibClone} from "../src/utils/LibClone.sol";
 import {MockUUPSImplementation} from "../test/utils/mocks/MockUUPSImplementation.sol";
+import "./utils/mocks/MockUUPSWithDifferentLayout.sol";
 
 contract UUPSUpgradeableTest is SoladyTest {
     MockUUPSImplementation impl1;
@@ -70,4 +71,18 @@ contract UUPSUpgradeableTest is SoladyTest {
         vm.expectRevert(MockUUPSImplementation.Unauthorized.selector);
         MockUUPSImplementation(proxy).upgradeToAndCall(address(0xABCD), "");
     }
+
+    function testUpgradeRevertsOnStorageLayoutMismatch() public {
+        MockUUPSWithDifferentLayout bad = new MockUUPSWithDifferentLayout();
+
+        vm.prank(address(this));
+        vm.expectRevert(UUPSUpgradeable.StorageLayoutMismatch.selector);
+        MockUUPSImplementation(proxy).upgradeToAndCall(address(bad), "");
+    }
+
+    function testUpgradeAllowsCompatibleImplementation() public {
+        MockUUPSImplementation compatible = new MockUUPSImplementation();
+        MockUUPSImplementation(proxy).upgradeToAndCall(address(compatible), "");
+        // passes silently
+    }
 }

test/utils/mocks/MockUUPSImplementation.sol

@@ -11,6 +11,11 @@ contract MockUUPSImplementation is UUPSUpgradeable, Brutalizer {
 
     address public owner;
 
+    /// @dev Storage layout version hash for upgrade compatibility checks.
+    bytes32 private constant _STORAGE_LAYOUT_VERSION =
+        0x7b86195f14c03aa21c7b6881091a51da39658e244432de468f62f9387ca79dba;
+    // keccak256("solady.mock.uups.v1")
+
     error Unauthorized();
 
     error CustomError(address owner_);
@@ -26,6 +31,32 @@ contract MockUUPSImplementation is UUPSUpgradeable, Brutalizer {
 
     function _authorizeUpgrade(address) internal override onlyOwner {}
 
+    /// @dev Returns the storage layout version for compatibility verification.
+    function STORAGE_LAYOUT_VERSION() external pure returns (bytes32) {
+        return _STORAGE_LAYOUT_VERSION;
+    }
+
+    /// @dev Checks storage layout compatibility by comparing version hashes.
+    function _checkStorageLayout(address newImpl) internal view override {
+        bytes32 expected = _STORAGE_LAYOUT_VERSION;
+
+        /// @solidity memory-safe-assembly
+        assembly {
+            // Skip check if `newImpl` has no code to preserve standard UUPS errors.
+            if extcodesize(newImpl) {
+                mstore(0x00, 0xa19f05fc) // `STORAGE_LAYOUT_VERSION()`
+                if iszero(staticcall(gas(), newImpl, 0x1c, 0x04, 0x00, 0x20)) {
+                    mstore(0x00, 0xc2fbd48f) // `StorageLayoutMismatch()`
+                    revert(0x1c, 0x04)
+                }
+                if iszero(eq(mload(0x00), expected)) {
+                    mstore(0x00, 0xc2fbd48f) // `StorageLayoutMismatch()`
+                    revert(0x1c, 0x04)
+                }
+            }
+        }
+    }
+
     function revertWithError() public view {
         revert CustomError(owner);
     }

test/utils/mocks/MockUUPSWithDifferentLayout.sol

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import {UUPSUpgradeable} from "../../../src/utils/UUPSUpgradeable.sol";
+
+contract MockUUPSWithDifferentLayout is UUPSUpgradeable {
+    bool public extraVar; // Breaks layout compatibility
+    uint256 public value;
+    address public owner;
+
+    /// @dev Intentionally different version to trigger mismatch.
+    bytes32 private constant _STORAGE_LAYOUT_VERSION =
+        0xbe1c9475ae85c040745b734fec444fcbc3d6e069d8e0b5027394bd8b0a614f94;
+    // keccak256("solady.mock.uups.incompatible.v1")
+
+    function _authorizeUpgrade(address) internal view override {
+        require(msg.sender == address(this), "unauthorized");
+    }
+
+    /// @dev Returns incompatible version hash to test mismatch detection.
+    function STORAGE_LAYOUT_VERSION() external pure returns (bytes32) {
+        return _STORAGE_LAYOUT_VERSION;
+    }
+
+    /// @dev No-op override to skip layout check (for negative testing).
+    function _checkStorageLayout(address) internal pure override {}
+}