From 347263e2d177229e57342534cbefda0704331844 Mon Sep 17 00:00:00 2001 From: Artjom Galaktionov Date: Thu, 11 Jan 2024 08:08:21 +0200 Subject: [PATCH] Vesting Refactor --- contracts/finance/Vesting.sol | 422 +++++++++++++++++----------------- 1 file changed, 213 insertions(+), 209 deletions(-) diff --git a/contracts/finance/Vesting.sol b/contracts/finance/Vesting.sol index b516d70f..3955bd7a 100644 --- a/contracts/finance/Vesting.sol +++ b/contracts/finance/Vesting.sol @@ -1,288 +1,292 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.4; -import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; -import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; -import {Address} from "@openzeppelin/contracts/utils/Address.sol"; -import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol"; -import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol"; -import {EnumerableSet} from "../libs/arrays/SetHelper.sol"; +import {OwnableUpgradeable, Initializable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol"; +import {SafeERC20, IERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; +import {MathUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/math/MathUpgradeable.sol"; -// should we move revoke to a separate contract as extension/preset -// should we allow the owner of the contract configure cliff -// should we make this contract ownable or make it as a separate preset -// add additional exetenstions with different formula of vesting (linear, exponential, etc) -abstract contract VestingWallet is Initializable { - using EnumerableSet for EnumerableSet.AddressSet; +abstract contract VestingWallet is Initializable, OwnableUpgradeable { + using MathUpgradeable for uint256; + using SafeERC20 for IERC20; - event BeneficiaryAdded(address account, uint256 shares); + // examples of vesting types + enum VestingScheduleType { + VAULT + // ANGELROUND, + // SEEDROUND, + // PRIVATEROUND, + // LISTINGS, + // GROWTH, + // OPERATIONAL, + // FOUNDERS, + // DEVELOPERS, + // BUGFINDING, + } - event AssetReleased(address indexed beneficiary, address indexed token, uint256 amount); - event AssetRevoked(address indexed beneficiary, address indexed token, uint256 amount); + // structure of vesting object + struct Vesting { + bool isActive; + address beneficiary; + uint256 totalAmount; + VestingScheduleType vestingScheduleType; + uint256 paidAmount; + bool isRevocable; + } - address public constant ETH = address(0); + // properties for linear vesting schedule + struct LinearVestingSchedule { + uint256 portionOfTotal; + uint256 startDate; + uint256 periodInSeconds; + uint256 portionPerPeriod; + uint256 cliffInPeriods; + } - EnumerableSet.AddressSet private _beneficiaries; + uint256 public constant SECONDS_IN_MONTH = 60 * 60 * 24 * 30; + uint256 public constant PORTION_OF_TOTAL_PRECISION = 10 ** 10; + uint256 public constant PORTION_PER_PERIOD_PRECISION = 10 ** 10; - uint64 private _cliff; - uint64 private _start; - uint64 private _duration; + uint256 public _activationTimestamp; - bool private _revocable; + IERC20 public _vestingToken; - uint256 private _totalShares; + Vesting[] public _vestings; + uint256 public _amountInVestings; + mapping(VestingScheduleType => LinearVestingSchedule[]) public _vestingSchedules; - struct VestingData { - uint256 shares; - mapping(address asset => AssetInfo) assetsInfo; - } + event VestingTokenSet(IERC20 token); + event VestingAdded(uint256 vestingId, address beneficiary); + event VestingRevoked(uint256 vestingId); + event VestingWithdraw(uint256 vestingId, uint256 amount); - struct AssetInfo { - bool isRevoked; - uint256 releasedAmount; + // initialization + function __VestingWallet_init(uint256 activationTimestamp_) internal onlyInitializing { + __Ownable_init(); + + _activationTimestamp = activationTimestamp_; + + _initializeVestingSchedules(); } - mapping(address beneficiary => VestingData) private _vestingData; - mapping(address asset => uint256) private _totalReleased; - - function __VestingWallet_init( - address[] memory beneficiaries_, - uint256[] memory shares_, - uint64 startTimestamp_, - uint64 cliffSeconds_, - uint64 durationSeconds_, - bool revocable_ - ) internal onlyInitializing { - require( - beneficiaries_.length == shares_.length, - "Vesting: beneficiaries and shares length mismatch" - ); - require(beneficiaries_.length > 0, "Vesting: no beneficiaries"); - require(cliffSeconds_ <= durationSeconds_, "Vesting: cliff is longer than duration"); - require( - startTimestamp_ + durationSeconds_ > block.timestamp, - "Vesting: final time is before current time" - ); + // get available amount to withdraw + function getWithdrawableAmount(uint256 vestingId_) external view virtual returns (uint256) { + Vesting memory _vesting = getVesting(vestingId_); - for (uint256 i = 0; i < beneficiaries_.length; i++) { - _addBeneficiary(beneficiaries_[i], shares_[i]); - } + require(_vesting.isActive, "VestingWallet: vesting is canceled"); - _start = startTimestamp_; - _duration = durationSeconds_; - _cliff = startTimestamp_ + cliffSeconds_; - _revocable = revocable_; + return _getWithdrawableAmount(_vesting); } - receive() external payable virtual {} + // withdraw available tokens from multiple vestings + function withdrawFromVestingBulk(uint256 offset_, uint256 limit_) external virtual { + uint256 _to = (offset_ + limit_).min(_vestings.length).max(offset_); - function releasable(address account_) public view virtual returns (uint256) { - return releasable(account_, ETH); + for (uint256 i = offset_; i < _to; i++) { + Vesting storage vesting = _getVesting(i); + if (vesting.isActive) { + _withdrawFromVesting(vesting, i); + } + } } - function releasable(address account_, address token_) public view virtual returns (uint256) { - return - vestedAmount(account_, token_, uint64(block.timestamp)) - released(account_, token_); - } + // withdraw available tokens from vesting + function withdrawFromVesting(uint256 vestingId_) external virtual { + Vesting storage _vesting = _getVesting(vestingId_); - function vestedAmount( - address account_, - uint64 timestamp_ - ) public view virtual returns (uint256) { - return vestedAmount(account_, ETH, timestamp_); - } + require(_vesting.isActive, "VestingWallet: vesting is canceled"); - function vestedAmount( - address account_, - address token_, - uint64 timestamp_ - ) public view virtual returns (uint256) { - return - _vestingSchedule( - account_, - token_, - beneficiaryAllocation(account_, token_), - timestamp_ - ); + _withdrawFromVesting(_vesting, vestingId_); } - function beneficiaryAllocation(address account_) public view virtual returns (uint256) { - return beneficiaryAllocation(account_, ETH); + // get vesting info by vesting id + function getVesting(uint256 vestingId_) public view virtual returns (Vesting memory) { + return _getVesting(vestingId_); } - function beneficiaryAllocation( - address account_, - address token_ - ) public view virtual returns (uint256) { - return (totalAllocation(token_) * shares(account_)) / totalShares(); + // get amount that is present on the contract but not allocated to vesting + function getAvailableTokensAmount() public view virtual returns (uint256) { + return _vestingToken.balanceOf(address(this)) - (_amountInVestings); } - function totalAllocation() public view virtual returns (uint256) { - return totalAllocation(ETH); + // initilizes default vesting schedules (here used as an example) + function _initializeVestingSchedules() internal virtual { + _addLinearVestingSchedule( + VestingScheduleType.VAULT, + LinearVestingSchedule({ + portionOfTotal: PORTION_OF_TOTAL_PRECISION, + startDate: _activationTimestamp, + periodInSeconds: SECONDS_IN_MONTH, + portionPerPeriod: PORTION_PER_PERIOD_PRECISION / 2, + cliffInPeriods: 1 + }) + ); } - function totalAllocation(address token_) public view virtual returns (uint256) { - return - (token_ == ETH ? address(this).balance : IERC20(token_).balanceOf(address(this))) + - totalReleased(token_); + // add linear vesting schedule with your own properties + function _addLinearVestingSchedule( + VestingScheduleType type_, + LinearVestingSchedule memory schedule_ + ) internal virtual onlyOwner { + _vestingSchedules[type_].push(schedule_); } - function totalReleased() public view virtual returns (uint256) { - return totalReleased(ETH); - } + // set vesting token + function _setVestingToken(IERC20 token_) internal virtual onlyOwner { + require(address(token_) == address(0), "VestingWallet: token is already set"); - function totalReleased(address token_) public view virtual returns (uint256) { - return _totalReleased[token_]; - } + _vestingToken = token_; - function shares(address account_) public view virtual returns (uint256) { - return _vestingData[account_].shares; + emit VestingTokenSet(token_); } - function totalShares() public view virtual returns (uint256) { - return _totalShares; - } + // create vesting for multiple beneficiaries + function _createVestingBulk( + address[] calldata beneficiaries_, + uint256[] calldata amounts_, + VestingScheduleType[] calldata vestingSchedules_, + bool[] calldata isRevokable_ + ) internal virtual onlyOwner { + require( + beneficiaries_.length == amounts_.length && + beneficiaries_.length == vestingSchedules_.length && + beneficiaries_.length == isRevokable_.length, + "VestingWallet: parameters length mismatch" + ); - function released(address account_) public view virtual returns (uint256) { - return released(account_, ETH); + for (uint256 i = 0; i < beneficiaries_.length; i++) { + _createVesting(beneficiaries_[i], amounts_[i], vestingSchedules_[i], isRevokable_[i]); + } } - function released(address account_, address token_) public view virtual returns (uint256) { - return _vestingData[account_].assetsInfo[token_].releasedAmount; - } + // create vesting + function _createVesting( + address beneficiary_, + uint256 amount_, + VestingScheduleType vestingSchedule_, + bool isRevokable_ + ) internal virtual onlyOwner returns (uint256 _vestingId) { + require( + getAvailableTokensAmount() >= amount_, + "VestingWallet: not enough tokens in vesting contract" + ); + require( + beneficiary_ != address(0), + "VestingWallet: cannot create vesting for zero address" + ); + require(amount_ > 0, "VestingWallet: cannot create vesting for zero amount"); - function cliff() public view virtual returns (uint256) { - return _cliff; - } + _amountInVestings += amount_; - function start() public view virtual returns (uint256) { - return _start; - } + _vestingId = _vestings.length; - function duration() public view virtual returns (uint256) { - return _duration; - } + _vestings.push( + Vesting({ + isActive: true, + beneficiary: beneficiary_, + totalAmount: amount_, + vestingScheduleType: vestingSchedule_, + paidAmount: 0, + isRevocable: isRevokable_ + }) + ); - function end() public view virtual returns (uint256) { - return start() + duration(); + emit VestingAdded(_vestingId, beneficiary_); } - function isBeneficiary(address account_) public view virtual returns (bool) { - return _beneficiaries.contains(account_); - } + // revoke vesting and release locked tokens + function _revokeVesting(uint256 vestingId_) internal virtual onlyOwner { + Vesting storage _vesting = _getVesting(vestingId_); - function revocable() public view virtual returns (bool) { - return _revocable; - } + require(_vesting.isActive, "VestingWallet: vesting is revoked"); + require(_vesting.isRevocable, "VestingWallet: vesting is not revokable"); - function revoked(address account_) public view virtual returns (bool) { - return revoked(account_, ETH); - } + _vesting.isActive = false; - function revoked(address account_, address token_) public view virtual returns (bool) { - return _vestingData[account_].assetsInfo[token_].isRevoked; - } + uint256 _amountReleased = _vesting.totalAmount - _vesting.paidAmount; + _amountInVestings -= _amountReleased; - function vestingData( - address account_ - ) public view virtual returns (uint256 _shares, AssetInfo memory _assetInfo) { - (_shares, _assetInfo) = vestingData(account_, ETH); + emit VestingRevoked(vestingId_); } - function vestingData( - address account_, - address token_ - ) public view virtual returns (uint256 _shares, AssetInfo memory _assetInfo) { - VestingData storage _accountVesting = _vestingData[account_]; + // withdraw tokens from vesting and transfer to beneficiary + function _withdrawFromVesting(Vesting storage vesting_, uint256 vestingId_) internal virtual { + uint256 _amountToPay = _getWithdrawableAmount(vesting_); - _shares = _accountVesting.shares; - _assetInfo = _accountVesting.assetsInfo[token_]; - } + vesting_.paidAmount += _amountToPay; + _amountInVestings -= _amountToPay; - function _vestingSchedule( - address account_, - uint256 totalAllocation_, - uint64 timestamp_ - ) internal view virtual returns (uint256) { - return _vestingSchedule(account_, ETH, totalAllocation_, timestamp_); - } + _vestingToken.safeTransfer(vesting_.beneficiary, _amountToPay); - function _vestingSchedule( - address account_, - address token_, - uint256 totalAllocation_, - uint64 timestamp_ - ) internal view virtual returns (uint256) { - if (timestamp_ < cliff()) { - return 0; - } else if (timestamp_ >= end() || revoked(account_, token_)) { - return totalAllocation_; - } else { - return (totalAllocation_ * (timestamp_ - start())) / duration(); - } + emit VestingWithdraw(vestingId_, _amountToPay); } - function _release(address account_) internal virtual { - _release(account_, ETH); + // get available amount to withdraw + function _getWithdrawableAmount( + Vesting memory _vesting + ) internal view virtual returns (uint256) { + return _calculateReleasableAmount(_vesting) - _vesting.paidAmount; } - function _release(address account_, address token_) internal virtual { - require(_beneficiaries.contains(account_), "Vesting: not a beneficiary"); - - VestingData storage _accountVesting = _vestingData[account_]; + // calculate releasable amount at the moment + function _calculateReleasableAmount( + Vesting memory vesting_ + ) internal view virtual returns (uint256) { + LinearVestingSchedule[] storage _vestingSchedulesByType = _vestingSchedules[ + vesting_.vestingScheduleType + ]; - require(_accountVesting.shares > 0, "Vesting: account has no shares"); + uint256 _releasableAmount; - uint256 _amount = releasable(token_); + for (uint256 i = 0; i < _vestingSchedulesByType.length; i++) { + LinearVestingSchedule storage _vestingSchedule = _vestingSchedulesByType[i]; - _accountVesting.assetsInfo[token_].releasedAmount += _amount; - _totalReleased[token_] += _amount; + if (_vestingSchedule.startDate > block.timestamp) return _releasableAmount; - token_ == ETH - ? Address.sendValue(payable(account_), _amount) - : SafeERC20.safeTransfer(IERC20(token_), account_, _amount); + uint256 _releasableAmountForThisSchedule = _calculateLinearVestingAvailableAmount( + _vestingSchedule, + vesting_.totalAmount + ); - emit AssetReleased(account_, token_, _amount); - } + _releasableAmount += _releasableAmountForThisSchedule; + } - function _revoke(address account_) internal virtual { - _revoke(account_, ETH); + return _releasableAmount; } - function _revoke(address account_, address token_) internal virtual { - require(_beneficiaries.contains(account_), "Vesting: not a beneficiary"); - require(revocable(), "Vesting: cannot revoke"); - - VestingData storage _accountVesting = _vestingData[account_]; - - require(_accountVesting.shares > 0, "Vesting: account has no shares"); - - AssetInfo storage _accountAssetInfo = _accountVesting.assetsInfo[token_]; - - require(!_accountAssetInfo.isRevoked, "Vesting: already revoked"); + // calculation of linear vesting + function _calculateLinearVestingAvailableAmount( + LinearVestingSchedule storage vestingSchedule_, + uint256 amount_ + ) internal view virtual returns (uint256) { + uint256 _elapsedPeriods = _calculateElapsedPeriods(vestingSchedule_); - uint256 _amount = beneficiaryAllocation(account_, token_) - releasable(account_, token_); + if (_elapsedPeriods <= vestingSchedule_.cliffInPeriods) return 0; - _accountAssetInfo.isRevoked = true; + uint256 _amountThisVestingSchedule = (amount_ * vestingSchedule_.portionOfTotal) / + (PORTION_OF_TOTAL_PRECISION); - token_ == ETH - ? Address.sendValue(payable(account_), _amount) - : SafeERC20.safeTransfer(IERC20(token_), account_, _amount); + uint256 _amountPerPeriod = (_amountThisVestingSchedule * + vestingSchedule_.portionPerPeriod) / (PORTION_PER_PERIOD_PRECISION); - emit AssetRevoked(account_, token_, _amount); + return (_amountPerPeriod * _elapsedPeriods).min(_amountThisVestingSchedule); } - function _addBeneficiary(address account_, uint256 shares_) private { - VestingData storage _accountVesting = _vestingData[account_]; + // withdraw tokens that is present on the contract but not allocated to vesting + function _withdrawExcessiveTokens() internal virtual onlyOwner { + _vestingToken.safeTransfer(owner(), getAvailableTokensAmount()); + } - require(account_ != ETH, "Vesting: account is the zero address"); - require(shares_ > 0, "Shares: shares are 0"); - require(_accountVesting.shares == 0, "Shares: account already has shares"); + // get vesting info + function _getVesting(uint256 _vestingId) internal view virtual returns (Vesting storage) { + require(_vestingId < _vestings.length, "VestingWallet: no vesting with such id"); - _beneficiaries.add(account_); - _accountVesting.shares = shares_; - _totalShares += shares_; + return _vestings[_vestingId]; + } - emit BeneficiaryAdded(account_, shares_); + // calculate elapsed periods + function _calculateElapsedPeriods( + LinearVestingSchedule storage _linearVesting + ) private view returns (uint256) { + return (block.timestamp - _linearVesting.startDate) / (_linearVesting.periodInSeconds); } }