Batch relayer chaining (#965)

* Add temp storage read/writes

* Add chained reference detection and key extraction

* Add chained reference functions

* Add chained references to simple swap

* Add Math.abs

* Add chained references to batch swap

* Add some docs

* Remove unnecessary distinction between references and keys

* Copy toChainedReference helper

* refactor: make _TEMP_STORAGE_PREFIX constant and use as a prefix

* fix: revert swapping ref and _TEMP_STORAGE_PREFIX and rename to _TEMP_STORAGE_SUFFIX

Swapping the order meant that we were no longer replicating the standard mapping mechanism

* revert: change _TEMP_STORAGE_SUFFIX to immutable

see: ethereum/solidity#9232

* Add joinPool output chained reference

* Add join support

* Add more join tests

* Add test helpers

* Fix linter

* Update pkg/standalone-utils/contracts/relayer/VaultActions.sol

Co-authored-by: Tom French <[email protected]>

* refactor: remove unused imports

Co-authored-by: Tom French <[email protected]>
Co-authored-by: Tom French <[email protected]>

pkg/solidity-utils/contracts/helpers/VaultHelpers.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+pragma solidity ^0.7.0;
+
+library VaultHelpers {
+    /**
+     * @dev Returns the address of a Pool's contract.
+     *
+     * This is the same code the Vault runs in `PoolRegistry._getPoolAddress`.
+     */
+    function toPoolAddress(bytes32 poolId) internal pure returns (address) {
+        // 12 byte logical shift left to remove the nonce and specialization setting. We don't need to mask,
+        // since the logical shift already sets the upper bits to zero.
+        return address(uint256(poolId) >> (12 * 8));
+    }
+}

pkg/solidity-utils/contracts/math/Math.sol

@@ -6,9 +6,16 @@ import "../helpers/BalancerErrors.sol";
 
 /**
  * @dev Wrappers over Solidity's arithmetic operations with added overflow checks.
- * Adapted from OpenZeppelin's SafeMath library
+ * Adapted from OpenZeppelin's SafeMath library.
  */
 library Math {
+    /**
+     * @dev Returns the absolute value of a signed integer.
+     */
+    function abs(int256 a) internal pure returns (uint256) {
+        return a > 0 ? uint256(a) : uint256(-a);
+    }
+
     /**
      * @dev Returns the addition of two unsigned integers of 256 bits, reverting on overflow.
      */

pkg/solidity-utils/contracts/test/MathMock.sol

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+pragma solidity ^0.7.0;
+
+import "../math/Math.sol";
+
+contract MathMock {
+    function abs(int256 a) public pure returns (uint256) {
+        return Math.abs(a);
+    }
+}

pkg/solidity-utils/test/Math.test.ts

@@ -0,0 +1,33 @@
+import { expect } from 'chai';
+
+import { Contract } from 'ethers';
+import { deploy } from '@balancer-labs/v2-helpers/src/contract';
+import { MAX_INT256, MIN_INT256 } from '@balancer-labs/v2-helpers/src/constants';
+
+describe('Math', () => {
+  let lib: Contract;
+
+  sharedBeforeEach('deploy lib', async () => {
+    lib = await deploy('MathMock', { args: [] });
+  });
+
+  it('handles zero', async () => {
+    expect(await lib.abs(0)).to.equal(0);
+  });
+
+  it('handles positive values', async () => {
+    expect(await lib.abs(42)).to.equal(42);
+  });
+
+  it('handles large positive values', async () => {
+    expect(await lib.abs(MAX_INT256)).to.equal(MAX_INT256);
+  });
+
+  it('handles negative values', async () => {
+    expect(await lib.abs(-3)).to.equal(3);
+  });
+
+  it('handles large negative values', async () => {
+    expect(await lib.abs(MIN_INT256)).to.equal(MIN_INT256.mul(-1));
+  });
+});

pkg/standalone-utils/contracts/interfaces/IBaseRelayerLibrary.sol

@@ -22,4 +22,10 @@ import "@balancer-labs/v2-vault/contracts/interfaces/IVault.sol";
  */
 abstract contract IBaseRelayerLibrary {
     function getVault() public view virtual returns (IVault);
+
+    function _isChainedReference(uint256 amount) internal pure virtual returns (bool);
+
+    function _setChainedReferenceValue(uint256 ref, uint256 value) internal virtual;
+
+    function _getChainedReferenceValue(uint256 ref) internal virtual returns (uint256);
 }

pkg/standalone-utils/contracts/relayer/BaseRelayerLibrary.sol

@@ -64,4 +64,54 @@ contract BaseRelayerLibrary is IBaseRelayerLibrary {
 
         address(_vault).functionCall(data);
     }
+
+    /**
+     * @dev Returns true if `amount` is not actually an amount, but rather a chained reference.
+     */
+    function _isChainedReference(uint256 amount) internal pure override returns (bool) {
+        return
+            (amount & 0xffff000000000000000000000000000000000000000000000000000000000000) ==
+            0xba10000000000000000000000000000000000000000000000000000000000000;
+    }
+
+    /**
+     * @dev Stores `value` as the amount referenced by chained reference `ref`.
+     */
+    function _setChainedReferenceValue(uint256 ref, uint256 value) internal override {
+        bytes32 slot = _getTempStorageSlot(ref);
+
+        // Since we do manual calculation of storage slots, it is easier (and cheaper) to rely on internal assembly to
+        // access it.
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+            sstore(slot, value)
+        }
+    }
+
+    /**
+     * @dev Returns the amount referenced by chained reference `ref`. Reading an amount clears it, so they can each
+     * only be read once.
+     */
+    function _getChainedReferenceValue(uint256 ref) internal override returns (uint256 value) {
+        bytes32 slot = _getTempStorageSlot(ref);
+
+        // Since we do manual calculation of storage slots, it is easier (and cheaper) to rely on internal assembly to
+        // access it.
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+            value := sload(slot)
+            sstore(slot, 0)
+        }
+    }
+
+    bytes32 private immutable _TEMP_STORAGE_SUFFIX = keccak256("balancer.base-relayer-library");
+
+    function _getTempStorageSlot(uint256 ref) private view returns (bytes32) {
+        // This replicates the mechanism Solidity uses to allocate storage slots for mappings, but using a hash as the
+        // mapping's storage slot, and subtracting 1 at the end. This should be more enough to prevent collisions with
+        // other state variables this or derived contracts might use.
+        // See https://docs.soliditylang.org/en/v0.8.9/internals/layout_in_storage.html
+
+        return bytes32(uint256(keccak256(abi.encodePacked(ref, _TEMP_STORAGE_SUFFIX))) - 1);
+    }
 }

pkg/standalone-utils/contracts/relayer/VaultActions.sol

@@ -15,9 +15,15 @@
 pragma solidity ^0.7.0;
 pragma experimental ABIEncoderV2;
 
-import "@balancer-labs/v2-solidity-utils/contracts/openzeppelin/IERC20Permit.sol";
-import "@balancer-labs/v2-solidity-utils/contracts/openzeppelin/IERC20PermitDAI.sol";
+import "@balancer-labs/v2-solidity-utils/contracts/helpers/InputHelpers.sol";
+import "@balancer-labs/v2-solidity-utils/contracts/helpers/VaultHelpers.sol";
+import "@balancer-labs/v2-solidity-utils/contracts/math/Math.sol";
+
 import "@balancer-labs/v2-vault/contracts/interfaces/IVault.sol";
+
+import "@balancer-labs/v2-pool-weighted/contracts/BaseWeightedPool.sol";
+import "@balancer-labs/v2-pool-weighted/contracts/WeightedPoolUserDataHelpers.sol";
+
 import "../interfaces/IBaseRelayerLibrary.sol";
 
 /**
@@ -30,28 +36,66 @@ import "../interfaces/IBaseRelayerLibrary.sol";
  * All functions must be payable so that it can be called as part of a multicall involving ETH
  */
 abstract contract VaultActions is IBaseRelayerLibrary {
+    using Math for uint256;
+
     function swap(
-        IVault.SingleSwap calldata singleSwap,
+        IVault.SingleSwap memory singleSwap,
         IVault.FundManagement calldata funds,
         uint256 limit,
         uint256 deadline,
-        uint256 value
+        uint256 value,
+        uint256 outputReference
     ) external payable returns (uint256) {
         require(funds.sender == msg.sender, "Incorrect sender");
-        return getVault().swap{ value: value }(singleSwap, funds, limit, deadline);
+
+        if (_isChainedReference(singleSwap.amount)) {
+            singleSwap.amount = _getChainedReferenceValue(singleSwap.amount);
+        }
+
+        uint256 result = getVault().swap{ value: value }(singleSwap, funds, limit, deadline);
+
+        if (_isChainedReference(outputReference)) {
+            _setChainedReferenceValue(outputReference, result);
+        }
+
+        return result;
     }
 
     function batchSwap(
         IVault.SwapKind kind,
-        IVault.BatchSwapStep[] calldata swaps,
+        IVault.BatchSwapStep[] memory swaps,
         IAsset[] calldata assets,
         IVault.FundManagement calldata funds,
         int256[] calldata limits,
         uint256 deadline,
-        uint256 value
+        uint256 value,
+        uint256[] memory outputReferences
     ) external payable returns (int256[] memory) {
         require(funds.sender == msg.sender, "Incorrect sender");
-        return getVault().batchSwap{ value: value }(kind, swaps, assets, funds, limits, deadline);
+        InputHelpers.ensureInputLengthMatch(assets.length, outputReferences.length);
+
+        for (uint256 i = 0; i < swaps.length; ++i) {
+            uint256 amount = swaps[i].amount;
+            if (_isChainedReference(amount)) {
+                swaps[i].amount = _getChainedReferenceValue(amount);
+            }
+        }
+
+        int256[] memory results = getVault().batchSwap{ value: value }(kind, swaps, assets, funds, limits, deadline);
+
+        for (uint256 i = 0; i < outputReferences.length; ++i) {
+            uint256 ref = outputReferences[i];
+            if (_isChainedReference(ref)) {
+                // Batch swap return values are signed, as they are Vault deltas (positive values stand for assets sent
+                // to the Vault, negatives for assets sent from the Vault). To simplify the chained reference value
+                // model, we simply store the absolute value.
+                // This should be fine for most use cases, as the caller can reason about swap results via the `limits`
+                // parameter.
+                _setChainedReferenceValue(ref, Math.abs(results[i]));
+            }
+        }
+
+        return results;
     }
 
     function manageUserBalance(IVault.UserBalanceOp[] calldata ops, uint256 value) external payable {
@@ -61,15 +105,80 @@ abstract contract VaultActions is IBaseRelayerLibrary {
         getVault().manageUserBalance{ value: value }(ops);
     }
 
+    enum PoolKind { WEIGHTED }
+
     function joinPool(
         bytes32 poolId,
+        PoolKind kind,
         address sender,
         address recipient,
-        IVault.JoinPoolRequest calldata request,
-        uint256 value
+        IVault.JoinPoolRequest memory request,
+        uint256 value,
+        uint256 outputReference
     ) external payable {
         require(sender == msg.sender, "Incorrect sender");
+
+        // The output of a join is expected to be balance in the Pool's token contract, typically known as BPT (Balancer
+        // Pool Tokens). Since the Vault is unaware of this (BPT is minted directly to the recipient), we manually
+        // measure this balance increase (but only if an output reference is provided).
+        IERC20 bpt = IERC20(VaultHelpers.toPoolAddress(poolId));
+        uint256 maybeInitialRecipientBPT = _isChainedReference(outputReference) ? bpt.balanceOf(recipient) : 0;
+
+        request.userData = _doJoinPoolChainedReferenceReplacements(kind, request.userData);
+
         getVault().joinPool{ value: value }(poolId, sender, recipient, request);
+
+        if (_isChainedReference(outputReference)) {
+            // In this context, `maybeInitialRecipientBPT` is guaranteed to have been initialized, so we can safely read
+            // from it. Note that we assume that the recipient balance change has a positive sign (i.e. the recipient
+            // received BPT).
+            uint256 finalRecipientBPT = bpt.balanceOf(recipient);
+            _setChainedReferenceValue(outputReference, finalRecipientBPT.sub(maybeInitialRecipientBPT));
+        }
+    }
+
+    function _doJoinPoolChainedReferenceReplacements(PoolKind kind, bytes memory userData)
+        private
+        returns (bytes memory)
+    {
+        if (kind == PoolKind.WEIGHTED) {
+            return _doWeightedJoinChainedReferenceReplacements(userData);
+        } else {
+            _revert(Errors.UNHANDLED_JOIN_KIND);
+        }
+    }
+
+    function _doWeightedJoinChainedReferenceReplacements(bytes memory userData) private returns (bytes memory) {
+        BaseWeightedPool.JoinKind kind = WeightedPoolUserDataHelpers.joinKind(userData);
+
+        if (kind == BaseWeightedPool.JoinKind.EXACT_TOKENS_IN_FOR_BPT_OUT) {
+            return _doWeightedExactTokensInForBPTOutReplacements(userData);
+        } else {
+            // All other join kinds are 'given out' (i.e the parameter is a BPT amount), so we don't do replacements for
+            // those.
+            return userData;
+        }
+    }
+
+    function _doWeightedExactTokensInForBPTOutReplacements(bytes memory userData) private returns (bytes memory) {
+        (uint256[] memory amountsIn, uint256 minBPTAmountOut) = WeightedPoolUserDataHelpers.exactTokensInForBptOut(
+            userData
+        );
+
+        bool replacedAmounts = false;
+        for (uint256 i = 0; i < amountsIn.length; ++i) {
+            uint256 amount = amountsIn[i];
+            if (_isChainedReference(amount)) {
+                amountsIn[i] = _getChainedReferenceValue(amount);
+                replacedAmounts = true;
+            }
+        }
+
+        // Save gas by only re-encoding the data if we actually performed a replacement
+        return
+            replacedAmounts
+                ? abi.encode(BaseWeightedPool.JoinKind.EXACT_TOKENS_IN_FOR_BPT_OUT, amountsIn, minBPTAmountOut)
+                : userData;
     }
 
     function exitPool(

pkg/standalone-utils/contracts/test/MockBaseRelayerLibrary.sol

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+// Copyright (C) 2015, 2016, 2017 Dapphub
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+pragma solidity ^0.7.0;
+
+import "../relayer/BaseRelayerLibrary.sol";
+
+contract MockBaseRelayerLibrary is BaseRelayerLibrary {
+    event ChainedReferenceValueRead(uint256 value);
+
+    constructor(IVault vault) BaseRelayerLibrary(vault) {}
+
+    function isChainedReference(uint256 amount) public pure returns (bool) {
+        return _isChainedReference(amount);
+    }
+
+    function setChainedReferenceValue(uint256 ref, uint256 value) public returns (uint256) {
+        _setChainedReferenceValue(ref, value);
+    }
+
+    function getChainedReferenceValue(uint256 ref) public {
+        emit ChainedReferenceValueRead(_getChainedReferenceValue(ref));
+    }
+}