Get rid of pruning logic

Author: adamdossa

contracts/libraries/TokenLib.sol

@@ -25,12 +25,10 @@ library TokenLib {
     }
 
     struct InvestorDataStorage {
-        // List of investors (may not be pruned to remove old investors with current zero balances)
-        mapping (address => bool) investorListed;
+        // Index of investors
+        mapping (address => uint256) investorIndex;
         // List of token holders
         address[] investors;
-        // Total number of non-zero token holders
-        uint256 investorCount;
     }
 
     // Emit when Module is archived from the SecurityToken
@@ -163,45 +161,22 @@ library TokenLib {
         address _from,
         address _to,
         uint256 _value,
-        uint256 _balanceTo,
-        uint256 _balanceFrom
+        uint256 _balanceFrom,
+        uint256 _balanceTo
         ) 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);
+            _investorData.investorIndex[_to] = _investorData.investors.length;
+            _investorData.investors.push(_to);
         }
         // 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 is the date related to investor metrics
-    * 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) public {
-        for (uint256 i = 0; i < _investorData.investors.length; i++) {
-            if (_investorData.investors[i] == address(0)) {
-                while(_investorData.investors.length > i) {
-                    if (_investorData.investors[_investorData.investors.length - 1] != address(0)) {
-                        _investorData.investors[i] = _investorData.investors[_investorData.investors.length - 1];
-                        _investorData.investors.length--;
-                        break;
-                    }
-                    _investorData.investors.length--;
-                }
-            }
+            _investorData.investors[_investorData.investorIndex[_from]] = _investorData.investors[_investorData.investors.length - 1];
+            _investorData.investorIndex[_investorData.investors[_investorData.investorIndex[_from]]] = _investorData.investorIndex[_from];
+            _investorData.investors.length--;
         }
     }
 

contracts/tokens/SecurityToken.sol

@@ -38,7 +38,7 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
 
     SemanticVersion securityTokenVersion;
 
-    // off-chain data 
+    // off-chain data
     string public tokenDetails;
 
     uint8 constant PERMISSION_KEY = 1;
@@ -393,24 +393,7 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
     * @param _value value of transfer
     */
     function _adjustInvestorCount(address _from, address _to, uint256 _value) internal {
-        TokenLib.adjustInvestorCount(investorData, _from, _to, _value, balanceOf(_to), balanceOf(_from));
-    }
-
-    /**
-    * @notice removes addresses with zero balances from the investors list
-    * @param _start Index in investor list at which to start removing zero balances
-    * @param _iters Max number of iterations of the for loop
-    * 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 {
-        uint iter = Math.min256(_start.add(_iters), investorData.investors.length);
-        for (uint256 i = _start; i < iter; i++) {
-            if (balanceOf(investorData.investors[i]) == 0) {
-                investorData.investorListed[investorData.investors[i]] = false;
-                investorData.investors[i] = address(0);
-            }
-        }
-        TokenLib.pruneInvestors(investorData);
+        TokenLib.adjustInvestorCount(investorData, _from, _to, _value, balanceOf(_from), balanceOf(_to));
     }
 
     /**
@@ -426,7 +409,7 @@ contract SecurityToken is StandardToken, DetailedERC20, ReentrancyGuard, Registr
      * @notice returns the investor count
      */
     function getInvestorCount() external view returns(uint256) {
-        return investorData.investorCount;
+        return investorData.investors.length;
     }
 
     /**

test/o_security_token.js

@@ -3,7 +3,7 @@ import { duration, ensureException, promisifyLogWatch, latestBlock } from "./hel
 import takeSnapshot, { increaseTime, revertToSnapshot } from "./helpers/time";
 import { encodeProxyCall, encodeModuleCall } from "./helpers/encodeCall";
 import { catchRevert } from "./helpers/exceptions";
-import { 
+import {
     setUpPolymathNetwork,
     deployGPMAndVerifyed,
     deployCappedSTOAndVerifyed,
@@ -501,7 +501,7 @@ contract("SecurityToken", accounts => {
         it("Should fail to get the total supply -- because checkpoint id is greater than present", async() => {
             await catchRevert(
                 I_SecurityToken.totalSupplyAt.call(50)
-            ); 
+            );
         })
     });
 
@@ -855,7 +855,7 @@ contract("SecurityToken", accounts => {
         });
 
         it("Should prune investor length", async () => {
-            await I_SecurityToken.pruneInvestors(0, 10, { from: token_owner });
+            // await I_SecurityToken.pruneInvestors(0, 10, { from: token_owner });
             // Hardcode list of expected accounts based on transfers above
 
             let investors = await I_SecurityToken.getInvestors.call();
@@ -871,7 +871,7 @@ contract("SecurityToken", accounts => {
             let balance2 = await I_SecurityToken.balanceOf(account_investor1);
             await I_SecurityToken.transfer(account_affiliate1, balance, { from: account_affiliate2});
             await I_SecurityToken.transfer(account_affiliate1, balance2, { from: account_investor1});
-            await I_SecurityToken.pruneInvestors(0, 10, { from: token_owner });
+            // await I_SecurityToken.pruneInvestors(0, 10, { from: token_owner });
             let investors = await I_SecurityToken.getInvestors.call();
             let expectedAccounts = [account_affiliate1];
             for (let i = 0; i < expectedAccounts.length; i++) {
@@ -893,7 +893,7 @@ contract("SecurityToken", accounts => {
     });
 
     describe("Test cases for the Mock TrackedRedeemption", async() => {
-        
+
         it("Should add the tracked redeemption module successfully", async() => {
             [I_MockRedemptionManagerFactory] = await deployMockRedemptionAndVerifyed(account_polymath, I_MRProxied, I_PolyToken.address, 0);
             let tx = await I_SecurityToken.addModule(I_MockRedemptionManagerFactory.address, "", 0, 0, {from: token_owner });
@@ -959,7 +959,7 @@ contract("SecurityToken", accounts => {
             await I_SecurityToken.approve(I_MockRedemptionManager.address, web3.utils.toWei("500"), {from: account_investor1});
             // Transfer the tokens to module (Burn)
             await I_MockRedemptionManager.transferToRedeem(web3.utils.toWei("500"), { from: account_investor1});
-            
+
             await catchRevert(
                 // Redeem tokens
                 I_MockRedemptionManager.redeemTokenByOwner(web3.utils.toWei("250"), {from: account_investor1})
@@ -1105,4 +1105,4 @@ contract("SecurityToken", accounts => {
         });
     });
 
-});
+});