Add a VestingWallet

contracts/finance/VestingWallet.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.0;
 import "../token/ERC20/utils/SafeERC20.sol";
 import "../utils/Address.sol";
 import "../utils/Context.sol";
+import "../utils/math/Math.sol";
 
 /**
  * @title VestingWallet
@@ -16,9 +17,11 @@ import "../utils/Context.sol";
  * be immediately releasable.
  */
 contract VestingWallet is Context {
-    event TokensReleased(address token, uint256 amount);
+    event TokensReleased(uint256 amount);
+    event ERC20TokensReleased(address token, uint256 amount);
 
-    mapping(address => uint256) private _released;
+    uint256 private _released;
+    mapping(address => uint256) private _erc20Released;
     address private immutable _beneficiary;
     uint256 private immutable _start;
     uint256 private immutable _duration;
@@ -69,26 +72,17 @@ contract VestingWallet is Context {
     }
 
     /**
-     * @dev Amont of token already released
+     * @dev Amont of eth already released
      */
-    function released(address token) public view returns (uint256) {
-        return _released[token];
+    function released() public view returns (uint256) {
+        return _released;
     }
 
     /**
-     * @dev Release the tokens that have already vested.
-     *
-     * Emits a {TokensReleased} event.
+     * @dev Amont of token already released
      */
-    function release(address token) public virtual {
-        uint256 releasable = vestedAmount(token, block.timestamp) - released(token);
-        _released[token] += releasable;
-        emit TokensReleased(token, releasable);
-        if (token == address(0)) {
-            Address.sendValue(payable(beneficiary()), releasable);
-        } else {
-            SafeERC20.safeTransfer(IERC20(token), beneficiary(), releasable);
-        }
+    function released(address token) public view returns (uint256) {
+        return _erc20Released[token];
     }
 
     /**
@@ -97,33 +91,45 @@ contract VestingWallet is Context {
      * Emits a {TokensReleased} event.
      */
     function release() public virtual {
-        release(address(0));
+        uint256 releasable = vestedAmount(block.timestamp) - released();
+        _released += releasable;
+        emit TokensReleased(releasable);
+        Address.sendValue(payable(beneficiary()), releasable);
     }
 
     /**
-     * @dev Calculates the amount of tokens that has already vested. Default implementation is a linear vesting curve.
+     * @dev Release the tokens that have already vested.
+     *
+     * Emits a {TokensReleased} event.
      */
-    function vestedAmount(address token, uint256 timestamp) public view virtual returns (uint256) {
-        if (timestamp < start()) {
-            return 0;
-        } else if (timestamp >= start() + duration()) {
-            return _historicalBalance(token);
-        } else {
-            return (_historicalBalance(token) * (timestamp - start())) / duration();
-        }
+    function release(address token) public virtual {
+        uint256 releasable = vestedAmount(token, block.timestamp) - released(token);
+        _erc20Released[token] += releasable;
+        emit ERC20TokensReleased(token, releasable);
+        SafeERC20.safeTransfer(IERC20(token), beneficiary(), releasable);
     }
 
     /**
      * @dev Calculates the amount of ether that has already vested. Default implementation is a linear vesting curve.
      */
     function vestedAmount(uint256 timestamp) public view virtual returns (uint256) {
-        return vestedAmount(address(0), timestamp);
+        if (timestamp < start()) {
+            return 0;
+        } else {
+            uint256 historicalBalance = address(this).balance + released();
+            return Math.min(historicalBalance, historicalBalance * (timestamp - start()) / duration());
+        }
     }
 
     /**
-     * @dev Calculates the historical balance (current balance + already released balance).
+     * @dev Calculates the amount of tokens that has already vested. Default implementation is a linear vesting curve.
      */
-    function _historicalBalance(address token) internal view virtual returns (uint256) {
-        return (token == address(0) ? address(this).balance : IERC20(token).balanceOf(address(this))) + released(token);
+    function vestedAmount(address token, uint256 timestamp) public view virtual returns (uint256) {
+        if (timestamp < start()) {
+            return 0;
+        } else {
+            uint256 historicalBalance = IERC20(token).balanceOf(address(this)) + released(token);
+            return Math.min(historicalBalance, historicalBalance * (timestamp - start()) / duration());
+        }
     }
 }

test/finance/VestingWallet.behavior.js

@@ -20,10 +20,15 @@ function shouldBehaveLikeVesting (beneficiary) {
     {
       const receipt = await this.mock.release(...args);
 
-      await expectEvent.inTransaction(receipt.tx, this.mock, 'TokensReleased', {
-        token: this.token ? this.token.address : constants.ZERO_ADDRESS,
-        amount: '0',
-      });
+      await expectEvent.inTransaction(
+        receipt.tx,
+        this.mock,
+        this.token ? 'ERC20TokensReleased' : 'TokensReleased',
+        Object.fromEntries(Object.entries({
+          token: this.token && this.token.address,
+          amount: '0',
+        }).filter(x => x.every(Boolean))),
+      );
 
       await this.checkRelease(receipt, beneficiary, '0');
 
@@ -40,10 +45,15 @@ function shouldBehaveLikeVesting (beneficiary) {
 
       const receipt = await this.mock.release(...args);
 
-      await expectEvent.inTransaction(receipt.tx, this.mock, 'TokensReleased', {
-        token: this.token ? this.token.address : constants.ZERO_ADDRESS,
-        amount: vested.sub(released),
-      });
+      await expectEvent.inTransaction(
+        receipt.tx,
+        this.mock,
+        this.token ? 'ERC20TokensReleased' : 'TokensReleased',
+        Object.fromEntries(Object.entries({
+          token: this.token && this.token.address,
+          amount: vested.sub(released),
+        }).filter(x => x.every(Boolean))),
+      );
 
       await this.checkRelease(receipt, beneficiary, vested.sub(released));