feat: operatorStateWithSocket (#418)

0x0aa0 · pakim249CAL · pakim249CAL · commit 7b4031cdc5ad · 2025-04-07T15:54:04.000Z
**Motivation:**

*Explain here the context, and why you're making that change. What is
the problem you're trying to solve.*

**Modifications:**

*Describe the modifications you've done.*

**Result:**

*After your change, what will change.*

---------

Co-authored-by: Patrick Kim &lt;pakim249@gmail.com&gt;
diff --git a/src/OperatorStateRetriever.sol b/src/OperatorStateRetriever.sol
@@ -2,6 +2,7 @@
 pragma solidity ^0.8.27;
 
 import {ISlashingRegistryCoordinator} from "./interfaces/ISlashingRegistryCoordinator.sol";
+import {ISocketRegistry} from "./interfaces/ISocketRegistry.sol";
 import {IBLSApkRegistry} from "./interfaces/IBLSApkRegistry.sol";
 import {IStakeRegistry} from "./interfaces/IStakeRegistry.sol";
 import {IIndexRegistry} from "./interfaces/IIndexRegistry.sol";
@@ -94,6 +95,81 @@ contract OperatorStateRetriever {
         return operators;
     }
 
+    /**
+     * @notice This function is intended to to be called by AVS operators every time a new task is created (i.e.)
+     * the AVS coordinator makes a request to AVS operators. Since all of the crucial information is kept onchain, 
+     * operators don't need to run indexers to fetch the data.
+     * @param registryCoordinator is the registry coordinator to fetch the AVS registry information from
+     * @param operatorId the id of the operator to fetch the quorums lists 
+     * @param blockNumber is the block number to get the operator state for
+     * @return quorumBitmap the quorumBitmap of the operator at the given blockNumber
+     * @return operators a 2d array of Operators. For each quorum, an ordered list of Operators
+     * @return sockets a 2d array of sockets. For each quorum, an ordered list of sockets
+     */
+    function getOperatorStateWithSocket(
+        ISlashingRegistryCoordinator registryCoordinator, 
+        bytes32 operatorId, 
+        uint32 blockNumber
+    ) external view returns (uint256 quorumBitmap, Operator[][] memory operators, string[][] memory sockets) {
+        bytes32[] memory operatorIds = new bytes32[](1);
+        operatorIds[0] = operatorId;
+        uint256 index = registryCoordinator.getQuorumBitmapIndicesAtBlockNumber(blockNumber, operatorIds)[0];
+    
+        quorumBitmap = registryCoordinator.getQuorumBitmapAtBlockNumberByIndex(operatorId, blockNumber, index);
+
+        bytes memory quorumNumbers = BitmapUtils.bitmapToBytesArray(quorumBitmap);
+
+        (operators, sockets) =
+            getOperatorStateWithSocket(registryCoordinator, quorumNumbers, blockNumber);
+    }
+
+    /// @dev Used below to avoid stack too deep.
+    struct Registries {
+        IStakeRegistry stakeRegistry;
+        IIndexRegistry indexRegistry;
+        IBLSApkRegistry blsApkRegistry;
+        ISocketRegistry socketRegistry;
+    }
+
+    /**
+     * @notice returns the ordered list of operators (id, stake, socket) for each quorum. The AVS coordinator 
+     * may call this function directly to get the operator state for a given block number
+     * @param registryCoordinator is the registry coordinator to fetch the AVS registry information from
+     * @param quorumNumbers are the ids of the quorums to get the operator state for
+     * @param blockNumber is the block number to get the operator state for
+     * @return operators a 2d array of Operators. For each quorum, an ordered list of Operators
+     * @return sockets a 2d array of sockets. For each quorum, an ordered list of sockets
+     */
+     function getOperatorStateWithSocket(
+        ISlashingRegistryCoordinator registryCoordinator,
+        bytes memory quorumNumbers,
+        uint32 blockNumber
+    ) public view returns (Operator[][] memory operators, string[][] memory sockets) {
+        Registries memory registries = Registries({
+            stakeRegistry: registryCoordinator.stakeRegistry(),
+            indexRegistry: registryCoordinator.indexRegistry(),
+            blsApkRegistry: registryCoordinator.blsApkRegistry(),
+            socketRegistry: registryCoordinator.socketRegistry()
+        });
+
+        operators = new Operator[][](quorumNumbers.length);
+        sockets = new string[][](quorumNumbers.length);
+        for (uint256 i = 0; i < quorumNumbers.length; i++) {
+            uint8 quorumNumber = uint8(quorumNumbers[i]);
+            bytes32[] memory operatorIds = registries.indexRegistry.getOperatorListAtBlockNumber(quorumNumber, blockNumber);
+            operators[i] = new Operator[](operatorIds.length);
+            sockets[i] = new string[](operatorIds.length);
+            for (uint256 j = 0; j < operatorIds.length; j++) {
+                operators[i][j] = Operator({
+                    operator: registries.blsApkRegistry.getOperatorFromPubkeyHash(operatorIds[j]),
+                    operatorId: bytes32(operatorIds[j]),
+                    stake: registries.stakeRegistry.getStakeAtBlockNumber(bytes32(operatorIds[j]), quorumNumber, blockNumber)
+                });
+                sockets[i][j] = registries.socketRegistry.getOperatorSocket(bytes32(operatorIds[j]));
+            }
+        }
+    }
+
     /**
      * @notice this is called by the AVS operator to get the relevant indices for the checkSignatures function
      * if they are not running an indexer
diff --git a/test/unit/OperatorStateRetrieverUnit.t.sol b/test/unit/OperatorStateRetrieverUnit.t.sol
@@ -2,7 +2,7 @@
 pragma solidity ^0.8.27;
 
 import "../utils/MockAVSDeployer.sol";
-import {IStakeRegistryErrors} from "../../src/interfaces/IStakeRegistry.sol";
+import {IStakeRegistryErrors, IStakeRegistryTypes} from "../../src/interfaces/IStakeRegistry.sol";
 import {ISlashingRegistryCoordinatorTypes} from "../../src/interfaces/IRegistryCoordinator.sol";
 
 contract OperatorStateRetrieverUnitTests is MockAVSDeployer {
@@ -144,6 +144,136 @@ contract OperatorStateRetrieverUnitTests is MockAVSDeployer {
         assertEq(operators[1][0].stake, defaultStake - 1);
     }
 
+    function test_getOperatorStateWithSocket_revert_neverRegistered() public {
+        cheats.expectRevert(
+            "RegistryCoordinator.getQuorumBitmapIndexAtBlockNumber: no bitmap update found for operatorId"
+        );
+        operatorStateRetriever.getOperatorStateWithSocket(
+            registryCoordinator, defaultOperatorId, uint32(block.number)
+        );
+    }
+
+    
+
+    function test_getOperatorStateWithSocket_revert_registeredFirstAfterReferenceBlockNumber() public {
+        cheats.roll(registrationBlockNumber);
+        _registerOperatorWithCoordinator(defaultOperator, 1, defaultPubKey);
+
+        // should revert because the operator was registered for the first time after the reference block number
+        cheats.expectRevert(
+            "RegistryCoordinator.getQuorumBitmapIndexAtBlockNumber: no bitmap update found for operatorId"
+        );
+        operatorStateRetriever.getOperatorStateWithSocket(
+            registryCoordinator, defaultOperatorId, registrationBlockNumber - 1
+        );
+    }
+
+    function test_getOperatorStateWithSocket_deregisteredBeforeReferenceBlockNumber() public {
+        uint256 quorumBitmap = 1;
+        cheats.roll(registrationBlockNumber);
+        _registerOperatorWithCoordinator(defaultOperator, quorumBitmap, defaultPubKey);
+
+        cheats.roll(registrationBlockNumber + 10);
+        cheats.prank(defaultOperator);
+        registryCoordinator.deregisterOperator(BitmapUtils.bitmapToBytesArray(quorumBitmap));
+
+        (uint256 fetchedQuorumBitmap, OperatorStateRetriever.Operator[][] memory operators,) =
+        operatorStateRetriever.getOperatorStateWithSocket(
+            registryCoordinator, defaultOperatorId, uint32(block.number)
+        );
+        assertEq(fetchedQuorumBitmap, 0);
+        assertEq(operators.length, 0);
+    }
+
+    function test_getOperatorStateWithSocket_registeredAtReferenceBlockNumber() public {
+        uint256 quorumBitmap = 1;
+        cheats.roll(registrationBlockNumber);
+        _registerOperatorWithCoordinator(defaultOperator, quorumBitmap, defaultPubKey);
+
+        (uint256 fetchedQuorumBitmap, OperatorStateRetriever.Operator[][] memory operators,) =
+        operatorStateRetriever.getOperatorStateWithSocket(
+            registryCoordinator, defaultOperatorId, uint32(block.number)
+        );
+        assertEq(fetchedQuorumBitmap, 1);
+        assertEq(operators.length, 1);
+        assertEq(operators[0].length, 1);
+        assertEq(operators[0][0].operator, defaultOperator);
+        assertEq(operators[0][0].operatorId, defaultOperatorId);
+        assertEq(operators[0][0].stake, defaultStake);
+    }
+
+    function test_getOperatorStateWithSocket_revert_quorumNotCreatedAtCallTime() public {
+        cheats.expectRevert(
+            "IndexRegistry._operatorCountAtBlockNumber: quorum did not exist at given block number"
+        );
+        operatorStateRetriever.getOperatorStateWithSocket(
+            registryCoordinator,
+            BitmapUtils.bitmapToBytesArray(1 << numQuorums),
+            uint32(block.number)
+        );
+    }
+
+    function test_getOperatorStateWithSocket_revert_quorumNotCreatedAtReferenceBlockNumber() public {
+        cheats.roll(registrationBlockNumber);
+        ISlashingRegistryCoordinator.OperatorSetParam memory operatorSetParams = ISlashingRegistryCoordinatorTypes
+            .OperatorSetParam({
+            maxOperatorCount: defaultMaxOperatorCount,
+            kickBIPsOfOperatorStake: defaultKickBIPsOfOperatorStake,
+            kickBIPsOfTotalStake: defaultKickBIPsOfTotalStake
+        });
+        uint96 minimumStake = 1;
+        IStakeRegistry.StrategyParams[] memory strategyParams =
+            new IStakeRegistry.StrategyParams[](1);
+        strategyParams[0] =
+            IStakeRegistryTypes.StrategyParams({strategy: IStrategy(address(1000)), multiplier: 1e16});
+
+        cheats.prank(registryCoordinator.owner());
+        registryCoordinator.createTotalDelegatedStakeQuorum(operatorSetParams, minimumStake, strategyParams);
+
+        cheats.expectRevert(
+            "IndexRegistry._operatorCountAtBlockNumber: quorum did not exist at given block number"
+        );
+        operatorStateRetriever.getOperatorStateWithSocket(
+            registryCoordinator,
+            BitmapUtils.bitmapToBytesArray(1 << numQuorums),
+            uint32(registrationBlockNumber - 1)
+        );
+    }
+
+    
+    function test_getOperatorStateWithSocket_returnsCorrect() public {
+        uint256 quorumBitmapOne = 1;
+        uint256 quorumBitmapThree = 3;
+        cheats.roll(registrationBlockNumber);
+        _registerOperatorWithCoordinator(defaultOperator, quorumBitmapOne, defaultPubKey);
+
+        address otherOperator = _incrementAddress(defaultOperator, 1);
+        BN254.G1Point memory otherPubKey = BN254.G1Point(1, 2);
+        bytes32 otherOperatorId = BN254.hashG1Point(otherPubKey);
+        _registerOperatorWithCoordinator(
+            otherOperator, quorumBitmapThree, otherPubKey, defaultStake - 1
+        );
+
+        (OperatorStateRetriever.Operator[][] memory operators,) = operatorStateRetriever
+            .getOperatorStateWithSocket(
+            registryCoordinator,
+            BitmapUtils.bitmapToBytesArray(quorumBitmapThree),
+            uint32(block.number)
+        );
+        assertEq(operators.length, 2);
+        assertEq(operators[0].length, 2);
+        assertEq(operators[1].length, 1);
+        assertEq(operators[0][0].operator, defaultOperator);
+        assertEq(operators[0][0].operatorId, defaultOperatorId);
+        assertEq(operators[0][0].stake, defaultStake);
+        assertEq(operators[0][1].operator, otherOperator);
+        assertEq(operators[0][1].operatorId, otherOperatorId);
+        assertEq(operators[0][1].stake, defaultStake - 1);
+        assertEq(operators[1][0].operator, otherOperator);
+        assertEq(operators[1][0].operatorId, otherOperatorId);
+        assertEq(operators[1][0].stake, defaultStake - 1);
+    }
+
     function test_getCheckSignaturesIndices_revert_neverRegistered() public {
         bytes32[] memory nonSignerOperatorIds = new bytes32[](1);
         nonSignerOperatorIds[0] = defaultOperatorId;
@@ -630,7 +760,7 @@ contract OperatorStateRetrieverUnitTests is MockAVSDeployer {
         OperatorStateRetriever.Operator[][] memory operators,
         uint256[][] memory expectedOperatorOverallIndices,
         OperatorMetadata[] memory operatorMetadatas
-    ) internal {
+    ) internal pure {
         // for each quorum
         for (uint256 j = 0; j < quorumNumbers.length; j++) {
             // make sure the each operator id and stake is correct
