feat: apply PR reviews

packages/contracts/contracts/L1/AssetManager.sol

@@ -199,6 +199,26 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
  return kghAssets + kroAssets;
  }
 
+ /**
+ * @inheritdoc IAssetManager
+ */
+ function getKghReward(
+ address validator,
+ address delegator,
+ uint256[] calldata tokenIds
+ ) external view 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);
+
+ (, uint128 kghBaseReward) = _calculateBaseRewardForKgh(validator, kghDelegator, tokenIds);
+
+ return totalBoostedReward + kghBaseReward;
+ }
+
  /**
  * @inheritdoc IAssetManager
  */
@@ -343,7 +363,9 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
  uint128 assets
  ) external isRegistered(validator) returns (uint128) {
  if (assets == 0) revert NotAllowedZeroInput();
+ ASSET_TOKEN.safeTransferFrom(msg.sender, address(this), assets);
  uint128 shares = _delegate(validator, msg.sender, assets);
+
  emit KroDelegated(validator, msg.sender, assets, shares);
  return shares;
  }
@@ -669,6 +691,37 @@ 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();
+
+ uint128 claimedRewards = _claimBoostedReward(validator, msg.sender);
+ if (claimedRewards == 0) revert InsufficientAsset();
+
+ (uint128 kroSharesToBurn, uint128 kghBaseReward) = _calculateBaseRewardForKgh(
+ validator,
+ kghDelegator,
+ tokenIds
+ );
+
+ unchecked {
+ claimedRewards += kghBaseReward;
+ vault.asset.totalKroShares -= kroSharesToBurn;
+ vault.asset.totalKro -= kghBaseReward;
+ }
+
+ 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 +882,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.
  *
@@ -948,6 +926,38 @@ contract AssetManager is ISemver, IERC721Receiver, IAssetManager {
  vault.kghDelegators[msg.sender].undelegateRequestTimes.push(block.timestamp);
  }
 
+ /**
+ * @notice Internal function to calculate base rewards from the given validator address,
+ * delegator struct, and token ids of KGHs.
+ *
+ * @param validator Address of the validator.
+ * @param delegator KghDelegator struct of the delegator.
+ * @param tokenIds Array of token ids of the KGH to calculate base rewards.
+ *
+ * @return The amount of shares corresponding to the reward.
+ * @return The amount of base rewards.
+ */
+ function _calculateBaseRewardForKgh(
+ address validator,
+ KghDelegator storage delegator,
+ uint256[] calldata tokenIds
+ ) internal view returns (uint128, uint128) {
+ uint128 kroSharesToBurn;
+ for (uint256 i = 0; i < tokenIds.length; ) {
+ if (kghDelegator.delegatedAt[tokenIds[i]] == 0) revert InvalidTokenIdsInput();
+ uint128 kroShares = kghDelegator.kroShares[tokenIds[i]] -
+ _convertToKroShares(validator, KGH_MANAGER.totalKroInKgh(tokenIds[i]));
+
+ unchecked {
+ kroSharesToBurn += kroShares;
+ ++i;
+ }
+ }
+
+ uint128 kghBaseReward = _convertToKroAssets(validator, kroSharesToBurn);
+ return (kroSharesToBurn, kghBaseReward);
+ }
+
  /**
  * @notice Internal conversion function for KRO (from assets to shares).
  *
@@ -1051,8 +1061,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;
@@ -1272,6 +1280,30 @@ 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.
+ *
+ * @return The amount of the claimed boosted reward.
+ */
+ function _claimBoostedReward(
+ address validator,
+ address delegator,
+ ) 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);
+
+ kghDelegator.rewardPerKghPaid = rewardPerKghStored;
+
+ return totalBoostedReward;
+ }
+
  /**
  * @inheritdoc IERC721Receiver
  */

packages/contracts/contracts/L1/interfaces/IAssetManager.sol

@@ -428,10 +428,15 @@ interface IAssetManager {
  *
  * @param validator The address of the validator.
  * @param delegator The address of the delegator.
+ * @param tokenIds The token id array of KGH to check the base reward.
  *
  * @return The amount of claimable reward of KGH delegation.
  */
- function getKghReward(address validator, address delegator) external view returns (uint128);
+ function getKghReward(
+ address validator,
+ address delegator,
+ uint256[] calldata tokenIds
+ ) external view returns (uint128);
 
  /**
  * @notice Deposit KRO. To deposit KRO, the validator should be initiated.
@@ -508,7 +513,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.