Multi round vesting wallet example

MultiRoundVestingWallet.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import {IERC20} from "../token/ERC20/IERC20.sol";
import {SafeERC20} from "../token/ERC20/utils/SafeERC20.sol";
import {Address} from "../utils/Address.sol";
import {Context} from "../utils/Context.sol";

contract MultiRoundVestingWallet is Context {
    using SafeERC20 for IERC20;

    struct InvestmentRound {
        IERC20 token;
        uint64 start;
        uint64 duration;
        uint256 amount;
    }

    event Invested(address investor, address token, uint256 amount, uint64 start, uint64 duration, uint256 roundIndex);
    event Released(address investor, address token, uint256 amount);

    mapping(address => mapping(address => InvestmentRound[])) private _investmentRounds;
    mapping(address => mapping(address => uint256)) private _releasedAmounts;

    function invest(IERC20 token, uint64 startTimestamp, uint64 durationSeconds, uint256 amount) external {
        require(durationSeconds > 0, "Duration cannot be zero");
        require(startTimestamp >= block.timestamp, "Start timestamp cannot be in the past");

        token.safeTransferFrom(msg.sender, address(this), amount);
        _investmentRounds[msg.sender][address(token)].push(InvestmentRound(token, startTimestamp, durationSeconds, amount));

        emit Invested(msg.sender, address(token), amount, startTimestamp, durationSeconds, _investmentRounds[msg.sender][address(token)].length - 1);
    }

    function vestedAmount(address investor, IERC20 token, uint64 timestamp) public view returns (uint256) {
        InvestmentRound[] storage rounds = _investmentRounds[investor][address(token)];
        uint256 total = 0;
        for (uint i = 0; i < rounds.length; i++) {
            InvestmentRound storage round = rounds[i];
            total += _vestingSchedule(round.amount, round.start, round.duration, timestamp);
        }
        return total;
    }

    function release(address investor, IERC20 token) public {
        InvestmentRound[] storage rounds = _investmentRounds[investor][address(token)];
        uint256 totalVested = 0;
        for (uint i = 0; i < rounds.length; i++) {
            InvestmentRound storage round = rounds[i];
            uint256 vested = _vestingSchedule(round.amount, round.start, round.duration, uint64(block.timestamp));
            totalVested += vested;
            if(vested > 0) round.amount -= vested; // update the amount left to vest
            if(round.amount == 0) delete rounds[i]; // remove the round if it is fully vested
        }

        uint256 previouslyReleased = _releasedAmounts[investor][address(token)];
        uint256 releasable = totalVested > previouslyReleased ? totalVested - previouslyReleased : 0;

        if (releasable > 0) {
            uint256 balance = token.balanceOf(address(this));
            require(balance >= releasable, "Not enough tokens available for release");

            _releasedAmounts[investor][address(token)] += releasable;
            token.safeTransfer(investor, releasable);
            emit Released(investor, address(token), releasable);
        }
    }

    function _vestingSchedule(uint256 totalAllocation, uint64 start, uint64 duration, uint64 timestamp) internal pure returns (uint256) {
        if (timestamp < start) {
            return 0;
        } else if (timestamp >= start + duration) {
            return totalAllocation;
        } else {
            return (totalAllocation * (timestamp - start)) / duration;
        }
    }
}