Created sepaerate mixins for interacting with the different vaults

contracts/staking/contracts/src/Staking.sol

@@ -25,10 +25,14 @@ import "./wrappers/MixinEpochWrapper.sol";
 import "./wrappers/MixinRewardsWrapper.sol";
 import "./wrappers/MixinFeesWrapper.sol";
 import "./wrappers/MixinExchangeWrapper.sol";
+import "./wrappers/MixinZrxVaultWrapper.sol";
+import "./wrappers/MixinRewardVaultWrapper.sol";
 
 
 contract Staking is
     MixinExchangeWrapper,
+    MixinZrxVaultWrapper,
+    MixinRewardVaultWrapper,
     MixinEpochWrapper,
     MixinRewardsWrapper,
     MixinStakeBalancesWrapper,
@@ -37,18 +41,9 @@ contract Staking is
     MixinFeesWrapper
 {
 
-    function setZrxVault(address _zrxVault)
+    // this contract can receive ETH
+    function ()
         external
-    {
-        zrxVault = IVault(_zrxVault);
-    }
-
-    function setRewardVault(address payable _rewardVault)
-        external
-    {
-        rewardVault = IRewardVault(_rewardVault);
-    }
-
-    ///// CAN RECEIVE FUNDS /////
-    function () external payable {}
+        payable
+    {}
 }

contracts/staking/contracts/src/core/MixinFees.sol

@@ -27,6 +27,7 @@ import "../interfaces/IStakingEvents.sol";
 import "./MixinStakeBalances.sol";
 import "./MixinEpoch.sol";
 import "./MixinPools.sol";
+import "./MixinRewardVault.sol";
 import "../interfaces/IStructs.sol";
 import "../libs/LibMath.sol";
 
@@ -37,6 +38,7 @@ contract MixinFees is
     IStructs,
     MixinConstants,
     MixinStorage,
+    MixinRewardVault,
     MixinEpoch,
     MixinStakeBalances,
     MixinPools
@@ -137,7 +139,7 @@ contract MixinFees is
             );
 
             // record reward in vault
-            rewardVault.recordDepositFor(activePoolIds[i].poolId, reward);
+            _recordDepositInRewardVault(activePoolIds[i].poolId, reward);
             totalRewardsRecordedInVault = _safeAdd(totalRewardsRecordedInVault, reward);
 
             // clear state for refunds
@@ -152,8 +154,7 @@ contract MixinFees is
             "MISCALCULATED_REWARDS"
         );
         if (totalRewardsRecordedInVault > 0) {
-            address payable rewardVaultAddress = address(uint160(address(rewardVault)));
-            rewardVaultAddress.transfer(totalRewardsRecordedInVault);
+            _depositIntoRewardVault(totalRewardsRecordedInVault);
         }
 
         // Notify finalization

contracts/staking/contracts/src/core/MixinPools.sol

@@ -22,12 +22,14 @@ import "../immutable/MixinStorage.sol";
 import "@0x/contracts-utils/contracts/src/SafeMath.sol";
 import "../immutable/MixinConstants.sol";
 import "../interfaces/IStakingEvents.sol";
+import "./MixinRewardVault.sol";
 
 contract MixinPools is
     SafeMath,
     IStakingEvents,
     MixinConstants,
-    MixinStorage
+    MixinStorage,
+    MixinRewardVault
 {
 
     function _getNextPoolId()
@@ -53,10 +55,7 @@ contract MixinPools is
         poolById[poolId] = pool;
 
         // create pool in reward vault
-        rewardVault.createPool(
-            poolId,
-            operatorShare
-        );
+        _createPoolInRewardVault(poolId, operatorShare);
 
         // 
         emit PoolCreated(poolId, operatorAddress, operatorShare);

contracts/staking/contracts/src/core/MixinRewardVault.sol

@@ -0,0 +1,100 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../interfaces/IRewardVault.sol";
+import "../immutable/MixinStorage.sol";
+
+
+contract MixinRewardVault is
+    MixinStorage
+{
+
+    function _setRewardVault(address payable _rewardVault)
+        internal
+    {
+        rewardVault = IRewardVault(_rewardVault);
+    }
+
+    function _getRewardVault()
+        internal
+        view
+        returns (address)
+    {
+        return address(rewardVault);
+    }
+
+    function _createPoolInRewardVault(bytes32 poolId, uint8 operatorShare)
+        internal
+    {
+        rewardVault.createPool(
+            poolId,
+            operatorShare
+        );
+    }
+
+     function _balanceInRewardVault(bytes32 poolId)
+        internal
+        view
+        returns (uint256)
+    {
+        return rewardVault.balanceOf(poolId);
+    }
+
+    function _balanceOfOperatorInRewardVault(bytes32 poolId)
+        internal
+        view
+        returns (uint256)
+    {
+        return rewardVault.balanceOfOperator(poolId);
+    }
+
+    function _balanceOfPoolInRewardVault(bytes32 poolId)
+        internal
+        view
+        returns (uint256)
+    {
+        return rewardVault.balanceOfPool(poolId);
+    }
+
+    function _withdrawFromPoolInRewardVault(bytes32 poolId, uint256 amount)
+        internal
+    {
+        rewardVault.withdrawFromPool(poolId, amount);
+    }
+
+    function _withdrawFromOperatorInRewardVault(bytes32 poolId, uint256 amount)
+        internal
+    {
+        rewardVault.withdrawFromOperator(poolId, amount);
+    }
+
+    function _depositIntoRewardVault(uint256 amountInWei)
+        internal
+    {
+        address payable rewardVaultAddress = address(uint160(address(rewardVault)));
+        rewardVaultAddress.transfer(amountInWei);
+    }
+
+    function _recordDepositInRewardVault(bytes32 poolId, uint256 amount)
+        internal
+    {
+        rewardVault.recordDepositFor(poolId, amount);
+    }
+}

contracts/staking/contracts/src/core/MixinRewards.sol

@@ -24,6 +24,7 @@ import "@0x/contracts-utils/contracts/src/SafeMath.sol";
 import "../immutable/MixinConstants.sol";
 import "../interfaces/IStakingEvents.sol";
 import "./MixinStakeBalances.sol";
+import "./MixinRewardVault.sol";
 
 contract MixinRewards is
     SafeMath,
@@ -31,6 +32,7 @@ contract MixinRewards is
     IStakingEvents,
     MixinConstants,
     MixinStorage,
+    MixinRewardVault,
     MixinStakeBalances
 {
     // Pinciple - design any Mixin such that internal members are callable without messing up internal state
@@ -38,36 +40,12 @@ contract MixinRewards is
 
     // @TODO -- add a MixinZrxVault and a MixinRewardVault that interact with the vaults.
 
-     function _getRewardBalance(bytes32 poolId)
-        internal
-        view
-        returns (uint256)
-    {
-        return rewardVault.balanceOf(poolId);
-    }
-
-    function _getRewardBalanceOfOperator(bytes32 poolId)
-        internal
-        view
-        returns (uint256)
-    {
-        return rewardVault.balanceOfOperator(poolId);
-    }
-
-    function _getRewardBalanceOfPool(bytes32 poolId)
-        internal
-        view
-        returns (uint256)
-    {
-        return rewardVault.balanceOfPool(poolId);
-    }
-
     function _computeRewardBalance(bytes32 poolId, address owner)
         internal
         view
         returns (uint256)
     {
-        uint256 poolBalance = rewardVault.balanceOfPool(poolId);
+        uint256 poolBalance = _balanceOfPoolInRewardVault(poolId);
         return _computePayoutDenominatedInRealAsset(
             delegatedStakeToPoolByOwner[owner][poolId],
             delegatedStakeByPoolId[poolId],
@@ -96,7 +74,7 @@ contract MixinRewards is
     function _withdrawOperatorReward(bytes32 poolId, uint256 amount)
         internal
     {
-        rewardVault.withdrawFromOperator(poolId, amount);
+        _withdrawFromOperatorInRewardVault(poolId, amount);
         poolById[poolId].operatorAddress.transfer(amount);
     }
 
@@ -112,7 +90,7 @@ contract MixinRewards is
         shadowRewardsInPoolByOwner[owner][poolId] = _safeAdd(shadowRewardsInPoolByOwner[owner][poolId], amount);
         shadowRewardsByPoolId[poolId] = _safeAdd(shadowRewardsByPoolId[poolId], amount);
 
-        rewardVault.withdrawFromPool(poolId, amount);
+        _withdrawFromPoolInRewardVault(poolId, amount);
         owner.transfer(amount);
     }
 
@@ -120,8 +98,8 @@ contract MixinRewards is
         internal
         returns (uint256)
     {
-        uint256 amount = rewardVault.balanceOfOperator(poolId);
-        rewardVault.withdrawFromOperator(poolId, amount);
+        uint256 amount = _balanceOfOperatorInRewardVault(poolId);
+        _withdrawFromOperatorInRewardVault(poolId, amount);
         poolById[poolId].operatorAddress.transfer(amount);
 
         return amount;
@@ -136,7 +114,7 @@ contract MixinRewards is
         shadowRewardsInPoolByOwner[owner][poolId] = _safeAdd(shadowRewardsInPoolByOwner[owner][poolId], amount);
         shadowRewardsByPoolId[poolId] = _safeAdd(shadowRewardsByPoolId[poolId], amount);
 
-        rewardVault.withdrawFromPool(poolId, amount);
+        _withdrawFromPoolInRewardVault(poolId, amount);
         owner.transfer(amount);
 
         return amount;

contracts/staking/contracts/src/core/MixinStake.sol

@@ -18,7 +18,6 @@
 
 pragma solidity ^0.5.5;
 
-import "../interfaces/IVault.sol";
 import "../libs/LibZrxToken.sol";
 import "@0x/contracts-utils/contracts/src/SafeMath.sol";
 import "../immutable/MixinStorage.sol";
@@ -26,6 +25,8 @@ import "../immutable/MixinConstants.sol";
 import "../interfaces/IStakingEvents.sol";
 import "./MixinStakeBalances.sol";
 import "./MixinEpoch.sol";
+import "./MixinZrxVault.sol";
+import "./MixinRewardVault.sol";
 import "../libs/LibRewards.sol";
 
 
@@ -35,6 +36,8 @@ contract MixinStake is
     IStakingEvents,
     MixinConstants,
     MixinStorage,
+    MixinZrxVault,
+    MixinRewardVault,
     MixinEpoch,
     MixinStakeBalances
 {
@@ -126,7 +129,7 @@ contract MixinStake is
         private
     {
         // deposit equivalent amount of ZRX into vault
-        zrxVault.depositFrom(owner, amount);
+        _depositFromOwnerIntoZrxVault(owner, amount);
 
         // mint stake
         stakeByOwner[owner] = _safeAdd(stakeByOwner[owner], amount);
@@ -145,7 +148,7 @@ contract MixinStake is
         stakeByOwner[owner] = _safeSub(stakeByOwner[owner], amount);
 
         // withdraw equivalent amount of ZRX from vault
-        zrxVault.withdrawFrom(owner, amount);
+        _withdrawToOwnerFromZrxVault(owner, amount);
 
         // emit stake event
         emit StakeBurned(
@@ -173,7 +176,7 @@ contract MixinStake is
 
         // update delegator's share of reward pool
         // note that this uses the snapshot parameters
-        uint256 poolBalance = rewardVault.balanceOfPool(poolId);
+        uint256 poolBalance = _balanceOfPoolInRewardVault(poolId);
         uint256 buyIn = _computeBuyInDenominatedInShadowAsset(
             amount,
             _delegatedStakeByPoolId,
@@ -215,7 +218,7 @@ contract MixinStake is
 
         // get payout
         // TODO -- not full balance, just balance that belongs to delegators.
-        uint256 poolBalance = rewardVault.balanceOfPool(poolId);
+        uint256 poolBalance = _balanceOfPoolInRewardVault(poolId);
         uint256 payoutInRealAsset;
         uint256 payoutInShadowAsset;
         if (_delegatedStakeToPoolByOwner == amount) {
@@ -244,7 +247,7 @@ contract MixinStake is
 
         // withdraw payout for delegator
         if (payoutInRealAsset > 0) {
-            rewardVault.withdrawFromPool(poolId, payoutInRealAsset);
+            _withdrawFromPoolInRewardVault(poolId, payoutInRealAsset);
             owner.transfer(payoutInRealAsset);
         }
     }

contracts/staking/contracts/src/core/MixinZrxVault.sol

@@ -0,0 +1,54 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../interfaces/IVault.sol";
+import "../immutable/MixinStorage.sol";
+
+
+contract MixinZrxVault is
+    MixinStorage
+{
+
+    function _setZrxVault(address _zrxVault)
+        internal
+    {
+        zrxVault = IVault(_zrxVault);
+    }
+
+    function _getZrxVault()
+        internal
+        view
+        returns (address)
+    {
+        return address(zrxVault);
+    }
+
+    function _depositFromOwnerIntoZrxVault(address owner, uint256 amount)
+        internal
+    {
+        zrxVault.depositFrom(owner, amount);
+    }
+
+    function _withdrawToOwnerFromZrxVault(address owner, uint256 amount)
+        internal
+    {
+        zrxVault.withdrawFrom(owner, amount);
+    }
+}

contracts/staking/contracts/src/interfaces/IRewardVault.sol

@@ -21,8 +21,6 @@ pragma solidity ^0.5.5;
 
 interface IRewardVault {
 
-
-
     function depositFor(bytes32 poolId)
         external
         payable;

contracts/staking/contracts/src/wrappers/MixinRewardVaultWrapper.sol

@@ -0,0 +1,67 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinRewardVault.sol";
+
+
+contract MixinRewardVaultWrapper is
+    MixinRewardVault
+{
+
+    // NOTE that some names differ slightly from the internal function names for simplicity of the API to outside clients.
+
+    function setRewardVault(address payable _rewardVault)
+        external
+    {
+        _setRewardVault(_rewardVault);
+    }
+
+    function getRewardVault()
+        external
+        view
+        returns (address)
+    {
+        return _getRewardVault();
+    }
+
+    function getRewardBalance(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _balanceInRewardVault(poolId);
+    }
+
+    function getRewardBalanceOfOperator(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _balanceOfOperatorInRewardVault(poolId);
+    }
+
+    function getRewardBalanceOfPool(bytes32 poolId)
+        external
+        view
+        returns (uint256)
+    {
+        return _balanceOfPoolInRewardVault(poolId);
+    }
+}

contracts/staking/contracts/src/wrappers/MixinRewardsWrapper.sol

@@ -27,30 +27,6 @@ contract MixinRewardsWrapper is
     Modifiers
 {
 
-    function getRewardBalance(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getRewardBalance(poolId);
-    }
-
-    function getRewardBalanceOfOperator(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getRewardBalanceOfOperator(poolId);
-    }
-
-    function getRewardBalanceOfPool(bytes32 poolId)
-        external
-        view
-        returns (uint256)
-    {
-        return _getRewardBalanceOfPool(poolId);
-    }
-
      function computeRewardBalance(bytes32 poolId, address owner)
         external
         view

contracts/staking/contracts/src/wrappers/MixinZrxVaultWrapper.sol

@@ -0,0 +1,41 @@
+/*
+
+  Copyright 2018 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.5.5;
+
+import "../core/MixinZrxVault.sol";
+
+
+contract MixinZrxVaultWrapper is
+    MixinZrxVault
+{
+
+    function setZrxVault(address _zrxVault)
+        external
+    {
+        _setZrxVault(_zrxVault);
+    }
+
+    function getZrxVault()
+        external
+        view
+        returns (address)
+    {
+        return _getZrxVault();
+    }
+}