Add callbacks to Generator Registry

contracts/GeneratorRegistry.sol

@@ -17,21 +17,26 @@ import "./ProofMarketplace.sol";
 import {IStakingManager} from "./interfaces/staking/IStakingManager.sol";
 import "./interfaces/IGeneratorRegistry.sol";
 
+import "./interfaces/IGeneratorCallbacks.sol";
+import "./staking/l2_contracts/StakingManager.sol";
+
 contract GeneratorRegistry is
     Initializable,
     ContextUpgradeable,
     ERC165Upgradeable,
     AccessControlUpgradeable,
     UUPSUpgradeable,
     ReentrancyGuardUpgradeable,
-    IGeneratorRegistry
+    IGeneratorRegistry,
+    IGeneratorCallbacks
 {
     // in case we add more contracts in the inheritance chain
     uint256[500] private __gap_0;
 
     /// @custom:oz-upgrades-unsafe-allow constructor
-    constructor(EntityKeyRegistry _entityRegistry) initializer {
+    constructor(EntityKeyRegistry _entityRegistry, StakingManager _stakingManager) initializer {
         ENTITY_KEY_REGISTRY = _entityRegistry;
+        STAKING_MANAGER = _stakingManager;
     }
 
     using HELPER for bytes;
@@ -64,6 +69,8 @@ contract GeneratorRegistry is
 
     /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     EntityKeyRegistry public immutable ENTITY_KEY_REGISTRY;
+
+    StakingManager public immutable STAKING_MANAGER;
     //-------------------------------- Constants and Immutable start --------------------------------//
 
     //-------------------------------- State variables start --------------------------------//
@@ -514,7 +521,7 @@ contract GeneratorRegistry is
     }
 
     /**
-     * @notice Should be called by proof market place only, PMP is assigned SLASHER_ROLE
+     * @notice Should be called by proof market place only, PMP is assigned SLASHER_ROLE, called when generators is about to be slashed
      */
     function releaseGeneratorResources(
         address generatorAddress,
@@ -534,6 +541,7 @@ contract GeneratorRegistry is
         info.activeRequests--;
 
         generator.computeConsumed -= info.computePerRequestRequired;
+        emit ComputeLockReleased(generatorAddress, info.computePerRequestRequired);
     }
 
     function assignGeneratorTask(
@@ -607,6 +615,62 @@ contract GeneratorRegistry is
         return (generator.rewardAddress, info.proofGenerationCost);
     }
 
+    function addStakeCallback(address generatorAddress, address token, uint256 amount) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit AddedStake(generatorAddress, token, amount);
+    }
+
+    function intendToReduceStakeCallback(address generatorAddress, address token, uint256 amount) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit RequestStakeDecrease(generatorAddress, token, amount);
+    }
+
+    function removeStakeCallback(address generatorAddress, address token, uint256 amount) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit RemovedStake(generatorAddress, token, amount);
+    }
+
+    function stakeLockImposedCallback(address generatorAddress, address token, uint256 amount) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit StakeLockImposed(generatorAddress, token, amount);
+    }
+
+    function stakeLockReleasedCallback(address generatorAddress, address token, uint256 amount) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit StakeLockReleased(generatorAddress, token, amount);
+    }
+
+    function stakeSlashedCallback(address generatorAddress, address token, uint256 amount) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit StakeSlashed(generatorAddress, token, amount);
+    }
+
+    function symbioticCompleteSnapshotCallback(uint256 captureTimestamp) external override {
+        if(!STAKING_MANAGER.isEnabledPool(msg.sender)){
+            revert Error.InvalidContractAddress();
+        }
+
+        emit SymbioticCompleteSnapshot(captureTimestamp);
+    }
+
     // for further increase
     uint256[50] private __gap1_0;
 }

contracts/binding_generator.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "./EntityKeyRegistry.sol";
+import "./GeneratorRegistry.sol";
+import "./ProofMarketplace.sol";
+import "./staking/l2_contracts/NativeStaking.sol";
+import "./staking/l2_contracts/SymbioticStaking.sol";
+
+contract binding_generator {
+
+}

contracts/interfaces/IGeneratorCallbacks.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+interface IGeneratorCallbacks {
+    event AddedStake(address indexed generatorAddress, address indexed token, uint256 amount);
+    function addStakeCallback(address generatorAddress, address token, uint256 amount) external;
+
+    event RequestStakeDecrease(address indexed generatorAddress, address indexed token, uint256 amount);
+    function intendToReduceStakeCallback(address generatorAddress, address token, uint256 amount) external;
+
+    event RemovedStake(address indexed generatorAddress, address indexed token, uint256 amount);
+    function removeStakeCallback(address generatorAddress, address token, uint256 amount) external;
+
+    event StakeLockImposed(address indexed generatorAddress, address indexed token, uint256 stake);
+    function stakeLockImposedCallback(address generatorAddress, address token, uint256 amount) external;
+
+    event StakeLockReleased(address indexed generatorAddress, address indexed token, uint256 stake);
+    function stakeLockReleasedCallback(address generatorAddress, address token, uint256 amount) external;
+
+    event StakeSlashed(address indexed generatorAddress, address indexed token, uint256 stake);
+    function stakeSlashedCallback(address generatorAddress, address token, uint256 amount) external;
+
+    event SymbioticCompleteSnapshot(uint256 indexed captureTimestamp);
+    function symbioticCompleteSnapshotCallback(uint256 captureTimestamp) external;
+}

contracts/interfaces/IGeneratorRegistry.sol

@@ -20,9 +20,9 @@ contract IGeneratorRegistry {
     event RequestComputeDecrease(address indexed generator, uint256 intendedUtilization);
     event DecreaseCompute(address indexed generator, uint256 compute);
 
-    event ComputeLockImposed(address indexed generator, uint256 stake);
+    event ComputeLockImposed(address indexed generator, uint256 compute);
 
-    event ComputeLockReleased(address indexed generator, uint256 stake);
+    event ComputeLockReleased(address indexed generator, uint256 compute);
 
     //-------------------------------- Events end --------------------------------//
 }

contracts/staking/l2_contracts/NativeStaking.sol

@@ -21,6 +21,8 @@ import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 
+import "../../interfaces/IGeneratorCallbacks.sol";
+
 contract NativeStaking is
     ContextUpgradeable,
     ERC165Upgradeable,
@@ -32,6 +34,11 @@ contract NativeStaking is
     using EnumerableSet for EnumerableSet.AddressSet;
     using SafeERC20 for IERC20;
 
+    IGeneratorCallbacks public immutable I_GENERATOR_CALLBACK;
+    constructor(IGeneratorCallbacks _generator_callback) {
+        I_GENERATOR_CALLBACK = _generator_callback;
+    }
+
     /*===================================================================================================================*/
     /*================================================ state variable ===================================================*/
     /*===================================================================================================================*/
@@ -138,6 +145,8 @@ contract NativeStaking is
         operatorstakeAmounts[_stakeToken][_operator] += _amount;
 
         emit Staked(msg.sender, _operator, _stakeToken, _amount);
+
+        I_GENERATOR_CALLBACK.addStakeCallback(_operator, _stakeToken, _amount);
     }
 
     // TODO
@@ -154,6 +163,8 @@ contract NativeStaking is
         uint256 index = withdrawalRequests[msg.sender][_operator].length - 1;
 
         emit StakeWithdrawalRequested(msg.sender, _operator, _stakeToken, index, _amount);
+
+        I_GENERATOR_CALLBACK.intendToReduceStakeCallback(_operator, _stakeToken, _amount);
     }
 
     function withdrawStake(address _operator, uint256[] calldata _index) external nonReentrant {
@@ -172,6 +183,8 @@ contract NativeStaking is
             IERC20(request.stakeToken).safeTransfer(msg.sender, request.amount);
 
             emit StakeWithdrawn(msg.sender, _operator, request.stakeToken, _index[i], request.amount);
+
+            I_GENERATOR_CALLBACK.removeStakeCallback(_operator, request.stakeToken, request.amount);
         }
     }
 
@@ -187,6 +200,8 @@ contract NativeStaking is
         operatorLockedAmounts[_stakeToken][_operator] += _amountToLock;
 
         emit StakeLocked(_jobId, _operator, _stakeToken, _amountToLock);
+
+        I_GENERATOR_CALLBACK.stakeLockImposedCallback(_operator, _stakeToken, _amountToLock);
     }
 
     /// @notice unlock stake and distribute reward
@@ -203,6 +218,8 @@ contract NativeStaking is
         _unlockStake(_jobId, lock.token, _operator, lock.amount);
 
         emit StakeUnlocked(_jobId, _operator, lock.token, lock.amount);
+
+        I_GENERATOR_CALLBACK.stakeLockReleasedCallback(_operator, lock.token, lock.amount);
     }
 
     function slash(Struct.JobSlashed[] calldata _slashedJobs) external onlyStakingManager {
@@ -217,6 +234,8 @@ contract NativeStaking is
             IERC20(lock.token).safeTransfer(_slashedJobs[i].rewardAddress, lockedAmount);
 
             emit JobSlashed(_slashedJobs[i].jobId, _slashedJobs[i].operator, lock.token, lockedAmount);
+
+            I_GENERATOR_CALLBACK.stakeSlashedCallback(_slashedJobs[i].operator, lock.token, lockedAmount);
         }
     }
 

contracts/staking/l2_contracts/SymbioticStaking.sol

@@ -27,6 +27,8 @@ import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol
 
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 
+import "../../interfaces/IGeneratorCallbacks.sol";
+
 import {console} from "hardhat/console.sol";
 
 contract SymbioticStaking is
@@ -40,6 +42,11 @@ contract SymbioticStaking is
     using EnumerableSet for EnumerableSet.AddressSet;
     using SafeERC20 for IERC20;
 
+    IGeneratorCallbacks public immutable I_GENERATOR_CALLBACK;
+    constructor(IGeneratorCallbacks _generator_callback) {
+        I_GENERATOR_CALLBACK = _generator_callback;
+    }
+
     struct EnclaveImage {
         bytes PCR0;
         bytes PCR1;
@@ -252,6 +259,8 @@ contract SymbioticStaking is
         operatorLockedAmounts[_stakeToken][_operator] += _amountToLock;
 
         emit StakeLocked(_jobId, _operator, _stakeToken, _amountToLock);
+
+        I_GENERATOR_CALLBACK.stakeLockImposedCallback(_operator, _stakeToken, _amountToLock);
     }
 
     function onJobCompletion(uint256 _jobId, address _operator, uint256 _feeRewardAmount) external onlyStakingManager {
@@ -278,6 +287,8 @@ contract SymbioticStaking is
         operatorLockedAmounts[lock.stakeToken][_operator] -= amountToLock[lock.stakeToken];
 
         emit StakeUnlocked(_jobId, _operator, lock.stakeToken, amountToLock[lock.stakeToken]);
+
+        I_GENERATOR_CALLBACK.stakeLockReleasedCallback(_operator, lock.stakeToken, amountToLock[lock.stakeToken]);
     }
 
     /*------------------------------------- slash ------------------------------------*/
@@ -295,6 +306,8 @@ contract SymbioticStaking is
             delete lockInfo[_slashedJobs[i].jobId];
 
             emit JobSlashed(_slashedJobs[i].jobId, _slashedJobs[i].operator, lock.stakeToken, lockedAmount);
+
+            I_GENERATOR_CALLBACK.stakeSlashedCallback(_slashedJobs[i].operator, lock.stakeToken, lockedAmount);
         }
     }
 
@@ -351,6 +364,8 @@ contract SymbioticStaking is
         confirmedTimestamps.push(confirmedTimestamp);
 
         emit SnapshotConfirmed(msg.sender, _captureTimestamp);
+
+        I_GENERATOR_CALLBACK.symbioticCompleteSnapshotCallback(_captureTimestamp);
     }
 
     /*===================================================================================================================*/
@@ -379,6 +394,11 @@ contract SymbioticStaking is
         return operatorStakeAmounts[latestConfirmedTimestamp()][_stakeToken][_operator];
     }
 
+    function getOperatorStakeAmountAt(uint256 _captureTimestamp, address _stakeToken, address _operator) public view returns (uint256) {
+        return operatorStakeAmounts[_captureTimestamp][_stakeToken][_operator];
+    }
+
+
     function getOperatorActiveStakeAmount(address _stakeToken, address _operator) public view returns (uint256) {
         uint256 operatorStakeAmount = getOperatorStakeAmount(_stakeToken, _operator);
         uint256 operatorLockedAmount = operatorLockedAmounts[_stakeToken][_operator];