From e0044359fcb4d044bf192c4e243b0f6c3c2e9a05 Mon Sep 17 00:00:00 2001
From: Rongjian Lan <rongjian.lan@gmail.com>
Date: Tue, 5 Mar 2019 16:54:53 -0800
Subject: [PATCH] Add listLockedAddresses and make lock/unlock works

---
 contracts/StakeLockContract.sol | 74 +++++++++++++++++++++------------
 1 file changed, 48 insertions(+), 26 deletions(-)

diff --git a/contracts/StakeLockContract.sol b/contracts/StakeLockContract.sol
index a49715b61..f103c39af 100644
--- a/contracts/StakeLockContract.sol
+++ b/contracts/StakeLockContract.sol
@@ -5,7 +5,7 @@ contract StakeLockContract {
      * @dev Error messages for require statements
      */
     string internal constant ALREADY_LOCKED = 'Tokens already locked';
-    string internal constant NOT_LOCKED = 'No tokens locked';
+    string internal constant NO_TOKEN_UNLOCKABLE = 'No tokens unlockable';
     string internal constant AMOUNT_ZERO = 'Amount can not be 0';
 
     uint256 internal constant LOCK_PERIOD_IN_EPOCHS = 3;  // Final locking period TBD.
@@ -20,6 +20,7 @@ contract StakeLockContract {
         uint256 _blockNum;
         uint256 _epochNum;
         uint256 _lockPeriodCount;  // The number of locking period the token will be locked.
+        uint256 _index;  // The index in the addressList
     }
 
     /**
@@ -28,62 +29,54 @@ contract StakeLockContract {
      */
     mapping(address => lockedToken) private locked;
     mapping(address => uint) private indices;
-    address[] private accountList;
+    address[] private addressList;
 
     event Locked(address indexed _of, uint _amount, uint256 _epoch);
     event Unlocked(address indexed account, uint index);
 
-    /**
-     * @dev Returns total tokens held by an address (locked + transferable)
-     * @param _of The address to query the total balance of
-     */
-    function lockedBalanceOf(address _of)
-        public
-        view
-        returns (uint256 amount)
-    {
-        amount = locked[_of]._amount;
-    }
-
     /**
      * @dev Locks a specified amount of tokens against an address
      *      starting at the specific epoch
-     * @param _epoch The epoch where the lock happens
      */
-    function lock(uint256 _epoch)
+    function lock()
         public
         payable
         returns (bool)
     {
         // If tokens are already locked, then functions extendLock or
         // increaseLockAmount should be used to make any changes
-        require(lockedBalanceOf(msg.sender) == 0, ALREADY_LOCKED);
+        require(balanceOf(msg.sender) == 0, ALREADY_LOCKED);
         require(msg.value != 0, AMOUNT_ZERO);
 
-        locked[msg.sender] = lockedToken(msg.value, block.number, currentEpoch(), 1);
+        // By default, the tokens can only be locked for one locking period.
+        locked[msg.sender] = lockedToken(msg.value, block.number, currentEpoch(), 1, addressList.push(msg.sender) - 1);
 
-        emit Locked(msg.sender, msg.value, _epoch);
+        emit Locked(msg.sender, msg.value, currentEpoch());
         return true;
     }
 
     /**
      * @dev Unlocks the unlockable tokens of a specified address
-     * @param _requestor Address of user, claiming back unlockable tokens
      */
-    function unlock(address payable _requestor)
+    function unlock()
         public
         returns (uint256 unlockableTokens)
     {
-        unlockableTokens = getUnlockableTokens(_requestor);
+        unlockableTokens = getUnlockableTokens(msg.sender);  // For now the unlockableTokens is all the tokens for a address
+        require(unlockableTokens != 0, NO_TOKEN_UNLOCKABLE);
 
-        if (unlockableTokens != 0) {
-            _requestor.transfer(unlockableTokens);
-        }
+        uint indexToRemove = locked[msg.sender]._index;
+        delete locked[msg.sender];
+        addressList[indexToRemove] = addressList[addressList.length - 1];
+        locked[addressList[addressList.length - 1]]._index = indexToRemove;
+        addressList.length--;
+
+        msg.sender.transfer(unlockableTokens);
     }
 
     /**
      * @dev Gets the unlockable tokens of a specified address
-     * @param _of The address to query the the unlockable token count of
+     * @param _of The address to query the unlockable token
      */
     function getUnlockableTokens(address _of)
         public
@@ -96,6 +89,18 @@ contract StakeLockContract {
         }
     }
 
+    /**
+     * @dev Gets the token balance of a specified address
+     * @param _of The address to query the token balance
+     */
+    function balanceOf(address _of)
+        public
+        view
+        returns (uint256 balance)
+    {
+        balance = locked[_of]._amount;
+    }
+
     /**
      * @dev Gets the epoch number of the specified block number.
      */
@@ -106,4 +111,21 @@ contract StakeLockContract {
     {
         return block.number / numBlocksPerEpoch;
     }
+
+    /**
+     * @dev Lists all the locked address data.
+     */
+    function listLockedAddresses()
+        public
+        view
+        returns (address[] memory lockedAddresses, uint256[] memory blockNums, uint256[] memory lockPeriodCounts)
+    {
+        lockedAddresses = addressList;
+        blockNums = new uint256[](addressList.length);
+        lockPeriodCounts = new uint256[](addressList.length);
+        for (uint i = 0; i < lockedAddresses.length; i++) {
+            blockNums[i] = locked[lockedAddresses[i]]._blockNum;
+            lockPeriodCounts[i] = locked[lockedAddresses[i]]._lockPeriodCount;
+        }
+    }
 }