From 7520f30fab5a5824e310351ffd66250ffda5f0b4 Mon Sep 17 00:00:00 2001 From: sm-stack Date: Fri, 19 Jul 2024 10:47:38 +0900 Subject: [PATCH] feat: apply PR reviews --- .../contracts/contracts/L1/AssetManager.sol | 156 +++++++++--------- .../contracts/L1/interfaces/IAssetManager.sol | 2 +- 2 files changed, 80 insertions(+), 78 deletions(-) diff --git a/packages/contracts/contracts/L1/AssetManager.sol b/packages/contracts/contracts/L1/AssetManager.sol index d7f56275e..1f93bcd7a 100644 --- a/packages/contracts/contracts/L1/AssetManager.sol +++ b/packages/contracts/contracts/L1/AssetManager.sol @@ -343,6 +343,7 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager { uint128 assets ) external isRegistered(validator) returns (uint128) { if (assets == 0) revert NotAllowedZeroInput(); + ASSET_TOKEN.safeTransferFrom(owner, address(this), assets); uint128 shares = _delegate(validator, msg.sender, assets); emit KroDelegated(validator, msg.sender, assets, shares); return shares; @@ -669,6 +670,53 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager { emit RewardClaimFinalized(msg.sender, rewardsToClaim); } + /** + * @inheritdoc IAssetManager + */ + function claimKghReward(address validator, uint256[] calldata tokenIds) external { + if (validator == address(0)) revert ZeroAddress(); + if (tokenIds.length == 0) revert NotAllowedZeroInput(); + + Vault storage vault = _vaults[validator]; + KghDelegator storage kghDelegator = vault.kghDelegators[msg.sender]; + if (kghDelegator.kghNum < uint128(tokenIds.length)) revert InvalidTokenIdsInput(); + + // Auto compounding is not applied for claim function. + uint128 claimedRewards = _claimBoostedReward( + validator, + msg.sender, + uint128(tokenIds.length), + false + ); + + // Burn kroShares to claim the base reward. + uint128 kroSharesToBurn; + for (uint256 i = 0; i < tokenIds.length; ) { + uint128 kroShares = kghDelegator.kroShares[tokenIds[i]] - + _convertToKroShares(validator, KGH_MANAGER.totalKroInKgh(tokenIds[i])); + + unchecked { + kroSharesToBurn += kroShares; + + ++i; + } + } + + uint128 kghBaseReward = _convertToKroAssets(validator, kroSharesToBurn); + + unchecked { + claimedRewards += kghBaseReward; + vault.asset.totalKroShares -= kroSharesToBurn; + vault.asset.totalKro -= kghBaseReward; + } + + if (claimedRewards == 0) revert InsufficientAsset(); + + ASSET_TOKEN.safeTransfer(msg.sender, claimedRewards); + + emit KghRewardClaimed(validator, msg.sender, claimedRewards, kroSharesToBurn); + } + /** * @notice Bond KRO from validator KRO during output submission or challenge creation. This * function is only called by the ValidatorManager contract. @@ -829,81 +877,6 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager { } } - /** - * @inheritdoc IAssetManager - */ - function claimKghReward(address validator, uint256[] calldata tokenIds) external { - if (validator == address(0)) revert ZeroAddress(); - if (tokenIds.length == 0) revert NotAllowedZeroInput(); - - Vault storage vault = _vaults[validator]; - KghDelegator storage kghDelegator = vault.kghDelegators[msg.sender]; - if (kghDelegator.kghNum < uint128(tokenIds.length)) revert InvalidTokenIdsInput(); - - // Auto compounding is not applied for claim function. - uint128 claimedRewards = _claimBoostedReward(validator, msg.sender, uint128(tokenIds.length), false); - - // Burn kroShares to claim the base reward. - uint128 kroSharesToBurn; - for (uint256 i = 0; i < tokenIds.length; ) { - uint128 kroShares = kghDelegator.shares[tokenIds[i]].kro - - _convertToKroShares(validator, KGH_MANAGER.totalKroInKgh(tokenIds[i])); - if (kroShares == 0) revert ShareNotExists(); - - unchecked { - kroSharesToBurn += kroShares; - - ++i; - } - } - - uint128 kghBaseReward = _convertToKroAssets(validator, kroSharesToBurn); - - unchecked { - claimedRewards += kghBaseReward; - vault.asset.totalKroShares -= kroSharesToBurn; - vault.asset.totalKro -= claimedRewards; - } - - ASSET_TOKEN.safeTransfer(msg.sender, claimedRewards); - - emit KghRewardClaimed(validator, msg.sender, claimedRewards, kroSharesToBurn); - } - - /** - * @notice Claim the boosted reward of the delegator. - * - * @param validator Address of the validator. - * @param delegator Address of the delegator. - * @param kghNum The number of KGH to claim the boosted reward. - * @param isCompounding Flag to auto compound the boosted reward. - * - * @return The amount of the claimed boosted reward. - */ - function _claimBoostedReward( - address validator, - address delegator, - uint128 kghNum, - bool isCompounding - ) internal returns (uint128) { - Vault storage vault = _vaults[validator]; - KghDelegator storage kghDelegator = vault.kghDelegators[delegator]; - - uint128 rewardPerKghStored = vault.asset.rewardPerKghStored; - uint128 totalBoostedReward = kghDelegator.kghNum * - (rewardPerKghStored - kghDelegator.rewardPerKghPaid[delegator]); - - kghDelegator.rewardPerKghPaid[delegator] = rewardPerKghStored; - - if (isCompounding) { - _delegate(validator, delegator, totalBoostedReward); - } - - vault.asset.boostedReward -= totalBoostedReward; - - return boostedReward; - } - /** * @notice Add pending assets and shares when undelegating KRO. * @@ -1051,8 +1024,6 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager { uint128 shares = _convertToKroShares(validator, assets); Vault storage vault = _vaults[validator]; - ASSET_TOKEN.safeTransferFrom(owner, address(this), assets); - unchecked { vault.asset.totalKro += assets; vault.asset.totalKroShares += shares; @@ -1219,6 +1190,37 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager { } } + /** + * @notice Internal function to claim the boosted reward of the delegator. + * + * @param validator Address of the validator. + * @param delegator Address of the delegator. + * @param kghNum The number of KGH to claim the boosted reward. + * @param isCompounding Flag to auto compound the boosted reward. + * + * @return The amount of the claimed boosted reward. + */ + function _claimBoostedReward( + address validator, + address delegator, + uint128 kghNum, + bool isCompounding + ) internal returns (uint128) { + Vault storage vault = _vaults[validator]; + KghDelegator storage kghDelegator = vault.kghDelegators[delegator]; + + uint128 rewardPerKghStored = vault.asset.rewardPerKghStored; + uint128 totalBoostedReward = kghNum * (rewardPerKghStored - kghDelegator.rewardPerKghPaid); + + kghDelegator.rewardPerKghPaid = rewardPerKghStored; + + if (isCompounding) { + _delegate(validator, delegator, totalBoostedReward); + } + + return totalBoostedReward; + } + /** * @notice Internal function to undelegate KGHs from the validator. * diff --git a/packages/contracts/contracts/L1/interfaces/IAssetManager.sol b/packages/contracts/contracts/L1/interfaces/IAssetManager.sol index 9c161f903..347e0cd00 100644 --- a/packages/contracts/contracts/L1/interfaces/IAssetManager.sol +++ b/packages/contracts/contracts/L1/interfaces/IAssetManager.sol @@ -508,7 +508,7 @@ interface IAssetManager { function undelegateKghBatch(address validator, uint256[] calldata tokenIds) external; /** - * @notice Claim the reward of the KGH delegator from the given validator vault. + * @notice Claim the reward of the KGH delegator from the given validator vault and token ids. * * @param validator Address of the validator. * @param tokenIds Array of token ids of KGHs to claim base reward.