Remove the investor array and add isVolRestricted flag for investors

contracts/libraries/VolumeRestrictionLib.sol

@@ -1,12 +1,19 @@
 pragma solidity ^0.5.0;
 
+import "../interfaces/IDataStore.sol";
 import "openzeppelin-solidity/contracts/math/SafeMath.sol";
 import "../storage/modules/TransferManager/VolumeRestrictionTMStorage.sol";
 
 library VolumeRestrictionLib {
 
     using SafeMath for uint256;
 
+    uint256 internal constant ONE = uint256(1);
+    uint8 internal constant INDEX = uint8(2);
+    bytes32 internal constant INVESTORFLAGS = "INVESTORFLAGS";
+    bytes32 internal constant INVESTORSKEY = 0xdf3a8dd24acdd05addfc6aeffef7574d2de3f844535ec91e8e0f3e45dba96731; //keccak256(abi.encodePacked("INVESTORS"))
+    bytes32 internal constant WHITELIST = "WHITELIST";
+
     function _checkLengthOfArray(
         address[] memory _holders,
         uint256[] memory _allowedTokens,
@@ -28,60 +35,89 @@ library VolumeRestrictionLib {
         );
     }
 
-    function deleteHolderFromList(VolumeRestrictionTMStorage.RestrictedData storage data, address _holder, uint8 _typeOfPeriod) public {
+    function deleteHolderFromList(
+        mapping(address => uint8) storage holderToRestrictionType,
+        address _holder,
+        address _dataStore,
+        uint8 _typeOfPeriod
+    ) 
+        public
+    {
         // Deleting the holder if holder's type of Period is `Both` type otherwise
         // it will assign the given type `_typeOfPeriod` to the _holder typeOfPeriod
         // `_typeOfPeriod` it always be contrary to the removing restriction
         // if removing restriction is individual then typeOfPeriod is TypeOfPeriod.OneDay
         // in uint8 its value is 1. if removing restriction is daily individual then typeOfPeriod
         // is TypeOfPeriod.MultipleDays in uint8 its value is 0.
-        if (data.restrictedHolders[_holder].typeOfPeriod != uint8(VolumeRestrictionTMStorage.TypeOfPeriod.Both)) {
-            uint128 index = data.restrictedHolders[_holder].index;
-            uint256 _len = data.restrictedAddresses.length;
-            if (index != _len) {
-                data.restrictedHolders[data.restrictedAddresses[_len - 1]].index = index;
-                data.restrictedAddresses[index - 1] = data.restrictedAddresses[_len - 1];
-            }
-            delete data.restrictedHolders[_holder];
-            data.restrictedAddresses.length--;
+        if (holderToRestrictionType[_holder] != uint8(VolumeRestrictionTMStorage.TypeOfPeriod.Both)) {
+            IDataStore dataStore = IDataStore(_dataStore);
+            uint256 flags = dataStore.getUint256(_getKey(INVESTORFLAGS, _holder));
+            flags = flags & ~(ONE << INDEX);
+            dataStore.setUint256(_getKey(INVESTORFLAGS, _holder), flags);
         } else {
-            data.restrictedHolders[_holder].typeOfPeriod = _typeOfPeriod;
+            holderToRestrictionType[_holder] = _typeOfPeriod;
         }
     }
 
-    function addRestrictionData(VolumeRestrictionTMStorage.RestrictedData storage data, address _holder, uint8 _callFrom, uint256 _endTime) public {
-        uint128 index = data.restrictedHolders[_holder].index;
-        if (data.restrictedHolders[_holder].seen == 0) {
-            data.restrictedAddresses.push(_holder);
-            index = uint128(data.restrictedAddresses.length);
+    function addRestrictionData(
+        mapping(address => uint8) storage holderToRestrictionType,
+        address _holder,
+        uint8 _callFrom,
+        uint256 _endTime,
+        address _dataStore
+    ) 
+        public 
+    {
+        IDataStore dataStore = IDataStore(_dataStore);
+
+        uint256 flags = dataStore.getUint256(_getKey(INVESTORFLAGS, _holder));
+        if (!_isExistingInvestor(_holder, dataStore)) {
+            dataStore.insertAddress(INVESTORSKEY, _holder);
+            //KYC data can not be present if added is false and hence we can set packed KYC as uint256(1) to set added as true
+            dataStore.setUint256(_getKey(WHITELIST, _holder), uint256(1));
         }
-        uint8 _type = _getTypeOfPeriod(data.restrictedHolders[_holder].typeOfPeriod, _callFrom, _endTime);
-        data.restrictedHolders[_holder] = VolumeRestrictionTMStorage.RestrictedHolder(uint8(1), _type, index);
+        if (!_isVolRestricted(flags)) {
+            flags = flags | (ONE << INDEX);
+            dataStore.setUint256(_getKey(INVESTORFLAGS, _holder), flags);
+        }
+        uint8 _type = _getTypeOfPeriod(holderToRestrictionType[_holder], _callFrom, _endTime);
+        holderToRestrictionType[_holder] = _type;
     }
 
-    function _getTypeOfPeriod(uint8 _currentTypeOfPeriod, uint8 _callFrom, uint256 _endTime) internal pure returns(uint8) {
-        if (_currentTypeOfPeriod != _callFrom && _endTime != uint256(0))
-            return uint8(VolumeRestrictionTMStorage.TypeOfPeriod.Both);
-        else
-            return _callFrom;
-    }
 
+    /**
+     * @notice Provide the restriction details of all the restricted addresses
+     * @return address List of the restricted addresses
+     * @return uint256 List of the tokens allowed to the restricted addresses corresponds to restricted address
+     * @return uint256 List of the start time of the restriction corresponds to restricted address
+     * @return uint256 List of the rolling period in days for a restriction corresponds to restricted address.
+     * @return uint256 List of the end time of the restriction corresponds to restricted address.
+     * @return uint8 List of the type of restriction to validate the value of the `allowedTokens`
+     * of the restriction corresponds to restricted address
+     */
     function getRestrictionData(
-        VolumeRestrictionTMStorage.RestrictedData storage _holderData,
-        VolumeRestrictionTMStorage.IndividualRestrictions storage _individualRestrictions
-    ) public view returns(
-        address[] memory allAddresses,
-        uint256[] memory allowedTokens,
-        uint256[] memory startTime,
-        uint256[] memory rollingPeriodInDays,
-        uint256[] memory endTime,
-        uint8[] memory typeOfRestriction
-    )
+        mapping(address => uint8) storage holderToRestrictionType,
+        VolumeRestrictionTMStorage.IndividualRestrictions storage _individualRestrictions,
+        address _dataStore
+    )   
+        public
+        view 
+        returns(
+            address[] memory allAddresses,
+            uint256[] memory allowedTokens,
+            uint256[] memory startTime,
+            uint256[] memory rollingPeriodInDays,
+            uint256[] memory endTime,
+            uint8[] memory typeOfRestriction
+            )
     {
-        uint256 counter = 0;
-        uint256 i = 0;
-        for (i = 0; i < _holderData.restrictedAddresses.length; i++) {
-            counter = counter + (_holderData.restrictedHolders[_holderData.restrictedAddresses[i]].typeOfPeriod == uint8(2) ? 2 : 1);
+        address[] memory investors = IDataStore(_dataStore).getAddressArray(INVESTORSKEY);
+        uint256 counter;
+        uint256 i;
+        for (i = 0; i < investors.length; i++) {
+            if (_isVolRestricted(IDataStore(_dataStore).getUint256(_getKey(INVESTORFLAGS, investors[i])))) {
+                counter = counter + (holderToRestrictionType[investors[i]] == uint8(2) ? 2 : 1);
+            }
         }
         allAddresses = new address[](counter);
         allowedTokens = new uint256[](counter);
@@ -90,21 +126,23 @@ library VolumeRestrictionLib {
         endTime = new uint256[](counter);
         typeOfRestriction = new uint8[](counter);
         counter = 0;
-        for (i = 0; i < _holderData.restrictedAddresses.length; i++) {
-            allAddresses[counter] =  _holderData.restrictedAddresses[i];
-            if (_holderData.restrictedHolders[_holderData.restrictedAddresses[i]].typeOfPeriod == uint8(VolumeRestrictionTMStorage.TypeOfPeriod.MultipleDays)) {
-                _setValues(_individualRestrictions.individualRestriction[_holderData.restrictedAddresses[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
-            }
-            else if (_holderData.restrictedHolders[_holderData.restrictedAddresses[i]].typeOfPeriod == uint8(VolumeRestrictionTMStorage.TypeOfPeriod.OneDay)) {
-                _setValues(_individualRestrictions.individualDailyRestriction[_holderData.restrictedAddresses[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
-            }
-            else if (_holderData.restrictedHolders[_holderData.restrictedAddresses[i]].typeOfPeriod == uint8(VolumeRestrictionTMStorage.TypeOfPeriod.Both)) {
-                _setValues(_individualRestrictions.individualRestriction[_holderData.restrictedAddresses[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
+        for (i = 0; i < investors.length; i++) {
+            if (_isVolRestricted(IDataStore(_dataStore).getUint256(_getKey(INVESTORFLAGS, investors[i])))) {
+                allAddresses[counter] = investors[i];
+                if (holderToRestrictionType[investors[i]] == uint8(VolumeRestrictionTMStorage.TypeOfPeriod.MultipleDays)) {
+                    _setValues(_individualRestrictions.individualRestriction[investors[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
+                } 
+                else if (holderToRestrictionType[investors[i]] == uint8(VolumeRestrictionTMStorage.TypeOfPeriod.OneDay)) {
+                    _setValues(_individualRestrictions.individualDailyRestriction[investors[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
+                } 
+                else if (holderToRestrictionType[investors[i]] == uint8(VolumeRestrictionTMStorage.TypeOfPeriod.Both)) {
+                    _setValues(_individualRestrictions.individualRestriction[investors[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
+                    counter++;
+                    allAddresses[counter] = investors[i];
+                    _setValues(_individualRestrictions.individualDailyRestriction[investors[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
+                }
                 counter++;
-                allAddresses[counter] =  _holderData.restrictedAddresses[i];
-                _setValues(_individualRestrictions.individualDailyRestriction[_holderData.restrictedAddresses[i]], allowedTokens, startTime, rollingPeriodInDays, endTime, typeOfRestriction, counter);
             }
-            counter++;
         }
     }
 
@@ -127,4 +165,26 @@ library VolumeRestrictionLib {
         typeOfRestriction[index] = uint8(restriction.typeOfRestriction);
     }
 
+    function _isVolRestricted(uint256 _flags) internal pure returns(bool) {
+        uint256 volRestricted = (_flags >> INDEX) & ONE;
+        return (volRestricted > 0 ? true : false);
+    }
+
+    function _getTypeOfPeriod(uint8 _currentTypeOfPeriod, uint8 _callFrom, uint256 _endTime) internal pure returns(uint8) {
+        if (_currentTypeOfPeriod != _callFrom && _endTime != uint256(0))
+            return uint8(VolumeRestrictionTMStorage.TypeOfPeriod.Both);
+        else
+            return _callFrom;
+    }
+
+    function _isExistingInvestor(address _investor, IDataStore dataStore) internal view returns(bool) {
+        uint256 data = dataStore.getUint256(_getKey(WHITELIST, _investor));
+        //extracts `added` from packed `_whitelistData`
+        return uint8(data) == 0 ? false : true;
+    }
+
+    function _getKey(bytes32 _key1, address _key2) internal pure returns(bytes32) {
+        return bytes32(keccak256(abi.encodePacked(_key1, _key2)));
+    }
+
 }

contracts/modules/TransferManager/VolumeRestrictionTM.sol

@@ -245,7 +245,14 @@ contract VolumeRestrictionTM is VolumeRestrictionTMStorage, TransferManager {
             _endTime,
             RestrictionType(_restrictionType)
         );
-        VolumeRestrictionLib.addRestrictionData(holderData, _holder, uint8(TypeOfPeriod.MultipleDays), individualRestrictions.individualRestriction[_holder].endTime);
+        VolumeRestrictionLib
+            .addRestrictionData(
+                holderToRestrictionType,
+                _holder,
+                uint8(TypeOfPeriod.MultipleDays),
+                individualRestrictions.individualRestriction[_holder].endTime,
+                getDataStore()
+            );
         emit AddIndividualRestriction(
             _holder,
             _allowedTokens,
@@ -292,15 +299,22 @@ contract VolumeRestrictionTM is VolumeRestrictionTMStorage, TransferManager {
             _endTime,
             RestrictionType(_restrictionType)
         );
-        VolumeRestrictionLib.addRestrictionData(holderData, _holder, uint8(TypeOfPeriod.OneDay), individualRestrictions.individualRestriction[_holder].endTime);
+        VolumeRestrictionLib
+            .addRestrictionData(
+                holderToRestrictionType,
+                _holder,
+                uint8(TypeOfPeriod.OneDay),
+                individualRestrictions.individualRestriction[_holder].endTime,
+                getDataStore()
+            );
         emit AddIndividualDailyRestriction(
             _holder,
             _allowedTokens,
             _startTime,
             1,
             _endTime,
             _restrictionType
-      );
+        );
     }
 
     /**
@@ -457,7 +471,7 @@ contract VolumeRestrictionTM is VolumeRestrictionTMStorage, TransferManager {
         require(_holder != address(0));
         require(individualRestrictions.individualRestriction[_holder].endTime != 0);
         individualRestrictions.individualRestriction[_holder] = VolumeRestriction(0, 0, 0, 0, RestrictionType(0));
-        VolumeRestrictionLib.deleteHolderFromList(holderData, _holder, uint8(TypeOfPeriod.OneDay));
+        VolumeRestrictionLib.deleteHolderFromList(holderToRestrictionType, _holder, getDataStore(), uint8(TypeOfPeriod.OneDay));
         bucketData.userToBucket[_holder].lastTradedDayTime = 0;
         bucketData.userToBucket[_holder].sumOfLastPeriod = 0;
         bucketData.userToBucket[_holder].daysCovered = 0;
@@ -482,7 +496,7 @@ contract VolumeRestrictionTM is VolumeRestrictionTMStorage, TransferManager {
         require(_holder != address(0));
         require(individualRestrictions.individualDailyRestriction[_holder].endTime != 0);
         individualRestrictions.individualDailyRestriction[_holder] = VolumeRestriction(0, 0, 0, 0, RestrictionType(0));
-        VolumeRestrictionLib.deleteHolderFromList(holderData, _holder, uint8(TypeOfPeriod.MultipleDays));
+        VolumeRestrictionLib.deleteHolderFromList(holderToRestrictionType, _holder, getDataStore(), uint8(TypeOfPeriod.MultipleDays));
         bucketData.userToBucket[_holder].dailyLastTradedDayTime = 0;
         emit IndividualDailyRestrictionRemoved(_holder);
     }
@@ -1122,7 +1136,7 @@ contract VolumeRestrictionTM is VolumeRestrictionTMStorage, TransferManager {
         uint256[] memory endTime,
         uint8[] memory typeOfRestriction
     ) {
-        return VolumeRestrictionLib.getRestrictionData(holderData, individualRestrictions);
+        return VolumeRestrictionLib.getRestrictionData(holderToRestrictionType, individualRestrictions, getDataStore());
     }
 
     /**

contracts/modules/TransferManager/VolumeRestrictionTMFactory.sol

@@ -48,8 +48,9 @@ contract VolumeRestrictionTMFactory is ModuleFactory {
      * @notice Type of the Module factory
      */
     function getTypes() external view returns(uint8[] memory) {
-        uint8[] memory res = new uint8[](1);
+        uint8[] memory res = new uint8[](2);
         res[0] = 2;
+        res[1] = 6;
         return res;
     }
 

contracts/storage/modules/TransferManager/VolumeRestrictionTMStorage.sol

@@ -9,22 +9,8 @@ contract VolumeRestrictionTMStorage {
 
     enum TypeOfPeriod { MultipleDays, OneDay, Both }
 
-    struct RestrictedHolder {
-        // 1 represent true & 0 for false
-        uint8 seen;
-        // Type of period will be enum index of TypeOfPeriod enum
-        uint8 typeOfPeriod;
-        // Index of the array where the holder address lives
-        uint128 index;
-    }
-
-    struct RestrictedData {
-        mapping(address => RestrictedHolder) restrictedHolders;
-        address[] restrictedAddresses;
-    }
-
-    // Restricted data (refernce from the VolumeRestrictionLib library )
-    RestrictedData holderData;
+    // Store the type of restriction corresponds to token holder address
+    mapping(address => uint8) holderToRestrictionType;
 
     struct VolumeRestriction {
         // If typeOfRestriction is `Percentage` then allowedTokens will be in

docs/investor_flags.md

@@ -19,5 +19,10 @@
             <td> canNotBuyFromSto </td>
             <td> Defines if an Investor is restricted from participating in STOs</td>
         </tr>
+        <tr>
+            <td> 2 </td>
+            <td> isVolRestricted </td>
+            <td> Defines if an Investor has the trade volume restriction (VRTM) or not </td>
+        </tr>
     </tbody>
 </table>