checkpoint - implementing delegating + timelocks

contracts/staking/contracts/src/Staking.sol

@@ -57,10 +57,10 @@ contract Staking is
         _activateStake(msg.sender, amount);
     }
 
-    function activateAndDelegateStake(address owner, uint256 amount)
+    function activateAndDelegateStake(address owner, bytes32 poolId, uint256 amount)
         external
     {
-        _activateAndDelegateStake(msg.sender, amount);
+        _activateAndDelegateStake(msg.sender, poolId, amount);
     }
 
     function deactivateAndTimelockStake(address owner, uint256 amount)
@@ -112,7 +112,7 @@ contract Staking is
         view
         returns (uint256)
     {
-        return getWithdrawableStake(owner);
+        return _getWithdrawableStake(owner);
     }
 
     function getTimelockedStake(address owner)

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

@@ -24,7 +24,8 @@ import "@0x/contracts-utils/contracts/src/SafeMath.sol";
 import "./MixinStorage.sol";
 import "./MixinConstants.sol";
 import "../interfaces/IStakingEvents.sol";
-import "./StakingBalances.sol";
+import "./MixinStakeBalances.sol";
+import "./MixinEpoch.sol";
 
 
 contract MixinStake is
@@ -32,7 +33,8 @@ contract MixinStake is
     IStakingEvents,
     MixinConstants,
     MixinStorage,
-    MixinStakeBalances
+    MixinStakeBalances,
+    MixinEpoch
 {
     using LibZrxToken for uint256;
 
@@ -59,20 +61,19 @@ contract MixinStake is
     function _activateStake(address owner, uint256 amount)
         internal
     {
-        _syncTimelockedStake();
+        _syncTimelockedStake(owner);
-
-        Timelock memory ownerTimelock = timelocksByOwner[owner];
         require(
-           ownerTimelock.availableBalance >= amount,
+            _getDeactivatedStake(owner) >= amount,
-           "INSUFFICIENT_BALANCE"
+            "INSUFFICIENT_BALANCE"
         );
-        ownerTimelock.sub(amount);
+        activeStakeByOwner[owner] = _safeAdd(activeStakeByOwner[owner], amount);
-        timelocksByOwner[owner] = ownerTimelock;
-
-        _activateStake(owner, amount);
     }
 
-    function _activateAndDelegateStake(address owner, bytes32 poolId, uint256 amount)
+    function _activateAndDelegateStake(
+        address owner,
+        bytes32 poolId,
+        uint256 amount
+    )
         internal
     {
         _activateStake(owner, amount);
@@ -82,32 +83,29 @@ contract MixinStake is
     function _deactivateAndTimelockStake(address owner, uint256 amount)
         internal
     {
-        Timelock memory ownerTimelock = timelocksByOwner[owner];
+        _syncTimelockedStake(owner);
-        ownerTimelock.add(amount);
+        require(
-        timelocksByOwner[owner] = ownerTimelock;
+            _getActivatedStake(owner) >= amount,
-
+            "INSUFFICIENT_BALANCE"
+        );
         activeStakeByOwner[owner] = _safeSub(activeStakeByOwner[owner], amount);
     }
 
     function _deactivateAndTimelockDelegatedStake(address owner, bytes32 poolId, uint256 amount)
         internal
     {
-        _deactivateStake(owner, amount);
+        _deactivateAndTimelockStake(owner, amount);
         _undelegateStake(owner, poolId, amount);
     }
 
     function _withdraw(address owner, uint256 amount)
         internal
     {
-        Timelock memory ownerTimelock = timelocksByOwner[owner];
+         _syncTimelockedStake(owner);
         require(
-           ownerTimelock.availableBalance >= amount,
+            _getDeactivatedStake(owner) >= amount,
-           "INSUFFICIENT_BALANCE"
+            "INSUFFICIENT_BALANCE"
         );
-        ownerTimelock.sub(amount);
-        timelocksByOwner[owner] = ownerTimelock;
-
-        // burn stake
         _burnStake(owner, amount);
     }
 
@@ -145,12 +143,6 @@ contract MixinStake is
         );
     }
 
-    function _activateStake(address owner, uint256 amount)
-        private
-    {
-        activeStakeByOwner[owner] = _safeAdd(activeStakeByOwner[owner], amount);
-    }
-
     function _delegateStake(address owner, bytes32 poolId, uint256 amount)
         private
     {
@@ -177,7 +169,7 @@ contract MixinStake is
         delegatedStakeByPoolId[poolId] = _safeSub(delegatedStakeByPoolId[poolId], amount);
     }
 
-    // Epoch | lockedAt  | total | pending | current | timelock() | withdraw() | available()
+    // Epoch | lockedAt  | total | pending |  | timelock() | withdraw() | available()
     // 0     | 0         | 0     | 0       | 0       |            |            | 0
     // 1     | 1         | 5     | 0       | 0       | +5         |            | 0
     // 2     | 1         | 5     | 0       | 0       |            |            | 0
@@ -194,34 +186,46 @@ contract MixinStake is
     function _timelockStake(address owner, uint256 amount)
         private
     {
-        Timelock memory ownerTimelock = _getSynchronizedTimelock(owner, amount);
+        (Timelock memory ownerTimelock,) = _getSynchronizedTimelock(owner);
-        ownerTimelock.total = _safeAdd(ownerTimelock.total, amount);
+        require(
-         timelocksByOwner[owner] = ownerTimelock;
+            amount <= 2**96 - 1,
+            "AMOUNT_TOO_LARGE"
+        );
+        uint96 downcastAmount = uint96(amount);
+        ownerTimelock.total += downcastAmount;
+        timelockedStakeByOwner[owner] = ownerTimelock;
     }
 
-    function _syncTimelockedStake(address owner, uint256 amount)
+    function _syncTimelockedStake(address owner)
         private
     {
-        timelocksByOwner[owner] = _getSynchronizedTimelock(owner, amount);
+        (Timelock memory ownerTimelock, bool isOutOfSync) = _getSynchronizedTimelock(owner);
-    }
+        if (!isOutOfSync) {
-
+            return;
-    function _getSynchronizedTimelock(address owner, uint256 amount)
-        private
-        returns (Timelock memory ownerTimelock)
-    {
-        Timelock memory ownerTimelock = timelocksByOwner[owner];
-        uint64 currentTimelockPeriod = getCurrentTimelockPeriod();
-        if (currentTimelockPeriod == _safeAdd(ownerTimelock.lockedAt, 1)) {
-            // shift one period
-            ownerTimelock.current = ownerTimelock.pending;
-            ownerTimelock.pending = ownerTimelock.total;
-        } else if (currentTimelockPeriod > ownerTimelock.lockedAt) {
-            // shift n periods
-            ownerTimelock.current = ownerTimelock.total;
-            ownerTimelock.pending = ownerTimelock.total;
-        } else {
-            // do nothing
         }
-        return ownerTimelock;
+        timelockedStakeByOwner[owner] = ownerTimelock;
+    }
+
+    function _getSynchronizedTimelock(address owner)
+        private
+        returns (
+            Timelock memory ownerTimelock,
+            bool isOutOfSync
+        )
+    {
+        uint64 currentTimelockPeriod = _getCurrentTimelockPeriod();
+        ownerTimelock = timelockedStakeByOwner[owner];
+        isOutOfSync = false;
+        if (currentTimelockPeriod == _safeAdd(ownerTimelock.lockedAt, 1)) {
+            // shift n periods
+            ownerTimelock.pending = ownerTimelock.total;
+            isOutOfSync = true;
+        } else if(currentTimelockPeriod > ownerTimelock.lockedAt) {
+            // Timelock has expired - zero out
+            ownerTimelock.lockedAt = 0;
+            ownerTimelock.total = 0;
+            ownerTimelock.pending = 0;
+        }
+        return (ownerTimelock, isOutOfSync);
     }
 }

contracts/staking/contracts/src/core/MixinStakeBalances.sol

@@ -25,12 +25,12 @@ import "./MixinStorage.sol";
 import "./MixinConstants.sol";
 
 
-contract MixinStake is
+contract MixinStakeBalances is
     SafeMath,
     MixinConstants,
-    MixinStorage,
+    MixinStorage
 {
-    
+
     function _getTotalStake(address owner)
         internal
         view
@@ -39,7 +39,7 @@ contract MixinStake is
         return stakeByOwner[owner];
     }
 
-    function getActivatedStake(address owner)
+    function _getActivatedStake(address owner)
         internal
         view
         returns (uint256)
@@ -47,17 +47,25 @@ contract MixinStake is
         return activeStakeByOwner[owner];
     }
 
-    function getDeactivatedStake(address owner)
+    function _getDeactivatedStake(address owner)
         internal
         view
         returns (uint256)
     {
-        return _safeSub(_getTotalStake(owner), getActivatedStake(owner));
+        return _safeSub(_getTotalStake(owner), _getActivatedStake(owner));
     }
 
-    function getStakeAvailableForActivation()
+    /*
+    function _getStakeAvailableForActivation()
+    internal
+        view
+        returns (uint256)
+    {
 
-    function getWithdrawableStake(address owner)
+    }
+    */
+
+    function _getWithdrawableStake(address owner)
         internal
         view
         returns (uint256)
@@ -65,7 +73,7 @@ contract MixinStake is
         
     }
 
-    function getTimelockedStake(address owner)
+    function _getTimelockedStake(address owner)
         internal
         view
         returns (uint256)
@@ -73,7 +81,7 @@ contract MixinStake is
         return timelockedStakeByOwner[owner].total;
     }
 
-    function getStakeDelegatedByOwner(address owner)
+    function _getStakeDelegatedByOwner(address owner)
         internal
         view
         returns (uint256)
@@ -81,7 +89,7 @@ contract MixinStake is
         return delegatedStakeByOwner[owner];
     }
 
-    function getStakeDelegatedToPoolByOwner(address owner, bytes32 poolId)
+    function _getStakeDelegatedToPoolByOwner(address owner, bytes32 poolId)
         internal
         view
         returns (uint256)
@@ -89,7 +97,7 @@ contract MixinStake is
         return delegatedStakeToPoolByOwner[owner][poolId];
     }
 
-    function getStakeDelegatedToPool(bytes32 poolId)
+    function _getStakeDelegatedToPool(bytes32 poolId)
         internal
         view
         returns (uint256)

contracts/staking/contracts/src/core/MixinStorage.sol

@@ -40,15 +40,15 @@ contract MixinStorage is
     // mapping from Owner to Amount Timelocked
     mapping (address => Timelock) timelockedStakeByOwner;
 
-    // mapping from Pool Id to Amount Delegated
-    mapping (bytes32 => uint256) delegatedStakeByPoolId;
-
     // mapping from Owner to Amount Delegated
     mapping (address => uint256) delegatedStakeByOwner;
 
     // mapping from Owner to Pool Id to Amount Delegated
     mapping (address => mapping (bytes32 => uint256)) delegatedStakeToPoolByOwner;
 
+    // mapping from Pool Id to Amount Delegated
+    mapping (bytes32 => uint256) delegatedStakeByPoolId;
+
     // tracking Pool Id
     bytes32 nextPoolId = INITIAL_POOL_ID;
 

contracts/staking/contracts/src/core/MixinTimelock.sol

@@ -1,49 +0,0 @@
-/*
-
-  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 "@0x/contracts-utils/contracts/src/SafeMath.sol";
-import "./MixinStorage.sol";
-
-
-contract MixinStake is
-    SafeMath,
-    MixinStorage
-{
-
-
-contract MixinTimelock {
-
-    // Epoch | lockedAt  | total | pending | current | timelock() | withdraw() | available()
-    // 0     | 0         | 0     | 0       | 0       |            |            | 0
-    // 1     | 1         | 5     | 0       | 0       | +5         |            | 0
-    // 2     | 1         | 5     | 0       | 0       |            |            | 0
-    // 2     | 2         | 15    | 5       | 0       | +10        |            | 0
-    // 3     | 2         | 15    | 5       | 0       |            |            | 5
-    // 3     | 3         | 30    | 15      | 5       | +15        |            | 5
-    // 4     | 3         | 30    | 15      | 5       |            |            | 15
-    // 5     | 3         | 30    | 15      | 5       |            |            | 30
-    // 5     | 5         | 30    | 30      | 30      | +0 *       |            | 30
-    // 6     | 6         | 50    | 30      | 30      | +20        |            | 30
-    // 6     | 6         | 20    | 0       | 0       |            | -30        | 0
-    // 7     | 6         | 20    | 0       | 0       |            |            | 0
-    // 8     | 6         | 20    | 0       | 0       |            |            | 20
-
-    function _timelockStake()
-}

contracts/staking/contracts/src/interfaces/IStructs.sol

@@ -0,0 +1,34 @@
+/*
+
+  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;
+
+
+interface IStructs {
+
+    struct Timelock {
+        uint64 lockedAt;
+        uint96 total;
+        uint96 pending;
+    }
+
+    struct Pool {
+        address operatorAddress;
+        uint8 operatorShare;
+    }
+}