diff --git a/contracts/token/TokenVesting.sol b/contracts/token/TokenVesting.sol index a85f02dfe..c201a4770 100644 --- a/contracts/token/TokenVesting.sol +++ b/contracts/token/TokenVesting.sol @@ -57,14 +57,15 @@ contract TokenVesting is Ownable { */ function release(ERC20Basic token) public { uint256 vested = vestedAmount(token); + uint256 unreleased = releasableAmount(token); - require(vested > 0); + require(unreleased > 0); - token.safeTransfer(beneficiary, vested); + token.safeTransfer(beneficiary, unreleased); - released[token] = released[token].add(vested); + released[token] = released[token].add(unreleased); - Released(vested); + Released(unreleased); } /** @@ -78,12 +79,12 @@ contract TokenVesting is Ownable { uint256 balance = token.balanceOf(this); - uint256 vested = vestedAmount(token); - uint256 vesting = balance - vested; + uint256 unreleased = releasableAmount(token); + uint256 refund = balance.sub(unreleased); revoked[token] = true; - token.safeTransfer(owner, vesting); + token.safeTransfer(owner, refund); Revoked(); } @@ -92,20 +93,24 @@ contract TokenVesting is Ownable { * @dev Calculates the amount that has already vested but hasn't been released yet. * @param token ERC20 token which is being vested */ + function releasableAmount(ERC20Basic token) public constant returns (uint256) { + return vestedAmount(token).sub(released[token]); + } + + /** + * @dev Calculates the amount that has already vested. + * @param token ERC20 token which is being vested + */ function vestedAmount(ERC20Basic token) public constant returns (uint256) { + uint256 currentBalance = token.balanceOf(this); + uint256 totalBalance = currentBalance.add(released[token]); + if (now < cliff) { return 0; } else if (now >= start + duration || revoked[token]) { - return token.balanceOf(this); + return totalBalance; } else { - uint256 currentBalance = token.balanceOf(this); - uint256 totalBalance = currentBalance.add(released[token]); - - uint256 vested = totalBalance.mul(now - start).div(duration); - uint256 unreleased = vested.sub(released[token]); - - // currentBalance can be 0 in case of vesting being revoked earlier. - return Math.min256(currentBalance, unreleased); + return totalBalance.mul(now - start).div(duration); } } } diff --git a/test/TokenVesting.js b/test/TokenVesting.js index c391c6577..6a4d86206 100644 --- a/test/TokenVesting.js +++ b/test/TokenVesting.js @@ -83,12 +83,11 @@ contract('TokenVesting', function ([_, owner, beneficiary]) { await increaseTimeTo(this.start + this.cliff + duration.weeks(12)); const vested = await this.vesting.vestedAmount(this.token.address); - const balance = await this.token.balanceOf(this.vesting.address); await this.vesting.revoke(this.token.address, { from: owner }); const ownerBalance = await this.token.balanceOf(owner); - ownerBalance.should.bignumber.equal(balance.sub(vested)); + ownerBalance.should.bignumber.equal(amount.sub(vested)); }); it('should keep the vested tokens when revoked by owner', async function () {