Multi Flow Pump v1.1

.gitignore

@@ -17,3 +17,6 @@ stash
 lcov.info
 .DS_Store
 test/output/
+
+# Python
+__pycache__

mocks/wells/MockReserveWell.sol

@@ -5,17 +5,28 @@
 pragma solidity ^0.8.20;
 
 import {IPump} from "src/interfaces/pumps/IPump.sol";
+import {Call} from "src/interfaces/IWell.sol";
 
 /**
  * @notice Mock Well that allows setting of reserves.
  */
 contract MockReserveWell {
     uint256[] reserves;
+    Call _wellFunction;
+
 
     constructor() {
         reserves = new uint256[](2);
     }
 
+    function setWellFunction(Call calldata __wellFunction) external {
+        _wellFunction = __wellFunction;
+    }
+
+    function wellFunction() external view returns (Call memory) {
+        return _wellFunction;
+    }
+
     function setReserves(uint256[] memory _reserves) public {
         reserves = _reserves;
     }

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@beanstalk/wells",
-  "version": "0.4.0",
+  "version": "1.1.0-prerelease0",
   "description": "A [{Well}](/src/Well.sol) is a constant function AMM that allows the provisioning of liquidity into a single pooled on-chain liquidity position.",
   "main": "index.js",
   "directories": {

src/functions/ConstantProduct.sol

@@ -23,6 +23,8 @@ import {LibMath} from "src/libraries/LibMath.sol";
 contract ConstantProduct is ProportionalLPToken, IBeanstalkWellFunction {
     using LibMath for uint256;
 
+    uint256 constant CALC_RATE_PRECISION = 1e18;
+
     /// @dev `s = π(b_i)^(1/n) * n`
     function calcLpTokenSupply(
         uint256[] calldata reserves,
@@ -92,4 +94,13 @@ contract ConstantProduct is ProportionalLPToken, IBeanstalkWellFunction {
         }
         reserve /= reserves.length - 1;
     }
+
+    function calcRate(
+        uint256[] calldata reserves,
+        uint256 i,
+        uint256 j,
+        bytes calldata
+    ) external pure returns (uint256 rate) {
+        return reserves[i] * CALC_RATE_PRECISION / reserves[j];
+    }
 }

src/functions/ConstantProduct2.sol

@@ -5,6 +5,7 @@ pragma solidity ^0.8.20;
 import {IBeanstalkWellFunction} from "src/interfaces/IBeanstalkWellFunction.sol";
 import {ProportionalLPToken2} from "src/functions/ProportionalLPToken2.sol";
 import {LibMath} from "src/libraries/LibMath.sol";
+import {Math} from "oz/utils/math/Math.sol";
 
 /**
  * @title ConstantProduct2
@@ -18,9 +19,10 @@ import {LibMath} from "src/libraries/LibMath.sol";
  *  `b_i` is the reserve at index `i`
  */
 contract ConstantProduct2 is ProportionalLPToken2, IBeanstalkWellFunction {
-    using LibMath for uint256;
+    using Math for uint256;
 
     uint256 constant EXP_PRECISION = 1e12;
+    uint256 constant CALC_RATE_PRECISION = 1e18;
 
     /**
      * @dev `s = (b_0 * b_1)^(1/2)`
@@ -103,4 +105,13 @@ contract ConstantProduct2 is ProportionalLPToken2, IBeanstalkWellFunction {
         uint256 i = j == 1 ? 0 : 1;
         reserve = reserves[i] * ratios[j] / ratios[i];
     }
+
+    function calcRate(
+        uint256[] calldata reserves,
+        uint256 i,
+        uint256 j,
+        bytes calldata
+    ) external pure returns (uint256 rate) {
+        return reserves[i] * CALC_RATE_PRECISION / reserves[j];
+    }
 }

src/functions/ProportionalLPToken.sol

@@ -3,6 +3,7 @@
 pragma solidity ^0.8.20;
 
 import {IWellFunction} from "src/interfaces/IWellFunction.sol";
+import {Math} from "oz/utils/math/Math.sol";
 
 /**
  * @title ProportionalLPToken
@@ -12,6 +13,8 @@ import {IWellFunction} from "src/interfaces/IWellFunction.sol";
  * recieves `s * b_i / S` of each underlying token.
  */
 abstract contract ProportionalLPToken is IWellFunction {
+    using Math for uint256;
+
     function calcLPTokenUnderlying(
         uint256 lpTokenAmount,
         uint256[] calldata reserves,
@@ -20,7 +23,7 @@ abstract contract ProportionalLPToken is IWellFunction {
     ) external pure returns (uint256[] memory underlyingAmounts) {
         underlyingAmounts = new uint256[](reserves.length);
         for (uint256 i; i < reserves.length; ++i) {
-            underlyingAmounts[i] = lpTokenAmount * reserves[i] / lpTokenSupply;
+            underlyingAmounts[i] = lpTokenAmount.mulDiv(reserves[i], lpTokenSupply);
         }
     }
 }

src/functions/ProportionalLPToken2.sol

@@ -3,6 +3,7 @@
 pragma solidity ^0.8.20;
 
 import {IWellFunction} from "src/interfaces/IWellFunction.sol";
+import {Math} from "oz/utils/math/Math.sol";
 
 /**
  * @title ProportionalLPToken2
@@ -12,14 +13,16 @@ import {IWellFunction} from "src/interfaces/IWellFunction.sol";
  * recieves `s * b_i / S` of each underlying token.
  */
 abstract contract ProportionalLPToken2 is IWellFunction {
+    using Math for uint256;
+
     function calcLPTokenUnderlying(
         uint256 lpTokenAmount,
         uint256[] calldata reserves,
         uint256 lpTokenSupply,
         bytes calldata
     ) external pure returns (uint256[] memory underlyingAmounts) {
         underlyingAmounts = new uint256[](2);
-        underlyingAmounts[0] = lpTokenAmount * reserves[0] / lpTokenSupply;
-        underlyingAmounts[1] = lpTokenAmount * reserves[1] / lpTokenSupply;
+        underlyingAmounts[0] = lpTokenAmount.mulDiv(reserves[0], lpTokenSupply);
+        underlyingAmounts[1] = lpTokenAmount.mulDiv(reserves[1], lpTokenSupply);
     }
 }

src/interfaces/IBeanstalkWellFunction.sol

@@ -2,34 +2,18 @@
 
 pragma solidity ^0.8.20;
 
-import {IWellFunction} from "src/interfaces/IWellFunction.sol";
+import {IMultiFlowPumpWellFunction} from "src/interfaces/IMultiFlowPumpWellFunction.sol";
 
 /**
  * @title IBeanstalkWellFunction
  * @notice Defines all necessary functions for Beanstalk to support a Well Function in addition to functions defined in the primary interface.
- * This includes 2 functions to solve for a given reserve value suc that the average price between
+ * It extends `IMultiFlowPumpWellFunction` as Beanstalk requires Wells to use MultiFlowPump in order to have access to manipulation resistant oracles.
+ * Beanstalk requires 2 functions to solve for a given reserve value such that the average price between
  * the given reserve and all other reserves equals the average of the input ratios.
- * `calcReserveAtRatioSwap` assumes the target ratios are reached through executing a swap.
- * `calcReserveAtRatioLiquidity` assumes the target ratios are reached through adding/removing liquidity.
+ * - `calcReserveAtRatioSwap` assumes the target ratios are reached through executing a swap. Note: `calcReserveAtRatioSwap` is included in {IMultiFlowPumpWellFunction}.
+ * - `calcReserveAtRatioLiquidity` assumes the target ratios are reached through adding/removing liquidity.
  */
-interface IBeanstalkWellFunction is IWellFunction {
-    /**
-     * @notice Calculates the `j` reserve such that `π_{i | i != j} (d reserves_j / d reserves_i) = π_{i | i != j}(ratios_j / ratios_i)`.
-     * assumes that reserve_j is being swapped for other reserves in the Well.
-     * @dev used by Beanstalk to calculate the deltaB every Season
-     * @param reserves The reserves of the Well
-     * @param j The index of the reserve to solve for
-     * @param ratios The ratios of reserves to solve for
-     * @param data Well function data provided on every call
-     * @return reserve The resulting reserve at the jth index
-     */
-    function calcReserveAtRatioSwap(
-        uint256[] calldata reserves,
-        uint256 j,
-        uint256[] calldata ratios,
-        bytes calldata data
-    ) external view returns (uint256 reserve);
-
+interface IBeanstalkWellFunction is IMultiFlowPumpWellFunction {
     /**
      * @notice Calculates the `j` reserve such that `π_{i | i != j} (d reserves_j / d reserves_i) = π_{i | i != j}(ratios_j / ratios_i)`.
      * assumes that reserve_j is being added or removed in exchange for LP Tokens.

src/interfaces/IMultiFlowPumpWellFunction.sol

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.20;
+
+import {IWellFunction} from "src/interfaces/IWellFunction.sol";
+
+/**
+ * @title IMultiFlowPumpWellFunction
+ * @dev A Well Function must implement IMultiFlowPumpWellFunction to be supported by
+ * the Multi Flow Pump.
+ */
+interface IMultiFlowPumpWellFunction is IWellFunction {
+    /**
+     * @notice Calculates the `j` reserve such that `π_{i | i != j} (d reserves_j / d reserves_i) = π_{i | i != j}(ratios_j / ratios_i)`.
+     * assumes that reserve_j is being swapped for other reserves in the Well.
+     * @dev used by Beanstalk to calculate the deltaB every Season
+     * @param reserves The reserves of the Well
+     * @param j The index of the reserve to solve for
+     * @param ratios The ratios of reserves to solve for
+     * @param data Well function data provided on every call
+     * @return reserve The resulting reserve at the jth index
+     */
+    function calcReserveAtRatioSwap(
+        uint256[] calldata reserves,
+        uint256 j,
+        uint256[] calldata ratios,
+        bytes calldata data
+    ) external view returns (uint256 reserve);
+
+    /**
+     * @notice Calculates the rate at which j can be exchanged for i.
+     * @param reserves The reserves of the Well
+     * @param i The index of the token for which the output is being calculated
+     * @param j The index of the token for which 1 token is being exchanged
+     * @param data Well function data provided on every call
+     * @return rate The rate at which j can be exchanged for i
+     * @dev should return with 36 decimal precision
+     */
+    function calcRate(
+        uint256[] calldata reserves,
+        uint256 i,
+        uint256 j,
+        bytes calldata data
+    ) external view returns (uint256 rate);
+}

src/interfaces/pumps/IMultiFlowPumpErrors.sol

@@ -11,4 +11,6 @@ interface IMultiFlowPumpErrors {
     error NotInitialized();
 
     error NoTimePassed();
+
+    error TooManyTokens();
 }

src/libraries/LibLastReserveBytes.sol

@@ -2,6 +2,8 @@
 
 pragma solidity ^0.8.20;
 
+import {ABDKMathQuad} from "src/libraries/ABDKMathQuad.sol";
+
 /**
  * @title LibLastReserveBytes
  * @author Publius
@@ -15,15 +17,25 @@ pragma solidity ^0.8.20;
  * in reserves for manipulation resistance purposes, the gas savings is worth the lose of precision.
  */
 library LibLastReserveBytes {
+    using ABDKMathQuad for uint256;
+    using ABDKMathQuad for bytes16;
+
     function readNumberOfReserves(bytes32 slot) internal view returns (uint8 _numberOfReserves) {
         assembly {
             _numberOfReserves := shr(248, sload(slot))
         }
     }
 
-    function storeLastReserves(bytes32 slot, uint40 lastTimestamp, bytes16[] memory reserves) internal {
+    function storeLastReserves(bytes32 slot, uint40 lastTimestamp, uint256[] memory lastReserves) internal {
         // Potential optimization – shift reserve bytes left to perserve extra decimal precision.
-        uint8 n = uint8(reserves.length);
+        uint8 n = uint8(lastReserves.length);
+
+        bytes16[] memory reserves = new bytes16[](n);
+
+        for (uint256 i; i < n; ++i) {
+            reserves[i] = lastReserves[i].fromUInt();
+        }
+
         if (n == 1) {
             assembly {
                 sstore(slot, or(or(shl(208, lastTimestamp), shl(248, n)), shl(104, shr(152, mload(add(reserves, 32))))))
@@ -73,7 +85,7 @@ library LibLastReserveBytes {
     function readLastReserves(bytes32 slot)
         internal
         view
-        returns (uint8 n, uint40 lastTimestamp, bytes16[] memory reserves)
+        returns (uint8 n, uint40 lastTimestamp, uint256[] memory lastReserves)
     {
         // Shortcut: two reserves can be quickly unpacked from one slot
         bytes32 temp;
@@ -82,13 +94,17 @@ library LibLastReserveBytes {
             n := shr(248, temp)
             lastTimestamp := shr(208, temp)
         }
-        if (n == 0) return (n, lastTimestamp, reserves);
+        if (n == 0) return (n, lastTimestamp, lastReserves);
         // Initialize array with length `n`, fill it in via assembly
-        reserves = new bytes16[](n);
+        bytes16[] memory reserves = new bytes16[](n);
         assembly {
             mstore(add(reserves, 32), shl(152, shr(104, temp)))
         }
-        if (n == 1) return (n, lastTimestamp, reserves);
+        if (n == 1) {
+            lastReserves = new uint256[](1);
+            lastReserves[0] = reserves[0].toUInt();
+            return (n, lastTimestamp, lastReserves);
+        }
         assembly {
             mstore(add(reserves, 64), shl(152, temp))
         }
@@ -116,6 +132,11 @@ library LibLastReserveBytes {
                 }
             }
         }
+
+        lastReserves = new uint256[](n);
+        for (uint256 i; i < n; ++i) {
+            lastReserves[i] = reserves[i].toUInt();
+        }
     }
 
     function readBytes(bytes32 slot) internal view returns (bytes32 value) {