Add stake delegate functionality

contracts/solidity/contracts/StakeDelegatable.sol

@@ -0,0 +1,127 @@
+pragma solidity ^0.4.24;
+
+import "openzeppelin-solidity/contracts/math/SafeMath.sol";
+
+
+/**
+ * @title Stake Delegatable
+ * @dev A contract that allows delegate your stake balance to any address
+ * that is not already a staker. Delegator refers to a staker who has
+ * delegated their stake to another address, while delegate refers to an
+ * address that has had a stake delegated to it.
+ */
+contract StakeDelegatable {
+
+    /**
+     * @dev Only not null addresses can be passed to the functions with this modifier.
+     */
+    modifier notNull(address _address) {
+        require(_address != address(0), "Provided address can not be zero.");
+        _;
+    }
+
+    /**
+     * @dev Only non staker addresses can be passed to the functions with this modifier.
+     */
+    modifier notStaker(address _address) {
+        require(stakeBalances[_address] == 0, "Provided address is not a staker.");
+        _;
+    }
+
+    mapping(address => uint256) public stakeBalances;
+    mapping(address => address) public delegatorToDelegate;
+    mapping(address => address) public delegateToDelegator;
+
+    /**
+     * @dev Gets the stake balance of the specified address.
+     * @param _address The address to query the balance of.
+     * @return The amount staked by the passed address.
+     */
+    function stakeBalanceOf(address _address)
+        public
+        view
+        notNull(_address)
+        returns (uint256)
+    {
+        // If provided address is a delegate return its delegator balance.
+        address delegator = delegatorToDelegate[_address];
+        if (delegator != address(0) && delegateToDelegator[delegator] == _address) {
+            return stakeBalances[delegator];
+        }
+
+        // If provided address is a delegator return zero balance since the 
+        // balance is delegated to a delegate.
+        address delegate = delegateToDelegator[_address];
+        if (delegate != address(0) && delegatorToDelegate[delegate] == _address) {
+            return 0;
+        }
+
+        return stakeBalances[_address];
+    }
+
+    /**
+     * @dev Returns address of a delegate if it exists for the
+     * provided address or the provided address otherwise.
+     * @param _address The address to check.
+     * @return Delegate address or provided address.
+     */
+    function getDelegatorOrDelegate(address _address)
+        public
+        view
+        notNull(_address)
+        returns (address)
+    {
+        address delegate = delegateToDelegator[_address];
+        if (delegate != address(0) && delegatorToDelegate[delegate] == _address) {
+            return delegate;
+        }
+        return _address;
+    }
+
+    /**
+     * @dev Approves address to delegate your stake balance. You can only
+     * have one delegate address. Delegate must also request to operate
+     * your stake by calling requestDelegateFor() method.
+     * @param _address Address to where you want to delegate your balance.
+     */
+    function approveDelegateAt(address _address)
+        public
+        notNull(_address)
+        notStaker(_address)
+    {
+        delegateToDelegator[msg.sender] = _address;
+    }
+
+    /**
+     * @dev Requests to delegate stake for a specified address.
+     * An address must approve you first to delegate by calling
+     * requestDelegateFor() method.
+     * @param _address Address for which you request to delegate.
+     */
+    function requestDelegateFor(address _address)
+        public
+        notNull(_address)
+    {
+        delegatorToDelegate[msg.sender] = _address;
+    }
+
+    /**
+     * @dev Removes delegate for your stake balance.
+     */
+    function removeDelegate() public {
+        address delegate = delegateToDelegator[msg.sender];
+        delete delegatorToDelegate[delegate];
+        delete delegateToDelegator[msg.sender];
+    }
+
+    /**
+     * @dev Revert if a delegate try to stake.
+     * @param _address The address to check.
+     */
+    function revertIfDelegateStakes(address _address) internal {
+        address delegator = delegatorToDelegate[_address];
+        if (delegator != address(0)) {
+            revert("Provided address can not stake since it has stake delegated to it.");
+        }
+    }
+}

contracts/solidity/contracts/TokenGrant.sol

@@ -4,6 +4,7 @@ import "openzeppelin-solidity/contracts/token/ERC20/StandardToken.sol";
 import "openzeppelin-solidity/contracts/token/ERC20/SafeERC20.sol";
 import "openzeppelin-solidity/contracts/math/SafeMath.sol";
 import "./StakingProxy.sol";
+import "./StakeDelegatable.sol";
 
 
 /**
@@ -14,7 +15,7 @@ import "./StakingProxy.sol";
  * released gradually based on the vesting schedule cliff and vesting duration.
  * Optionally grant can be revoked by the token grant creator.
  */
-contract TokenGrant {
+contract TokenGrant is StakeDelegatable {
     using SafeMath for uint256;
     using SafeERC20 for StandardToken;
 
@@ -54,9 +55,6 @@ contract TokenGrant {
     // available to be released to the beneficiary
     mapping(address => uint256) public balances;
 
-    // Token grants stake balances.
-    mapping(address => uint256) public stakeBalances;
-
     // Token grants stake withdrawals.
     mapping(uint256 => uint256) public stakeWithdrawalStart;
 
@@ -82,15 +80,6 @@ contract TokenGrant {
         return balances[_owner];
     }
 
-    /**
-     * @dev Gets the grants stake balance of the specified address.
-     * @param _owner The address to query the grants balance of.
-     * @return An uint256 representing the grants stake balance owned by the passed address.
-     */
-    function stakeBalanceOf(address _owner) public view returns (uint256 balance) {
-        return stakeBalances[_owner];
-    }
-
     /**
      * @dev Gets grant by ID. Returns only basic grant data.
      * If you need vesting schedule for the grant you must call `getGrantVestingSchedule()`
@@ -240,6 +229,8 @@ contract TokenGrant {
         uint256 available = grants[_id].amount.sub(grants[_id].released);
         require(available > 0, "Must have available granted amount to stake.");
 
+        revertIfDelegateStakes(msg.sender);
+
         // Lock grant from releasing its balance.
         grants[_id].locked = true;
 

contracts/solidity/contracts/TokenStaking.sol

@@ -5,6 +5,7 @@ import "openzeppelin-solidity/contracts/token/ERC20/SafeERC20.sol";
 import "openzeppelin-solidity/contracts/math/SafeMath.sol";
 import "./StakingProxy.sol";
 import "./utils/UintArrayUtils.sol";
+import "./StakeDelegatable.sol";
 
 
 /**
@@ -13,7 +14,7 @@ import "./utils/UintArrayUtils.sol";
  * A holder of the specified token can stake its tokens to this contract
  * and unstake after withdrawal delay is over.
  */
-contract TokenStaking {
+contract TokenStaking is StakeDelegatable {
     using SafeMath for uint256;
     using SafeERC20 for StandardToken;
     using UintArrayUtils for uint256[];
@@ -35,7 +36,6 @@ contract TokenStaking {
     uint256 public withdrawalDelay;
     uint256 public numWithdrawals;
 
-    mapping(address => uint256) public balances;
     mapping(address => uint256[]) public withdrawalIndices;
     mapping(uint256 => Withdrawal) public withdrawals;
 
@@ -67,14 +67,19 @@ contract TokenStaking {
         require(StandardToken(_token) == token, "Token contract must be the same one linked to this contract.");
         require(_value <= token.balanceOf(_from), "Sender must have enough tokens.");
 
+        revertIfDelegateStakes(_from);
+
         // Transfer tokens to this contract.
         token.transferFrom(_from, this, _value);
 
         // Maintain a record of the stake amount by the sender.
-        balances[_from] = balances[_from].add(_value);
-        emit Staked(_from, _value);
+        stakeBalances[_from] = stakeBalances[_from].add(_value);
+
+        // Emit staked event. Check if staker works via operator first.
+        address delegatorOrDelegate = getDelegatorOrDelegate(_from);
+        emit Staked(delegatorOrDelegate, _value);
         if (address(stakingProxy) != address(0)) {
-            stakingProxy.emitStakedEvent(_from, _value);
+            stakingProxy.emitStakedEvent(delegatorOrDelegate, _value);
         }
     }
 
@@ -86,16 +91,19 @@ contract TokenStaking {
      */
     function initiateUnstake(uint256 _value) public returns (uint256 id) {
 
-        require(_value <= balances[msg.sender], "Staker must have enough tokens to unstake.");
+        require(_value <= stakeBalances[msg.sender], "Staker must have enough tokens to unstake.");
 
-        balances[msg.sender] = balances[msg.sender].sub(_value);
+        stakeBalances[msg.sender] = stakeBalances[msg.sender].sub(_value);
 
         id = numWithdrawals++;
         withdrawals[id] = Withdrawal(msg.sender, _value, now);
         withdrawalIndices[msg.sender].push(id);
         emit InitiatedUnstake(id);
+
+        // Emit unstaked event. Check if staker delegated its balance first.
+        address delegatorOrDelegate = getDelegatorOrDelegate(msg.sender);
         if (address(stakingProxy) != address(0)) {
-            stakingProxy.emitUnstakedEvent(msg.sender, _value);
+            stakingProxy.emitUnstakedEvent(delegatorOrDelegate, _value);
         }
         return id;
     }
@@ -123,15 +131,6 @@ contract TokenStaking {
         emit FinishedUnstake(_id);
     }
 
-    /**
-     * @dev Gets the stake balance of the specified address.
-     * @param _staker The address to query the balance of.
-     * @return An uint256 representing the amount owned by the passed address.
-     */
-    function stakeBalanceOf(address _staker) public view returns (uint256 balance) {
-        return balances[_staker];
-    }
-
     /**
      * @dev Gets withdrawal request by ID.
      * @param _id ID of withdrawal request.

contracts/solidity/test/TestStakeDelegate.js

@@ -0,0 +1,95 @@
+import increaseTime, { duration, increaseTimeTo } from './helpers/increaseTime';
+import exceptThrow from './helpers/expectThrow';
+const KeepToken = artifacts.require('./KeepToken.sol');
+const StakingProxy = artifacts.require('./StakingProxy.sol');
+const TokenStaking = artifacts.require('./TokenStaking.sol');
+
+
+contract('TestStakeDelegate', function(accounts) {
+
+  let token, stakingProxy, stakingContract,
+    account_one = accounts[0],
+    account_two = accounts[1],
+    account_three = accounts[2],
+    account_four = accounts[3];
+
+  beforeEach(async () => {
+    token = await KeepToken.new();
+    stakingProxy = await StakingProxy.new();
+    stakingContract = await TokenStaking.new(token.address, stakingProxy.address, duration.days(30));
+    await stakingProxy.authorizeContract(stakingContract.address, {from: account_one})
+
+    // Stake tokens as account one
+    await token.approveAndCall(stakingContract.address, 200, "", {from: account_one});
+
+    // Send tokens to the accounts
+    await token.transfer(account_two, 200, {from: account_one});
+    await token.transfer(account_three, 500, {from: account_one});
+
+    // Stake tokens as account two
+    await token.approveAndCall(stakingContract.address, 200, "", {from: account_two});
+  });
+
+  it("should not be able to delegate stake to a delegate address that is a staker", async function() {
+    await stakingContract.requestDelegateFor(account_one, {from: account_two});
+    await exceptThrow(stakingContract.approveDelegateAt(account_two));
+  });
+
+  it("should be able to delegate stake to a delegate address to represent your stake balance", async function() {
+    await stakingContract.requestDelegateFor(account_one, {from: account_three});
+    await stakingContract.approveDelegateAt(account_three, {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 0, "Delegator account stake balance should become zero.");
+  });
+
+  it("should not be able to delegate stake to a delegate address that is not approved", async function() {
+    await stakingContract.requestDelegateFor(account_two, {from: account_three});
+    await stakingContract.approveDelegateAt(account_three, {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 0, "Delegate account should be zero since there were no handshake with delegator.");
+  });
+
+  it("should be able to update delegate address if new delegate request exist", async function() {
+    await stakingContract.requestDelegateFor(account_one, {from: account_three});
+    await stakingContract.approveDelegateAt(account_three, {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 0, "Delegator account stake balance should become zero.");
+
+    await stakingContract.approveDelegateAt(account_four, {from: account_one});
+    await stakingContract.requestDelegateFor(account_one, {from: account_four});
+    assert.equal(await stakingProxy.balanceOf(account_three), 0, "Previous delegate account should stop representing delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_four), 200, "Updated delegate account should represent delegator's stake balance.");
+  });
+
+  it("should be able to remove delegated address that represents your stake balance", async function() {
+    await stakingContract.requestDelegateFor(account_one, {from: account_three});
+    await stakingContract.approveDelegateAt(account_three, {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+    await stakingContract.removeDelegate();
+    assert.equal(await stakingProxy.balanceOf(account_three), 0, "Delegate account should stop representing delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 200, "Delegator account should get its balance back.");
+  });
+
+  it("should be able to change stake and get delegate to reflect updated balance", async function() {
+    await stakingContract.requestDelegateFor(account_one, {from: account_three});
+    await stakingContract.approveDelegateAt(account_three, {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+
+    // Stake more tokens
+    await token.approveAndCall(stakingContract.address, 100, "", {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 300, "Delegate account should reflect delegator's updated stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 0, "Delegator account stake balance should be zero.");
+
+    // Unstake everything
+    await stakingContract.initiateUnstake(300);
+    assert.equal(await stakingProxy.balanceOf(account_three), 0, "Delegate account should reflect delegator's updated stake balance.");
+  });
+
+  it("should revert if delegate tries to stake", async function() {
+    await stakingContract.requestDelegateFor(account_one, {from: account_three});
+    await stakingContract.approveDelegateAt(account_three, {from: account_one});
+    assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+
+    // Stake tokens as account three
+    await exceptThrow(token.approveAndCall(stakingContract.address, 500, "", {from: account_three}));
+  });
+});