Add Decay to stake amount

Author: mikeghen

src/QueryTypeStakerFactory.sol

@@ -14,9 +14,15 @@ contract QueryTypeStakerFactory is Ownable {
   /// @notice The token that will be used for staking.
   IERC20 public immutable STAKING_TOKEN;
 
+  /// @notice Address that receives fees and decayed stake from pools
+  address public feeRecipient;
+
   /// @notice Maps query type bit fields to their corresponding staking pool addresses.
   mapping(bytes32 queryType => address poolAddress) public queryTypePools;
 
+  /// @notice Emitted when fee recipient is updated
+  event FeeRecipientUpdated(address indexed oldRecipient, address indexed newRecipient);
+
   /// @notice Emitted when a new staking pool is created for a query type bit field.
   event CreateQueryTypeStakingPool(bytes32 indexed queryType, address indexed poolAddress);
 
@@ -32,6 +38,7 @@ contract QueryTypeStakerFactory is Ownable {
   constructor(address _owner, address _stakingToken) Ownable(_owner) {
     if (_stakingToken == address(0)) revert QueryTypeStakerFactory__InvalidTokenAddress();
     STAKING_TOKEN = IERC20(_stakingToken);
+    _setFeeRecipient(_owner);
   }
 
   /// @notice Creates a new staking pool for a specific query type bit field.
@@ -49,12 +56,27 @@ contract QueryTypeStakerFactory is Ownable {
 
     // Deploy new staking pool with STAKING_TOKEN address and initial conversion entry
     QueryTypeStakingPool _newPool =
-      new QueryTypeStakingPool(_poolOwner, address(STAKING_TOKEN), _initialEntry);
+      new QueryTypeStakingPool(_poolOwner, address(STAKING_TOKEN), address(this), _initialEntry);
     _poolAddress = address(_newPool);
 
     // Store the pool address
     queryTypePools[_queryType] = _poolAddress;
 
     emit CreateQueryTypeStakingPool(_queryType, _poolAddress);
   }
+
+  /// @notice Updates the fee recipient address.
+  /// @param _newRecipient The new address to receive fees.
+  function setFeeRecipient(address _newRecipient) external {
+    _checkOwner();
+    _setFeeRecipient(_newRecipient);
+  }
+
+  /// @notice Internal function to update the fee recipient and emit an event.
+  /// @param _newRecipient The new address to receive fees.
+  function _setFeeRecipient(address _newRecipient) internal {
+    address old = feeRecipient;
+    feeRecipient = _newRecipient;
+    emit FeeRecipientUpdated(old, _newRecipient);
+  }
 }

src/QueryTypeStakingPool.sol

@@ -4,6 +4,7 @@ pragma solidity 0.8.26;
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {QueryTypeStakerFactory} from "src/QueryTypeStakerFactory.sol";
 
 /// @title QueryTypeStakingPool
 /// @author ScopeLift
@@ -29,6 +30,9 @@ contract QueryTypeStakingPool is Ownable {
   /// @notice The ERC20 token contract that can be staked in this pool.
   IERC20 public immutable STAKING_TOKEN;
 
+  /// @notice Address of the factory that deployed this pool. Provides the feeRecipient.
+  address public immutable FACTORY;
+
   /// @notice A struct containing information about a user's stake, including the amount staked, the
   /// index into the conversion table history at time of staking, and the lockup/access period end
   /// times.
@@ -37,6 +41,8 @@ contract QueryTypeStakingPool is Ownable {
     uint256 conversionTableIndex;
     uint48 lockupEnd;
     uint48 accessEnd;
+    uint48 lastClaimed;
+    uint256 decayed;
   }
 
   /// @notice A mapping that associates staker addresses with their stake information.
@@ -97,6 +103,9 @@ contract QueryTypeStakingPool is Ownable {
   /// @notice Emitted when an address is blocklisted for this pool
   event AddressBlocklisted(address indexed user);
 
+  /// @notice Emitted when decayed stake is claimed and forwarded to the fee recipient.
+  event DecayClaimed(address indexed staker, uint256 amount, address indexed feeRecipient);
+
   /// @notice Thrown when attempting to stake with an invalid lockup period.
   error QueryTypeStakingPool__LockupPeriodTooLow();
 
@@ -138,11 +147,16 @@ contract QueryTypeStakingPool is Ownable {
   /// @param _owner The address that will own the contract and have permission to update the
   /// conversion table.
   /// @param _stakingToken The address of the ERC20 token that will be staked.
+  /// @param _factory The address of the factory that deployed this pool.
   /// @param _initialConversionTableEntry The first entry in the conversion table history.
-  constructor(address _owner, address _stakingToken, bytes32 _initialConversionTableEntry)
-    Ownable(_owner)
-  {
+  constructor(
+    address _owner,
+    address _stakingToken,
+    address _factory,
+    bytes32 _initialConversionTableEntry
+  ) Ownable(_owner) {
     STAKING_TOKEN = IERC20(_stakingToken);
+    FACTORY = _factory;
 
     // Initialize the conversion table with the provided entry
     conversionTableHistory.push(_initialConversionTableEntry);
@@ -192,6 +206,8 @@ contract QueryTypeStakingPool is Ownable {
   /// @notice Allows users to stake tokens for the predefined lockup and access periods.
   /// @param _amount The amount of tokens to stake.
   function stake(uint256 _amount) external {
+    _claimDecay(msg.sender);
+
     if (_amount < minimumStake) revert QueryTypeStakingPool__AmountBelowMinimum();
     if (isBlocklisted[msg.sender]) revert QueryTypeStakingPool__AddressBlocklisted();
 
@@ -208,7 +224,9 @@ contract QueryTypeStakingPool is Ownable {
     if (stakeInfo.amount == 0) {
       // First-time stake
       stakeInfo.amount = _amount;
+      stakeInfo.decayed = 0;
       stakeInfo.conversionTableIndex = conversionTableHistory.length - 1;
+      stakeInfo.lastClaimed = uint48(block.timestamp);
       stakes[msg.sender] = stakeInfo;
       STAKING_TOKEN.safeTransferFrom(msg.sender, address(this), _amount);
       emit Staked(
@@ -222,6 +240,7 @@ contract QueryTypeStakingPool is Ownable {
     }
 
     stakeInfo.amount += _amount;
+    stakeInfo.lastClaimed = uint48(block.timestamp);
     stakes[msg.sender] = stakeInfo;
 
     STAKING_TOKEN.safeTransferFrom(msg.sender, address(this), _amount);
@@ -241,13 +260,17 @@ contract QueryTypeStakingPool is Ownable {
   /// @param _amount The amount of tokens the user wishes to unstake.
   function unstake(uint256 _amount) external {
     StakeInfo storage userStake = stakes[msg.sender];
+
+    _claimDecay(msg.sender);
+
     if (userStake.amount == 0) revert QueryTypeStakingPool__NoStakeFound();
     if (block.timestamp < userStake.lockupEnd) revert QueryTypeStakingPool__StillInLockupPeriod();
     if (_amount > userStake.amount) revert QueryTypeStakingPool__InsufficientBalance();
 
     if (isBlocklisted[msg.sender]) totalJailed -= _amount;
     else totalStaked -= _amount;
 
+    userStake.decayed = 0;
     userStake.amount -= _amount;
     STAKING_TOKEN.safeTransfer(msg.sender, _amount);
 
@@ -283,4 +306,44 @@ contract QueryTypeStakingPool is Ownable {
     isBlocklisted[_user] = true;
     emit AddressBlocklisted(_user);
   }
+
+  /// @notice Claims the decayed portion of the caller's stake and forwards it to the feeRecipient.
+  /// Any address can call this function on behalf of a staker.
+  /// @param _staker The address whose decayed stake should be claimed. If omitted, defaults to
+  /// msg.sender.
+  function claim(address _staker) public {
+    _claimDecay(_staker);
+  }
+
+  /// @notice Internal helper that settles the decayed portion of a stake.
+  /// @param _staker The address whose decayed stake should be claimed.
+  /// @return _claimed The amount of decayed stake claimed.
+  function _claimDecay(address _staker) internal returns (uint256 _claimed) {
+    StakeInfo storage stakeInfo = stakes[_staker];
+    if (stakeInfo.amount == 0) return 0;
+
+    uint256 elapsed = block.timestamp - stakeInfo.lastClaimed;
+    if (elapsed == 0) return 0;
+
+    uint256 totalPeriod = stakeInfo.accessEnd - stakeInfo.lastClaimed;
+    if (totalPeriod == 0) return 0;
+
+    uint256 decayed = (stakeInfo.amount * elapsed) / totalPeriod;
+    if (decayed == 0) return 0;
+
+    address _feeRecipient = QueryTypeStakerFactory(FACTORY).feeRecipient();
+    if (_feeRecipient == address(0)) return 0;
+
+    if (decayed > stakeInfo.amount) decayed = stakeInfo.amount;
+
+    // Apply decay and update accounting
+    stakeInfo.amount -= decayed;
+    stakeInfo.decayed += decayed;
+    stakeInfo.lastClaimed = uint48(block.timestamp);
+
+    STAKING_TOKEN.safeTransfer(_feeRecipient, decayed);
+
+    emit DecayClaimed(_staker, decayed, _feeRecipient);
+    return decayed;
+  }
 }

test/QueryTypeStakerFactory.t.sol

@@ -33,6 +33,7 @@ contract Constructor is QueryTypeStakerFactoryTest {
     QueryTypeStakerFactory newFactory = new QueryTypeStakerFactory(_owner, _stakingToken);
     assertEq(newFactory.owner(), _owner);
     assertEq(address(newFactory.STAKING_TOKEN()), _stakingToken);
+    assertEq(newFactory.feeRecipient(), _owner);
   }
 
   function testFuzz_RevertIf_StakingTokenAddressIsZero(address _owner) public {
@@ -41,6 +42,16 @@ contract Constructor is QueryTypeStakerFactoryTest {
     vm.expectRevert(QueryTypeStakerFactory.QueryTypeStakerFactory__InvalidTokenAddress.selector);
     new QueryTypeStakerFactory(_owner, address(0));
   }
+
+  function testFuzz_EmitsFeeRecipientUpdatedEventWithArbitraryOwner(
+    address _owner,
+    address _stakingToken
+  ) public {
+    vm.assume(_owner != address(0) && _stakingToken != address(0));
+    vm.expectEmit();
+    emit QueryTypeStakerFactory.FeeRecipientUpdated(address(0), _owner);
+    new QueryTypeStakerFactory(_owner, _stakingToken);
+  }
 }
 
 contract CreateStakingPool is QueryTypeStakerFactoryTest {
@@ -101,3 +112,31 @@ contract CreateStakingPool is QueryTypeStakerFactoryTest {
     vm.stopPrank();
   }
 }
+
+contract SetFeeRecipient is QueryTypeStakerFactoryTest {
+  function testFuzz_SetsFeeRecipientCorrectly(address _newRecipient) public {
+    vm.assume(_newRecipient != address(0));
+    vm.prank(owner);
+    factory.setFeeRecipient(_newRecipient);
+    assertEq(factory.feeRecipient(), _newRecipient);
+  }
+
+  function testFuzz_RevertIf_CallerIsNotOwner(address _notOwner, address _newRecipient) public {
+    vm.assume(_notOwner != owner && _notOwner != address(0));
+    vm.assume(_newRecipient != address(0));
+
+    vm.prank(_notOwner);
+    vm.expectRevert(abi.encodeWithSignature("OwnableUnauthorizedAccount(address)", _notOwner));
+    factory.setFeeRecipient(_newRecipient);
+  }
+
+  function testFuzz_EmitsFeeRecipientUpdatedEventWithArbitraryRecipient(address _newRecipient)
+    public
+  {
+    vm.assume(_newRecipient != address(0));
+    vm.expectEmit();
+    emit QueryTypeStakerFactory.FeeRecipientUpdated(owner, _newRecipient);
+    vm.prank(owner);
+    factory.setFeeRecipient(_newRecipient);
+  }
+}