change cec distributor

contracts/activity/CECDistributor.sol

@@ -16,40 +16,35 @@ import {Governable} from "../core/Governable.sol";
 contract CECDistributor is ReentrancyGuard, Pausable, Ownable, Governable {
   using SafeERC20 for IERC20;
 
-  mapping(address account => uint256 amount) public balanceMap;
+  struct ReleaseInfo {
-  mapping(address account => uint256 amount) public releaseMap;
+    uint256 amount;
+    uint256 releaseAllMonth;
+    uint256 tgeRatio;
+    uint256 lockDuration;
+  }
+  
+  mapping(address account => ReleaseInfo info) public infoMap;
+  mapping(address account => uint256 amount) public releasedMap;
 
   string public name;
   IERC20 public immutable cecToken;
   uint256 public constant DURATION = 86400 * 30;
-  uint256 private releaseAllMonth;
 
   uint256 public start = 0;
-  // release ratio when tge
-  uint256 public tgeRatio;
   uint256 public constant TGE_PRECISION = 1000000;
-  //
   uint256 public lockDuration;
-  address public wallet;
 
-  event EventBalanceUpdated(address indexed account, uint256 amount);
+  event EventInfoUpdated(address indexed account, ReleaseInfo info);
   event EventCECClaimed(address indexed user, address indexed to, uint256 amount);
   event EventChangeAddress(address oldAddr, address newAddr);
 
+  
   constructor(
     string memory _name,
-    address _cecToken, 
+    address _cecToken
-    address _wallet, 
-    uint256 _lockDuration, 
-    uint256 _releaseAllMonth, 
-    uint256 _tgeRatio
   ) {
     name = _name;
     cecToken = IERC20(_cecToken);
-    wallet = _wallet;
-    lockDuration = _lockDuration;
-    releaseAllMonth = _releaseAllMonth;
-    tgeRatio = _tgeRatio;
   }
 
   /**
@@ -84,32 +79,37 @@ contract CECDistributor is ReentrancyGuard, Pausable, Ownable, Governable {
     start = newStart;
   }
 
-  function updateBalances(address[] calldata accounts, uint256[] calldata amounts) external onlyOwner {
+  function withdrawToken(address to, uint256 amount) external onlyOwner {
-    require(accounts.length == amounts.length, "CECDistributor: invalid input");
+    cecToken.safeTransfer(to, amount);
+  }
+
+  function updateInfo(address[] calldata accounts, ReleaseInfo[] calldata infos) external onlyOwner {
+    require(accounts.length == infos.length, "CECDistributor: invalid input");
     for (uint256 i = 0; i < accounts.length; i++) {
-      balanceMap[accounts[i]] = amounts[i];
+      infoMap[accounts[i]] = infos[i];
-      emit EventBalanceUpdated(accounts[i], amounts[i]);
+      emit EventInfoUpdated(accounts[i], infos[i]);
     }
   }
 
   function calcClaimAmount(address user) public view whenNotPaused returns (uint256) {
-    require(balanceMap[user] > 0, "CECDistributor: not in whitelist");
+    require(infoMap[user].amount > 0, "CECDistributor: not in whitelist");
     require(block.timestamp >= start, "CECDistributor: not in claim time");
     uint256 claimAmount = 0;
     uint256 tgeAmount = 0;
-    if (tgeRatio > 0) {
+    ReleaseInfo memory info = infoMap[user];
-      tgeAmount = ((balanceMap[user] * tgeRatio) / TGE_PRECISION);
+    if (info.tgeRatio > 0) {
+      tgeAmount = ((info.amount * info.tgeRatio) / TGE_PRECISION);
       claimAmount += tgeAmount;
     }
-    if (block.timestamp > start + lockDuration) {
+    if (block.timestamp > start + info.lockDuration) {
-      uint256 monthNum = (block.timestamp - start - lockDuration) / DURATION;
+      uint256 monthNum = (block.timestamp - start - info.lockDuration) / DURATION;
-      if (monthNum <= releaseAllMonth) {
+      if (monthNum <= info.releaseAllMonth) {
-        claimAmount += (((balanceMap[user] - tgeAmount) * monthNum) / releaseAllMonth);
+        claimAmount += (((info.amount - tgeAmount) * monthNum) / info.releaseAllMonth);
       } else {
-        claimAmount = balanceMap[user];
+        claimAmount = info.amount;
       }
     }
-    claimAmount -= releaseMap[user];
+    claimAmount -= releasedMap[user];
     return claimAmount;
   }
 
@@ -119,22 +119,22 @@ contract CECDistributor is ReentrancyGuard, Pausable, Ownable, Governable {
     address _user = _msgSender();
     uint256 amount = calcClaimAmount(_user);
     if (amount > 0) {
-      releaseMap[_user] = amount;
+      releasedMap[_user] = amount;
-      cecToken.safeTransferFrom(wallet, to, amount);
+      cecToken.safeTransfer(to, amount);
       emit EventCECClaimed(_user, to, amount);
     }
     return amount;
   }
 
   function changeAddress(address from, address to) external {
-    require(balanceMap[to] == 0, "CECDistributor: new addr is in whitelist");
+    require(infoMap[to].amount == 0, "CECDistributor: new addr is in whitelist");
-    require(balanceMap[from] > 0, "CECDistributor: not in whitelist");
+    require(infoMap[from].amount > 0, "CECDistributor: not in whitelist");
     address _sender = _msgSender();
     require(_sender == owner() || _sender == gov || _sender == from, "CECDistributor: sender not allowed");
-    balanceMap[to] = balanceMap[from];
+    infoMap[to] = infoMap[from];
-    balanceMap[from] = 0;
+    delete infoMap[from];
-    releaseMap[to] = releaseMap[from];
+    releasedMap[to] = releasedMap[from];
-    releaseMap[from] = 0;
+    releasedMap[from] = 0;
     emit EventChangeAddress(from, to);
   }
 }

deploy/10_deploy_cecdistributor.ts

@@ -13,7 +13,8 @@ const deployCECDistributor: DeployFunction =
     const { cec } = config.staking
     // update distributor wallet  
     const wallet = admin;
-    const params: any[] = [ 'first', cec, wallet,  ONE_DAY * 3, 10, 300000]
+    const params: any[] = [ 'first', cec, wallet,  0, 12, 100000]
+    // const params: any[] = [ 'first', cec, wallet,  0, 10, 500000]
     const ret = await hre.deployments.deploy("CECDistributor", {
       from,
       args: params,

test/testCECDistributor.ts

@@ -19,18 +19,17 @@ describe('TestCECDistributor', function() {
     const Cec = await hre.ethers.getContractFactory("MintableBaseToken");
     const cec = await Cec.deploy("test cec", "cec");
     await cec.setMinter(owner.address, true)
-    await cec.mint(user0.address, expandDecimals(15000, 18))
 
     const CECDistributor = await hre.ethers.getContractFactory("CECDistributor");
     const lockDuration = ONE_DAY; // one day
-    const distributor = await CECDistributor.deploy("first", cec.target, user0.address, lockDuration, 10, 300000 );
+    const distributor = await CECDistributor.deploy("first", cec.target);
-    //@ts-ignore
+    await cec.mint(distributor.target, expandDecimals(15000, 18))
-    await cec.connect(user0).approve(distributor.target, expandDecimals(10000, 18))
+
     const chainId = hre.network.config.chainId
     const start = (Date.now() / 1000 + 3600) | 0 // one hour later
     await distributor.setStart(start)
     expect(await distributor.name()).to.equal("first")
-    await distributor.updateBalances([user1.address], [expandDecimals(1000, 18)])
+    await distributor.updateInfo([user1.address], [[expandDecimals(1000, 18), 10, 300000, lockDuration]])
     return { distributor, owner, user0, user1, user2, chainId, cec, start };
   }
   describe("Deployment", function () {