feat: implement claimKghReward function

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.
@@ -873,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).
  *
@@ -976,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;
@@ -1197,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

@@ -54,7 +54,7 @@ interface IAssetManager {
  */
  struct KghDelegator {
  uint128 rewardPerKghPaid;
- uint256 kghNum;
+ uint128 kghNum;
  mapping(uint256 => uint128) delegatedAt;
  mapping(uint256 => uint128) kroShares;
  }
@@ -260,6 +260,11 @@ interface IAssetManager {
  */
  error ShareNotExists();
 
+ /**
+ * @notice Reverts when the given token ids are invalid.
+ */
+ error InvalidTokenIdsInput();
+
  /**
  * @notice Returns the address of withdraw account of given validator.
  *
@@ -423,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.
@@ -503,9 +513,10 @@ 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.
  */
- function claimKghReward(address validator) external;
+ function claimKghReward(address validator, uint256[] calldata tokenIds) external;
 }