Migrate the investors logic from ST to library

contracts/interfaces/ISecurityToken.sol

@@ -116,12 +116,6 @@ interface ISecurityToken {
     */
     function investors(uint256 _index) external view returns (address);
 
-    /**
-    * @notice gets the number of investors
-    * @return count of investors
-    */
-    function investorCount() external view returns (uint256);
-
     /**
     * @notice allows the owner to withdraw unspent POLY stored by them on the ST.
     * @dev Owner can transfer POLY to the ST which will be used to pay for modules that require a POLY fee.
@@ -245,4 +239,9 @@ interface ISecurityToken {
      * @notice Use to get the version of the securityToken
      */
      function getVersion() external view returns(uint8[]);
+
+     /**
+     * @notice gets the investor count
+     */
+     function getInvestorCount() external view returns(uint256);
 }

contracts/libraries/TokenLib.sol

@@ -1,9 +1,12 @@
 pragma solidity ^0.4.24;
 
 import "../modules/PermissionManager/IPermissionManager.sol";
+import "./KindMath.sol";
 
 library TokenLib {
 
+    using KindMath for uint256;
+
     // Struct for module data
     struct ModuleData {
         bytes32 name;
@@ -21,6 +24,15 @@ library TokenLib {
         uint256 value;
     }
 
+    struct InvestorDataStorage {
+        // List of investors (may not be pruned to remove old investors with current zero balances)
+        mapping (address => bool) investorListed;
+        // List of token holders
+        address[] investors;
+        // Total number of non-zero token holders
+        uint256 investorCount;
+    }
+
     // Emit when Module get archived from the securityToken
     event ModuleArchived(uint8[] _types, address _module, uint256 _timestamp);
     // Emit when Module get unarchived from the securityToken
@@ -137,4 +149,45 @@ library TokenLib {
         );
     }
 
+    /**
+    * @notice keeps track of the number of non-zero token holders
+    * @param _investorData Date releated to investor metrics
+    * @param _from sender of transfer
+    * @param _to receiver of transfer
+    * @param _value value of transfer
+    * @param _balanceTo balance of the _to address
+    * @param _balanceFrom balance of the _from address
+    */
+    function adjustInvestorCount(InvestorDataStorage storage _investorData, address _from, address _to, uint256 _value, uint256 _balanceTo, uint256 _balanceFrom) public  {
+        if ((_value == 0) || (_from == _to)) {
+            return;
+        }
+        // Check whether receiver is a new token holder
+        if ((_balanceTo == 0) && (_to != address(0))) {
+            _investorData.investorCount = (_investorData.investorCount).add(1);
+        }
+        // Check whether sender is moving all of their tokens
+        if (_value == _balanceFrom) {
+            _investorData.investorCount = (_investorData.investorCount).sub(1);
+        }
+        //Also adjust investor list
+        if (!_investorData.investorListed[_to] && (_to != address(0))) {
+            _investorData.investors.push(_to);
+            _investorData.investorListed[_to] = true;
+        }
+
+    }
+
+     /**
+    * @notice removes addresses with zero balances from the investors list
+    * @param _investorData Date releated to investor metrics
+    * @param _index Index in investor list
+    * NB - pruning this list will mean you may not be able to iterate over investors on-chain as of a historical checkpoint
+    */
+    function pruneInvestors(InvestorDataStorage storage _investorData, uint256 _index) public {
+        _investorData.investorListed[_investorData.investors[_index]] = false;
+        _investorData.investors[_index] = _investorData.investors[_investorData.investors.length - 1];
+        _investorData.investors.length--;
+    }
+
 }

contracts/modules/TransferManager/CountTransferManager.sol

@@ -28,7 +28,7 @@ contract CountTransferManager is ITransferManager {
     /// @notice Used to verify the transfer transaction according to the rule implemented in the trnasfer managers
     function verifyTransfer(address /* _from */, address _to, uint256 /* _amount */, bool /* _isTransfer */) public returns(Result) {
         if (!paused) {
-            if (maxHolderCount < ISecurityToken(securityToken).investorCount()) {
+            if (maxHolderCount < ISecurityToken(securityToken).getInvestorCount()) {
                 // Allow transfers to existing maxHolders
                 if (ISecurityToken(securityToken).balanceOf(_to) != 0) {
                     return Result.NA;

contracts/tokens/SecurityToken.sol

@@ -27,6 +27,8 @@ import "../libraries/TokenLib.sol";
 contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, RegistryUpdater {
     using SafeMath for uint256;
 
+    TokenLib.InvestorDataStorage investorData;
+
     // Use to hold the version
     struct SemanticVersion {
         uint8 major;
@@ -50,12 +52,6 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
     // Value of current checkpoint
     uint256 public currentCheckpointId;
 
-    // Total number of non-zero token holders
-    uint256 public investorCount;
-
-    // List of token holders
-    address[] investors;
-
     // Use to temporarily halt all transactions
     bool public transfersFrozen;
 
@@ -86,9 +82,6 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
     // Times at which each checkpoint was created
     uint256[] checkpointTimes;
 
-    // List of investors (may not be pruned to remove old investors with current zero balances)
-    mapping (address => bool) investorListed;
-
     // Emit at the time when module get added
     event ModuleAdded(
         uint8[] _types,
@@ -401,23 +394,7 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
     * @param _value value of transfer
     */
     function _adjustInvestorCount(address _from, address _to, uint256 _value) internal {
-        if ((_value == 0) || (_from == _to)) {
-            return;
-        }
-        // Check whether receiver is a new token holder
-        if ((balanceOf(_to) == 0) && (_to != address(0))) {
-            investorCount = investorCount.add(1);
-        }
-        // Check whether sender is moving all of their tokens
-        if (_value == balanceOf(_from)) {
-            investorCount = investorCount.sub(1);
-        }
-        //Also adjust investor list
-        if (!investorListed[_to] && (_to != address(0))) {
-            investors.push(_to);
-            investorListed[_to] = true;
-        }
-
+        TokenLib.adjustInvestorCount(investorData, _from, _to, _value, balanceOf(_to), balanceOf(_from));
     }
 
     /**
@@ -427,11 +404,9 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
     * NB - pruning this list will mean you may not be able to iterate over investors on-chain as of a historical checkpoint
     */
     function pruneInvestors(uint256 _start, uint256 _iters) external onlyOwner {
-        for (uint256 i = _start; i < Math.min256(_start.add(_iters), investors.length); i++) {
-            if ((i < investors.length) && (balanceOf(investors[i]) == 0)) {
-                investorListed[investors[i]] = false;
-                investors[i] = investors[investors.length - 1];
-                investors.length--;
+        for (uint256 i = _start; i < Math.min256(_start.add(_iters), investorData.investors.length); i++) {
+            if ((i < investorData.investors.length) && (balanceOf(investorData.investors[i]) == 0)) {
+                TokenLib.pruneInvestors(investorData, i);
             }
         }
     }
@@ -442,7 +417,14 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
      * @return length
      */
     function getInvestors() external view returns(address[]) {
-        return investors;
+        return investorData.investors;
+    }
+
+    /**
+     * @notice gets the investor count
+     */
+    function getInvestorCount() external view returns(uint256) {
+        return investorData.investorCount;
     }
 
     /**

test/o_security_token.js

@@ -1154,7 +1154,7 @@ contract('SecurityToken', accounts => {
           it("Should force burn the tokens - value too high", async ()=> {
               let errorThrown = false;
               await I_GeneralTransferManager.changeAllowAllBurnTransfers(true, {from : token_owner});
-              let currentInvestorCount = await I_SecurityToken.investorCount();
+              let currentInvestorCount = await I_SecurityToken.getInvestorCount.call();
               let currentBalance = await I_SecurityToken.balanceOf(account_temp);
               try {
                   let tx = await I_SecurityToken.forceBurn(account_temp, currentBalance + web3.utils.toWei("500", "ether"), "", { from: account_controller });
@@ -1168,7 +1168,7 @@ contract('SecurityToken', accounts => {
           it("Should force burn the tokens - wrong caller", async ()=> {
                let errorThrown = false;
                await I_GeneralTransferManager.changeAllowAllBurnTransfers(true, {from : token_owner});
-               let currentInvestorCount = await I_SecurityToken.investorCount();
+               let currentInvestorCount = await I_SecurityToken.getInvestorCount.call();
                let currentBalance = await I_SecurityToken.balanceOf(account_temp);
                try {
                    let tx = await I_SecurityToken.forceBurn(account_temp, currentBalance, "", { from: token_owner });
@@ -1181,13 +1181,13 @@ contract('SecurityToken', accounts => {
            });
 
            it("Should burn the tokens", async ()=> {
-                let currentInvestorCount = await I_SecurityToken.investorCount();
+                let currentInvestorCount = await I_SecurityToken.getInvestorCount.call();
                 let currentBalance = await I_SecurityToken.balanceOf(account_temp);
                 // console.log(currentInvestorCount.toString(), currentBalance.toString());
                 let tx = await I_SecurityToken.forceBurn(account_temp, currentBalance, "", { from: account_controller });
                 // console.log(tx.logs[0].args._value.toNumber(), currentBalance.toNumber());
                 assert.equal(tx.logs[0].args._value.toNumber(), currentBalance.toNumber());
-                let newInvestorCount = await I_SecurityToken.investorCount();
+                let newInvestorCount = await I_SecurityToken.getInvestorCount.call();
                 // console.log(newInvestorCount.toString());
                 assert.equal(newInvestorCount.toNumber() + 1, currentInvestorCount.toNumber(), "Investor count drops by one");
            });
@@ -1298,13 +1298,13 @@ contract('SecurityToken', accounts => {
                 let sender = account_investor1;
                 let receiver = account_investor2;
 
-                let start_investorCount = await I_SecurityToken.investorCount.call();
+                let start_investorCount = await I_SecurityToken.getInvestorCount.call();
                 let start_balInv1 = await I_SecurityToken.balanceOf.call(account_investor1);
                 let start_balInv2 = await I_SecurityToken.balanceOf.call(account_investor2);
 
                 let tx = await I_SecurityToken.forceTransfer(account_investor1, account_investor2, web3.utils.toWei("10", "ether"), "reason", {from: account_controller});
 
-                let end_investorCount = await I_SecurityToken.investorCount.call();
+                let end_investorCount = await I_SecurityToken.getInvestorCount.call();
                 let end_balInv1 = await I_SecurityToken.balanceOf.call(account_investor1);
                 let end_balInv2 = await I_SecurityToken.balanceOf.call(account_investor2);