feat: apply PR reviews

kroma-network · Jul 19, 2024 · 7520f30 · 7520f30
1 parent f9df342
commit 7520f30
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 78 deletions.
diff --git 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
@@ -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.