Remove storage use for BeaconProxy deployment

contracts/proxy/ERC1967/ERC1967Utils.sol

@@ -145,16 +145,25 @@ library ERC1967Utils {
      * @dev Stores a new beacon in the ERC-1967 beacon slot.
      */
     function _setBeacon(address newBeacon) private {
+        _validateBeacon(newBeacon);
+
+        StorageSlot.getAddressSlot(BEACON_SLOT).value = newBeacon;
+
+        _validateBeaconImplementation(newBeacon);
+    }
+
+    function _validateBeacon(address newBeacon) private view {
         if (newBeacon.code.length == 0) {
             revert ERC1967InvalidBeacon(newBeacon);
         }
+    }
 
-        StorageSlot.getAddressSlot(BEACON_SLOT).value = newBeacon;
-
+    function _validateBeaconImplementation(address newBeacon) private view returns(address) {
         address beaconImplementation = IBeacon(newBeacon).implementation();
         if (beaconImplementation.code.length == 0) {
             revert ERC1967InvalidImplementation(beaconImplementation);
         }
+        return beaconImplementation;
     }
 
     /**
@@ -179,6 +188,18 @@ library ERC1967Utils {
         }
     }
 
+    function validateBeaconAndCall(address newBeacon, bytes memory data) internal {
+        _validateBeacon(newBeacon);
+        address implementation = _validateBeaconImplementation(newBeacon);
+        emit BeaconUpgraded(newBeacon);
+
+        if (data.length > 0) {
+            Address.functionDelegateCall(implementation, data);
+        } else {
+            _checkNonPayable();
+        }
+    }
+
     /**
      * @dev Reverts if `msg.value` is not zero. It can be used to avoid `msg.value` stuck in the contract
      * if an upgrade doesn't perform an initialization call.

contracts/proxy/beacon/BeaconProxy.sol

@@ -11,14 +11,10 @@ import {ERC1967Utils} from "../ERC1967/ERC1967Utils.sol";
  * @dev This contract implements a proxy that gets the implementation address for each call from an {UpgradeableBeacon}.
  *
  * The beacon address can only be set once during construction, and cannot be changed afterwards. It is stored in an
- * immutable variable to avoid unnecessary storage reads, and also in the beacon storage slot specified by
- * https://eips.ethereum.org/EIPS/eip-1967[ERC-1967] so that it can be accessed externally.
+ * immutable variable to avoid unnecessary storage reads.
  *
  * CAUTION: Since the beacon address can never be changed, you must ensure that you either control the beacon, or trust
  * the beacon to not upgrade the implementation maliciously.
- *
- * IMPORTANT: Do not use the implementation logic to modify the beacon storage slot. Doing so would leave the proxy in
- * an inconsistent state where the beacon storage slot does not match the beacon address.
  */
 contract BeaconProxy is Proxy {
     // An immutable address for the beacon to avoid unnecessary SLOADs before each delegate call.
@@ -37,7 +33,7 @@ contract BeaconProxy is Proxy {
      * - If `data` is empty, `msg.value` must be zero.
      */
     constructor(address beacon, bytes memory data) payable {
-        ERC1967Utils.upgradeBeaconToAndCall(beacon, data);
+        ERC1967Utils.validateBeaconAndCall(beacon, data);
         _beacon = beacon;
     }
 

test/proxy/beacon/BeaconProxy.test.js

@@ -49,7 +49,7 @@ describe('BeaconProxy', function () {
   describe('initialization', function () {
     async function assertInitialized({ value, balance }) {
       const beaconAddress = await getAddressInSlot(this.proxy, BeaconSlot);
-      expect(beaconAddress).to.equal(this.beacon);
+      expect(beaconAddress).to.equal(ethers.ZeroAddress);
 
       const dummy = this.v1.attach(this.proxy);
       expect(await dummy.value()).to.equal(value);