Registry of Balancer standard contracts (v2)

pkg/vault/contracts/BalancerContractRegistry.sol

@@ -59,6 +59,9 @@ contract BalancerContractRegistry is IBalancerContractRegistry, SingletonAuthent
     // contract state by overwriting a registry entry with an "alias" that matches a different contract).
     mapping(bytes32 contractAliasId => address addr) private _contractAliases;
 
+    /// @dev A `_contractRegistry` entry has no corresponding `_contractInfo`. Should never happen.
+    error InconsistentState();
+
     constructor(IVault vault) SingletonAuthentication(vault) {
         // solhint-disable-previous-line no-empty-blocks
     }
@@ -127,18 +130,7 @@ contract BalancerContractRegistry is IBalancerContractRegistry, SingletonAuthent
         }
 
         // Ensure name isn't already in use (including as an alias).
-        bytes32 contractId = _getContractId(contractName);
-        address existingContract = _contractRegistry[contractId];
-
-        if (existingContract == address(0)) {
-            existingContract = _contractAliases[contractId];
-        }
-
-        if (existingContract != address(0)) {
-            info = _contractInfo[existingContract];
-
-            revert ContractAlreadyRegistered(info.contractType, contractName);
-        }
+        bytes32 contractId = _ensureUniqueName(contractName, true);
 
         // Store the address in the registry, under the unique name.
         _contractRegistry[contractId] = contractAddress;
@@ -166,7 +158,9 @@ contract BalancerContractRegistry is IBalancerContractRegistry, SingletonAuthent
 
         ContractInfo memory info = _contractInfo[contractAddress];
         // This should be impossible: the registry and info mappings must be in sync.
-        require(info.isRegistered, "Registry entry with no contract info");
+        if (info.isRegistered == false) {
+            revert InconsistentState();
+        }
 
         delete _contractRegistry[contractId];
         delete _contractInfo[contractAddress];
@@ -218,17 +212,10 @@ contract BalancerContractRegistry is IBalancerContractRegistry, SingletonAuthent
         }
 
         // Ensure the proposed alias is not in use (i.e., no collision with existing registered contracts).
-        // It can already be in `_contractAliases`; that's the "update" case. For instance, when we want to
-        // migrate the `WeightedPool` alias from v2 to v3. If the name is not already in `_contractAliases`,
-        // we are adding a new alias.
-        bytes32 contractId = _getContractId(contractAlias);
-        address registeredContract = _contractRegistry[contractId];
-
-        if (registeredContract != address(0)) {
-            info = _contractInfo[registeredContract];
-
-            revert ContractAlreadyRegistered(info.contractType, contractAlias);
-        }
+        // It can match an existing alias: that's the "update" case. For instance, if we wanted to migrate
+        // the `WeightedPool` alias from v2 to v3. If the name is not already in `_contractAliases`, we are
+        // adding a new alias.
+        bytes32 contractId = _ensureUniqueName(contractAlias, false);
 
         // This will either add a new or overwrite an existing alias.
         _contractAliases[contractId] = contractAddress;
@@ -272,6 +259,25 @@ contract BalancerContractRegistry is IBalancerContractRegistry, SingletonAuthent
         return _contractInfo[contractAddress];
     }
 
+    function _ensureUniqueName(
+        string memory contractName,
+        bool includeAliases
+    ) internal view returns (bytes32 contractId) {
+        contractId = _getContractId(contractName);
+
+        // Check both the registered names and aliases to ensure this name is not a duplicate.
+        address existingContract = _contractRegistry[contractId];
+        if (includeAliases && existingContract == address(0)) {
+            existingContract = _contractAliases[contractId];
+        }
+
+        if (existingContract != address(0)) {
+            ContractInfo memory info = _contractInfo[existingContract];
+
+            revert ContractAlreadyRegistered(info.contractType, contractName);
+        }
+    }
+
     function _getContractId(string memory contractName) internal pure returns (bytes32) {
         return keccak256(abi.encodePacked(contractName));
     }