Add Generalized Script and Task for Well Component Upgrades

contracts/beanstalk/init/InitUpgradeWell.sol

@@ -0,0 +1,110 @@
+/*
+ SPDX-License-Identifier: MIT
+*/
+pragma solidity ^0.8.20;
+pragma experimental ABIEncoderV2;
+
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import {LibWellDeployer} from "contracts/libraries/Basin/LibWellDeployer.sol";
+import {IWellUpgradeable} from "contracts/interfaces/basin/IWellUpgradeable.sol";
+import {IAquifer} from "contracts/interfaces/basin/IAquifer.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import {IWell, Call} from "contracts/interfaces/basin/IWell.sol";
+
+/**
+ * @title InitUpgradeWell
+ * Upgrades a set of upgradeable wells to use new arbitrary components.
+ * Intended for when fixes or updates to a well's components but keeping the same data.
+ */
+contract InitUpgradeWell {
+    // A default well salt is used to prevent front-running attacks as boring wells also uses msg.sender with non-zero salt.
+    bytes32 internal constant DEFAULT_WELL_SALT =
+        0x0000000000000000000000000000000000000000000000000000000000000010;
+
+    /**
+     * @notice Upgrades a set of upgradeable wells to use new arbitrary components.
+     * @param wellsToUpgrade The addresses of the wells to upgrade.
+     * For now only the Upgradeable well implementation is supported.
+     * @param newWellFunctionTarget The new well function target to use or address(0) for no change.
+     * @param newPumpTarget The new pump target to use or address(0) for no change.
+     */
+    function init(
+        address[] memory wellsToUpgrade,
+        address newWellFunctionTarget,
+        address newPumpTarget
+    ) external {
+        for (uint256 i; i < wellsToUpgrade.length; i++) {
+            address wellToUpgrade = wellsToUpgrade[i];
+
+            // get well components
+            (
+                IERC20[] memory tokens,
+                Call memory wellFunction,
+                Call[] memory pumps, // well data
+                ,
+                address aquifer
+            ) = IWell(wellToUpgrade).well();
+
+            // get the upgradeable well implementation component address
+            address wellImplementation = getWellUpgradeableImplementation(wellToUpgrade, aquifer);
+
+            // if specified, create new well function with updated target but same data
+            wellFunction = newWellFunctionTarget != address(0)
+                ? Call(newWellFunctionTarget, wellFunction.data)
+                : wellFunction;
+
+            // if specified, update the well's first pump with updated target but same data
+            pumps[0] = newPumpTarget != address(0) ? Call(newPumpTarget, pumps[0].data) : pumps[0];
+
+            deployAndUpgradeWell(
+                tokens,
+                wellFunction,
+                pumps,
+                aquifer,
+                wellImplementation,
+                wellToUpgrade
+            );
+        }
+    }
+
+    /**
+     * @notice Deploys a minimal proxy well with the upgradeable well implementation and a
+     * ERC1967Proxy in front of it to allow for future upgrades.
+     * Upgrades the existing well to the new implementation.
+     */
+    function deployAndUpgradeWell(
+        IERC20[] memory tokens,
+        Call memory wellFunction,
+        Call[] memory pumps,
+        address aquifer,
+        address wellImplementation,
+        address wellToUpgrade
+    ) internal {
+        // Encode well data
+        (bytes memory immutableData, bytes memory initData) = LibWellDeployer
+            .encodeUpgradeableWellDeploymentData(aquifer, tokens, wellFunction, pumps);
+
+        // Bore upgradeable well with the same salt for reproducibility.
+        address _well = IAquifer(aquifer).boreWell(
+            wellImplementation,
+            immutableData,
+            initData,
+            DEFAULT_WELL_SALT
+        );
+
+        // Upgrade the well to the new implementation
+        IWellUpgradeable(payable(wellToUpgrade)).upgradeTo(_well);
+    }
+
+    /**
+     * @notice Returns the WellUpgradeable standalone component address of an upgradeable Well.
+     * by traversing through the ERC1967Proxy and the minimal proxy chain and finally fetching the Aquifer registry.
+     */
+    function getWellUpgradeableImplementation(
+        address wellToUpgrade,
+        address aquifer
+    ) internal view returns (address) {
+        address minimalProxy = IWellUpgradeable(wellToUpgrade).getImplementation();
+        return IAquifer(aquifer).wellImplementation(minimalProxy);
+    }
+}

contracts/interfaces/basin/IWellUpgradeable.sol

@@ -19,4 +19,6 @@ interface IWellUpgradeable is IWell {
     function upgradeTo(address newImplementation) external;
 
     function upgradeToAndCall(address newImplementation, bytes memory data) external;
+
+    function getImplementation() external view returns (address);
 }

hardhat.config.js

@@ -30,7 +30,9 @@ const {
   PINTO_WSOL_WELL_BASE,
   nameToAddressMap,
   addressToNameMap,
-  addressToBalanceSlotMap
+  addressToBalanceSlotMap,
+  STABLE_2_BASE,
+  ZERO_ADDRESS
 } = require("./test/hardhat/utils/constants.js");
 const { task } = require("hardhat/config");
 const { upgradeWithNewFacets, decodeDiamondCutAction } = require("./scripts/diamond.js");
@@ -989,6 +991,50 @@ task(
   });
 });
 
+task("upgrade-wells", "Upgrades wells to new well function and pump targets")
+  .addParam("wells", "Comma-separated list of well addresses to upgrade")
+  .addOptionalParam("newfunction", "Address of the new well function target", ZERO_ADDRESS)
+  .addOptionalParam("newpump", "Address of the new pump target", ZERO_ADDRESS)
+  .setAction(async function (taskArgs) {
+    let owner;
+    let mock = true;
+    if (mock) {
+      owner = await impersonateSigner(L2_PCM);
+      await mintEth(owner.address);
+    } else {
+      owner = (await ethers.getSigners())[0];
+    }
+
+    // Parse wells array from comma-separated string
+    const wellsToUpgrade = taskArgs.wells.split(",").map(well => well.trim());
+    const newWellFunctionTarget = taskArgs.newfunction;
+    const newPumpTarget = taskArgs.newpump;
+
+    // Validate that at least one target is not zero address
+    if (newWellFunctionTarget === ZERO_ADDRESS && newPumpTarget === ZERO_ADDRESS) {
+      throw new Error("At least one of newfunction or newpump must be specified (not zero address)");
+    }
+
+    console.log("Wells to upgrade:", wellsToUpgrade);
+    if (newWellFunctionTarget !== ZERO_ADDRESS) {
+      console.log("New well function target:", newWellFunctionTarget);
+    }
+    if (newPumpTarget !== ZERO_ADDRESS) {
+      console.log("New pump target:", newPumpTarget);
+    }
+
+    // upgrade facets
+    await upgradeWithNewFacets({
+      diamondAddress: L2_PINTO,
+      facetNames: [],
+      initFacetName: "InitUpgradeWell",
+      initArgs: [wellsToUpgrade, newWellFunctionTarget, newPumpTarget],
+      object: !mock,
+      verbose: true,
+      account: owner
+    });
+  });
+
 task("whitelist-rebalance", "Deploys whitelist rebalance").setAction(async function () {
   const mock = true;
   let owner;

test/hardhat/utils/constants.js

@@ -129,8 +129,8 @@ module.exports = {
 
   // Basin
   AQUIFER_BASE: "0xBA51AA60B3b8d9A36cc748a62Aa56801060183f8",
-  BEANSTALK_PUMP_BASE: "0xBA51AA73ab1b8720E6D5602Bd3cBaaedB6399133",
-  STABLE_2_BASE: "0xBA51055a97b40d7f41f3F64b57469b5D45B67c87",
+  BEANSTALK_PUMP_BASE: "0xBA51AAaA66DaB6c236B356ad713f759c206DcB93",
+  STABLE_2_BASE: "0xBA51A1151c76D8CA41DE379e16b8732620fb3c8D",
   WELL_IMPLEMENTATION_BASE: "0xBA5100A00aa91589b02D094AdbE38aA7F13B4b50",
   WELL_IMPLEMENTATION_UPGRADEABLE_BASE: "0xBA510990a720725Ab1F9a0D231F045fc906909f4",
   CONSTANT_PRODUCT_2_BASE: "0xBA510C289fD067EBbA41335afa11F0591940d6fe",