Stake locks extracted out of TokenStaking into a library

solidity/contracts/TokenStaking.sol

@@ -20,8 +20,8 @@ import "openzeppelin-solidity/contracts/math/SafeMath.sol";
 import "./StakeDelegatable.sol";
 import "./libraries/staking/MinimumStakeSchedule.sol";
 import "./libraries/staking/GrantStaking.sol";
+import "./libraries/staking/Locks.sol";
 import "./utils/PercentUtils.sol";
-import "./utils/LockUtils.sol";
 import "./utils/BytesLib.sol";
 import "./Authorizations.sol";
 import "./TokenStakingEscrow.sol";
@@ -39,6 +39,7 @@ contract TokenStaking is Authorizations, StakeDelegatable {
     using LockUtils for LockUtils.LockSet;
     using SafeERC20 for ERC20Burnable;
     using GrantStaking for GrantStaking.Storage;
+    using Locks for Locks.Storage;
 
     event Staked(
         address owner,
@@ -55,8 +56,6 @@ contract TokenStaking is Authorizations, StakeDelegatable {
     event LockReleased(address indexed operator, address lockCreator);
     event ExpiredLockReleased(address indexed operator, address lockCreator);
 
-    uint256 public constant maximumLockDuration = 86400 * 200; // 200 days in seconds
-
     uint256 public initializationPeriod;
     uint256 public undelegationPeriod;
 
@@ -68,14 +67,11 @@ contract TokenStaking is Authorizations, StakeDelegatable {
 
     GrantStaking.Storage internal grantStaking;
 
+    Locks.Storage internal locks;
+
     // KEEP token grant contract.
     TokenGrant internal tokenGrant;
 
-    // Locks placed on the operator.
-    // `operatorLocks[operator]` returns all locks placed on the operator.
-    // Each authorized operator contract can place one lock on an operator.
-    mapping(address => LockUtils.LockSet) internal operatorLocks;
-
     /// @notice Creates a token staking contract for a provided Standard ERC20Burnable token.
     /// @param _token KEEP token contract.
     /// @param _tokenGrant KEEP token grant contract.
@@ -281,7 +277,6 @@ contract TokenStaking is Authorizations, StakeDelegatable {
             isAuthorizedForOperator(operator, msg.sender),
             "Not authorized"
         );
-        require(duration <= maximumLockDuration, "Lock duration too long");
 
         uint256 operatorParams = operators[operator].packedParams;
 
@@ -294,11 +289,7 @@ contract TokenStaking is Authorizations, StakeDelegatable {
             "Operator undelegating"
         );
 
-        operatorLocks[operator].setLock(
-            msg.sender,
-            uint96(block.timestamp.add(duration))
-        );
-        emit StakeLocked(operator, msg.sender, block.timestamp.add(duration));
+        locks.lockStake(operator, duration);
     }
 
     /// @notice Removes a lock the caller had previously placed on the operator.
@@ -318,8 +309,7 @@ contract TokenStaking is Authorizations, StakeDelegatable {
             isAuthorizedForOperator(operator, msg.sender),
             "Not authorized"
         );
-        operatorLocks[operator].releaseLock(msg.sender);
-        emit LockReleased(operator, msg.sender);
+        locks.releaseLock(operator);
     }
 
     /// @notice Removes the lock of the specified operator contract
@@ -332,38 +322,14 @@ contract TokenStaking is Authorizations, StakeDelegatable {
         address operator,
         address operatorContract
     ) public {
-        LockUtils.LockSet storage locks = operatorLocks[operator];
-        require(
-            locks.contains(operatorContract),
-            "No matching lock present"
-        );
-        bool expired = block.timestamp >= locks.getLockTime(operatorContract);
-        bool disabled = !registry.isApprovedOperatorContract(operatorContract);
-        require(
-            expired || disabled,
-            "Lock still active and valid"
-        );
-        locks.releaseLock(operatorContract);
-        emit ExpiredLockReleased(operator, operatorContract);
+        locks.releaseExpiredLock(operator, operatorContract, registry);
     }
 
     /// @notice Check whether the operator has any active locks
     /// that haven't expired yet
     /// and whose creators aren't disabled by the panic button.
-    function isStakeLocked(
-        address operator
-    ) public view returns (bool) {
-        LockUtils.Lock[] storage _locks = operatorLocks[operator].locks;
-        LockUtils.Lock memory lock;
-        for (uint i = 0; i < _locks.length; i++) {
-            lock = _locks[i];
-            if (block.timestamp < lock.expiresAt) {
-                if (registry.isApprovedOperatorContract(lock.creator)) {
-                    return true;
-                }
-            }
-        }
-        return false;
+    function isStakeLocked(address operator) public view returns (bool) {
+        return locks.isStakeLocked(operator, registry);
     }
 
     /// @notice Get the locks placed on the operator.
@@ -375,15 +341,7 @@ contract TokenStaking is Authorizations, StakeDelegatable {
         public
         view
         returns (address[] memory creators, uint256[] memory expirations) {
-        uint256 lockCount = operatorLocks[operator].locks.length;
-        creators = new address[](lockCount);
-        expirations = new uint256[](lockCount);
-        LockUtils.Lock memory lock;
-        for (uint i = 0; i < lockCount; i++) {
-            lock = operatorLocks[operator].locks[i];
-            creators[i] = lock.creator;
-            expirations[i] = lock.expiresAt;
-        }
+        return locks.getLocks(operator);
     }
 
     /// @notice Slash provided token amount from every member in the misbehaved
@@ -609,10 +567,7 @@ contract TokenStaking is Authorizations, StakeDelegatable {
             return false;
         }
         // Undelegating finished, so check locks
-        LockUtils.LockSet storage locks = operatorLocks[_operator];
-        // `getLockTime` returns 0 if the lock doesn't exist,
-        // thus we don't need to check for its presence separately.
-        return block.timestamp >= locks.getLockTime(_operatorContract);
+        return locks.isStakeReleased(_operator, _operatorContract);
     }
 
     function transferOrDeposit(

solidity/contracts/libraries/staking/Locks.sol

@@ -0,0 +1,109 @@
+pragma solidity 0.5.17;
+
+import "openzeppelin-solidity/contracts/math/SafeMath.sol";
+import "../../KeepRegistry.sol";
+import "./LockUtils.sol";
+
+library Locks {
+    using SafeMath for uint256;
+    using LockUtils for LockUtils.LockSet;
+
+    event StakeLocked(address indexed operator, address lockCreator, uint256 until);
+    event LockReleased(address indexed operator, address lockCreator);
+    event ExpiredLockReleased(address indexed operator, address lockCreator);
+
+    uint256 public constant maximumLockDuration = 86400 * 200; // 200 days in seconds
+
+    struct Storage {
+        // Locks placed on the operator.
+        // `operatorLocks[operator]` returns all locks placed on the operator.
+        // Each authorized operator contract can place one lock on an operator.
+        mapping(address => LockUtils.LockSet) operatorLocks;
+    }
+
+    function lockStake(
+        Storage storage self,
+        address operator,
+        uint256 duration
+    ) public {
+        require(duration <= maximumLockDuration, "Lock duration too long");
+        self.operatorLocks[operator].setLock(
+            msg.sender,
+            uint96(block.timestamp.add(duration))
+        );
+        emit StakeLocked(operator, msg.sender, block.timestamp.add(duration));
+    }
+
+    function releaseLock(
+        Storage storage self,
+        address operator
+    ) public {
+        self.operatorLocks[operator].releaseLock(msg.sender);
+        emit LockReleased(operator, msg.sender);
+    }
+
+    function releaseExpiredLock(
+        Storage storage self,
+        address operator,
+        address operatorContract,
+        KeepRegistry registry
+    ) public {
+        LockUtils.LockSet storage locks = self.operatorLocks[operator];
+        require(
+            locks.contains(operatorContract),
+            "No matching lock present"
+        );
+        bool expired = block.timestamp >= locks.getLockTime(operatorContract);
+        bool disabled = !registry.isApprovedOperatorContract(operatorContract);
+        require(
+            expired || disabled,
+            "Lock still active and valid"
+        );
+        locks.releaseLock(operatorContract);
+        emit ExpiredLockReleased(operator, operatorContract);
+    }
+
+    function isStakeLocked(
+        Storage storage self,
+        address operator,
+        KeepRegistry registry
+    ) public view returns (bool) {
+        LockUtils.Lock[] storage _locks = self.operatorLocks[operator].locks;
+        LockUtils.Lock memory lock;
+        for (uint i = 0; i < _locks.length; i++) {
+            lock = _locks[i];
+            if (block.timestamp < lock.expiresAt) {
+                if (registry.isApprovedOperatorContract(lock.creator)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    function isStakeReleased(
+        Storage storage self,
+        address operator,
+        address operatorContract
+    ) public view returns (bool) {
+        LockUtils.LockSet storage locks = self.operatorLocks[operator];
+        // `getLockTime` returns 0 if the lock doesn't exist,
+        // thus we don't need to check for its presence separately.
+        return block.timestamp >= locks.getLockTime(operatorContract);
+    }
+
+    function getLocks(
+        Storage storage self,
+        address operator
+    ) public view returns (address[] memory creators, uint256[] memory expirations) {
+        uint256 lockCount = self.operatorLocks[operator].locks.length;
+        creators = new address[](lockCount);
+        expirations = new uint256[](lockCount);
+        LockUtils.Lock memory lock;
+        for (uint i = 0; i < lockCount; i++) {
+            lock = self.operatorLocks[operator].locks[i];
+            creators[i] = lock.creator;
+            expirations[i] = lock.expiresAt;
+        }
+    }
+} 

solidity/contracts/stubs/LockStub.sol

@@ -1,6 +1,6 @@
 pragma solidity 0.5.17;
 
-import "../utils/LockUtils.sol";
+import "../libraries/staking/LockUtils.sol";
 
 contract LockStub {
     using LockUtils for LockUtils.LockSet;

solidity/migrations/2_deploy_contracts.js

@@ -4,6 +4,7 @@ const AltBn128 = artifacts.require("./cryptography/AltBn128.sol");
 const BLS = artifacts.require("./cryptography/BLS.sol");
 const MinimumStakeSchedule = artifacts.require("./MinimumStakeSchedule.sol");
 const GrantStaking = artifacts.require("./GrantStaking.sol");
+const Locks = artifacts.require("./Locks.sol");
 const TokenStaking = artifacts.require("./TokenStaking.sol");
 const TokenStakingEscrow = artifacts.require("./TokenStakingEscrow.sol");
 const PermissiveStakingPolicy = artifacts.require('./PermissiveStakingPolicy.sol');
@@ -43,8 +44,10 @@ module.exports = async function(deployer, network) {
   await deployer.deploy(TokenStakingEscrow, KeepToken.address, TokenGrant.address);
   await deployer.deploy(MinimumStakeSchedule);
   await deployer.deploy(GrantStaking);
+  await deployer.deploy(Locks);
   await deployer.link(MinimumStakeSchedule, TokenStaking);
   await deployer.link(GrantStaking, TokenStaking);
+  await deployer.link(Locks, TokenStaking);
   await deployer.deploy(
     TokenStaking,
     KeepToken.address,

solidity/test/RolesLookupTest.js

@@ -1,5 +1,6 @@
 const {accounts, contract} = require('@openzeppelin/test-environment')
 const {time, expectRevert} = require('@openzeppelin/test-helpers')
+const {initTokenStaking} = require('./helpers/initContracts')
 const {grantTokens} = require('./helpers/grantTokens')
 const {
   delegateStake,
@@ -10,10 +11,6 @@ const {createSnapshot, restoreSnapshot} = require('./helpers/snapshot')
 const assert = require('chai').assert
 
 const KeepToken = contract.fromArtifact('KeepToken')
-const MinimumStakeSchedule = contract.fromArtifact('MinimumStakeSchedule')
-const TokenStaking = contract.fromArtifact('TokenStaking')
-const GrantStaking = contract.fromArtifact('GrantStaking')
-const TokenStakingEscrow = contract.fromArtifact('TokenStakingEscrow')
 const TokenGrant = contract.fromArtifact('TokenGrant')
 const KeepRegistry = contract.fromArtifact('KeepRegistry')
 const PermissiveStakingPolicy = contract.fromArtifact('PermissiveStakingPolicy')
@@ -47,7 +44,6 @@ describe('RolesLookup', () => {
   let token,
     tokenGrant,
     tokenStaking,
-    tokenStakingEscrow,
     tokenGrantStakingPolicy,
     managedGrantFactory,
     lookup
@@ -56,30 +52,16 @@ describe('RolesLookup', () => {
     const registry = await KeepRegistry.new({from: deployer})
     token = await KeepToken.new({from: deployer})
     tokenGrant = await TokenGrant.new(token.address, {from: deployer})
-    tokenStakingEscrow = await TokenStakingEscrow.new(
-      token.address, 
-      tokenGrant.address, 
-      {from: deployer}
-    )
-    await TokenStaking.detectNetwork()
-    await TokenStaking.link(
-      'MinimumStakeSchedule', 
-      (await MinimumStakeSchedule.new({from: deployer})).address
-    )
-    await TokenStaking.link(
-      'GrantStaking', 
-      (await GrantStaking.new({from: deployer})).address
-    );
-    tokenStaking = await TokenStaking.new(
+    const stakingContracts = await initTokenStaking(
       token.address,
       tokenGrant.address,
-      tokenStakingEscrow.address,
       registry.address,
       initializationPeriod,
       undelegationPeriod,
-      {from: deployer}
+      contract.fromArtifact('TokenStakingEscrow'),
+      contract.fromArtifact('TokenStaking')
     )
-    await tokenStakingEscrow.transferOwnership(tokenStaking.address, {from: deployer})
+    tokenStaking = stakingContracts.tokenStaking
     tokenGrantStakingPolicy = await PermissiveStakingPolicy.new()
     managedGrantFactory = await ManagedGrantFactory.new(
       token.address,