Merge pull request #566 from PolymathNetwork/takeSP-export

Refactor takeSnapShot

Author: maxsam4

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/Checkpoint/DividendCheckpoint.sol

@@ -79,15 +79,15 @@ contract DividendCheckpoint is DividendCheckpointStorage, ICheckpoint, Module {
      * @notice Creates a checkpoint on the security token
      * @return Checkpoint ID
      */
-    function createCheckpoint() public withPerm(CHECKPOINT) returns(uint256) {
+    function createCheckpoint() public withPerm(OPERATOR) returns(uint256) {
         return ISecurityToken(securityToken).createCheckpoint();
     }
 
     /**
      * @notice Function to clear and set list of excluded addresses used for future dividends
      * @param _excluded Addresses of investors
      */
-    function setDefaultExcluded(address[] memory _excluded) public withPerm(MANAGE) {
+    function setDefaultExcluded(address[] memory _excluded) public withPerm(ADMIN) {
         require(_excluded.length <= EXCLUDED_ADDRESS_LIMIT, "Too many excluded addresses");
         for (uint256 j = 0; j < _excluded.length; j++) {
             require(_excluded[j] != address(0), "Invalid address");
@@ -105,7 +105,7 @@ contract DividendCheckpoint is DividendCheckpointStorage, ICheckpoint, Module {
      * @param _investors Addresses of investors
      * @param _withholding Withholding tax for individual investors (multiplied by 10**16)
      */
-    function setWithholding(address[] memory _investors, uint256[] memory _withholding) public withPerm(MANAGE) {
+    function setWithholding(address[] memory _investors, uint256[] memory _withholding) public withPerm(ADMIN) {
         require(_investors.length == _withholding.length, "Mismatched input lengths");
         /*solium-disable-next-line security/no-block-members*/
         emit SetWithholding(_investors, _withholding);
@@ -120,7 +120,7 @@ contract DividendCheckpoint is DividendCheckpointStorage, ICheckpoint, Module {
      * @param _investors Addresses of investor
      * @param _withholding Withholding tax for all investors (multiplied by 10**16)
      */
-    function setWithholdingFixed(address[] memory _investors, uint256 _withholding) public withPerm(MANAGE) {
+    function setWithholdingFixed(address[] memory _investors, uint256 _withholding) public withPerm(ADMIN) {
         require(_withholding <= 10 ** 18, "Incorrect withholding tax");
         /*solium-disable-next-line security/no-block-members*/
         emit SetWithholdingFixed(_investors, _withholding);
@@ -139,7 +139,7 @@ contract DividendCheckpoint is DividendCheckpointStorage, ICheckpoint, Module {
         address payable[] memory _payees
     )
         public
-        withPerm(DISTRIBUTE)
+        withPerm(OPERATOR)
         validDividendIndex(_dividendIndex)
     {
         Dividend storage dividend = dividends[_dividendIndex];
@@ -161,7 +161,7 @@ contract DividendCheckpoint is DividendCheckpointStorage, ICheckpoint, Module {
         uint256 _start,
         uint256 _iterations
     ) public
-      withPerm(DISTRIBUTE)
+      withPerm(OPERATOR)
       validDividendIndex(_dividendIndex)
     {
         Dividend storage dividend = dividends[_dividendIndex];
@@ -391,8 +391,8 @@ contract DividendCheckpoint is DividendCheckpointStorage, ICheckpoint, Module {
      */
     function getPermissions() public view returns(bytes32[] memory) {
         bytes32[] memory allPermissions = new bytes32[](2);
-        allPermissions[0] = DISTRIBUTE;
-        allPermissions[1] = MANAGE;
+        allPermissions[0] = ADMIN;
+        allPermissions[1] = OPERATOR;
         return allPermissions;
     }
 

contracts/modules/Checkpoint/ERC20DividendCheckpoint.sol

@@ -56,7 +56,7 @@ contract ERC20DividendCheckpoint is ERC20DividendCheckpointStorage, DividendChec
         bytes32 _name
     )
         external
-        withPerm(MANAGE)
+        withPerm(ADMIN)
     {
         createDividendWithExclusions(_maturity, _expiry, _token, _amount, excluded, _name);
     }
@@ -79,7 +79,7 @@ contract ERC20DividendCheckpoint is ERC20DividendCheckpointStorage, DividendChec
         bytes32 _name
     )
         external
-        withPerm(MANAGE)
+        withPerm(ADMIN)
     {
         _createDividendWithCheckpointAndExclusions(_maturity, _expiry, _token, _amount, _checkpointId, excluded, _name);
     }
@@ -102,7 +102,7 @@ contract ERC20DividendCheckpoint is ERC20DividendCheckpointStorage, DividendChec
         bytes32 _name
     )
         public
-        withPerm(MANAGE)
+        withPerm(ADMIN)
     {
         uint256 checkpointId = ISecurityToken(securityToken).createCheckpoint();
         _createDividendWithCheckpointAndExclusions(_maturity, _expiry, _token, _amount, checkpointId, _excluded, _name);
@@ -128,7 +128,7 @@ contract ERC20DividendCheckpoint is ERC20DividendCheckpointStorage, DividendChec
         bytes32 _name
     )
         public
-        withPerm(MANAGE)
+        withPerm(ADMIN)
     {
         _createDividendWithCheckpointAndExclusions(_maturity, _expiry, _token, _amount, _checkpointId, _excluded, _name);
     }
@@ -251,7 +251,7 @@ contract ERC20DividendCheckpoint is ERC20DividendCheckpointStorage, DividendChec
      * @notice Issuer can reclaim remaining unclaimed dividend amounts, for expired dividends
      * @param _dividendIndex Dividend to reclaim
      */
-    function reclaimDividend(uint256 _dividendIndex) external withPerm(MANAGE) {
+    function reclaimDividend(uint256 _dividendIndex) external withPerm(OPERATOR) {
         require(_dividendIndex < dividends.length, "Invalid dividend");
         /*solium-disable-next-line security/no-block-members*/
         require(now >= dividends[_dividendIndex].expiry, "Dividend expiry in future");
@@ -267,7 +267,7 @@ contract ERC20DividendCheckpoint is ERC20DividendCheckpointStorage, DividendChec
      * @notice Allows issuer to withdraw withheld tax
      * @param _dividendIndex Dividend to withdraw from
      */
-    function withdrawWithholding(uint256 _dividendIndex) external withPerm(MANAGE) {
+    function withdrawWithholding(uint256 _dividendIndex) external withPerm(OPERATOR) {
         require(_dividendIndex < dividends.length, "Invalid dividend");
         Dividend storage dividend = dividends[_dividendIndex];
         uint256 remainingWithheld = dividend.totalWithheld.sub(dividend.totalWithheldWithdrawn);