feat: use napi-rs pubkey-index-map (#7091)

* feat: use napi-rs pubkey-index-map

* fix: export toMemoryEfficientHexStr()

* fix: state-transition CachedBeaconState unit test

* chore: remove commented code

packages/beacon-node/package.json

@@ -105,6 +105,7 @@
     "@chainsafe/prometheus-gc-stats": "^1.0.0",
     "@chainsafe/ssz": "^0.17.1",
     "@chainsafe/threads": "^1.11.1",
+    "@chainsafe/pubkey-index-map": "2.0.0",
     "@ethersproject/abi": "^5.7.0",
     "@fastify/bearer-auth": "^10.0.1",
     "@fastify/cors": "^10.0.1",

packages/beacon-node/src/api/impl/beacon/state/index.ts

@@ -9,6 +9,7 @@ import {
   getRandaoMix,
 } from "@lodestar/state-transition";
 import {EPOCHS_PER_HISTORICAL_VECTOR} from "@lodestar/params";
+import {fromHex} from "@lodestar/utils";
 import {ApiError} from "../../errors.js";
 import {ApiModules} from "../../types.js";
 import {
@@ -200,7 +201,7 @@ export function getBeaconStateApi({
             }
             balances.push({index: id, balance: state.balances.get(id)});
           } else {
-            const index = headState.epochCtx.pubkey2index.get(id);
+            const index = headState.epochCtx.pubkey2index.get(fromHex(id));
             if (index != null && index <= state.validators.length) {
               balances.push({index, balance: state.balances.get(index)});
             }

packages/beacon-node/src/api/impl/beacon/state/utils.ts

@@ -1,6 +1,7 @@
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {routes} from "@lodestar/api";
 import {FAR_FUTURE_EPOCH, GENESIS_SLOT} from "@lodestar/params";
-import {BeaconStateAllForks, PubkeyIndexMap} from "@lodestar/state-transition";
+import {BeaconStateAllForks} from "@lodestar/state-transition";
 import {BLSPubkey, Epoch, phase0, RootHex, Slot, ValidatorIndex} from "@lodestar/types";
 import {fromHex} from "@lodestar/utils";
 import {CheckpointWithHex, IForkChoice} from "@lodestar/fork-choice";
@@ -187,7 +188,7 @@ export function getStateValidatorIndex(
 
   // typeof id === Uint8Array
   const validatorIndex = pubkey2index.get(id);
-  if (validatorIndex === undefined) {
+  if (validatorIndex === null) {
     return {valid: false, code: 404, reason: "Validator pubkey not found in state"};
   }
   if (validatorIndex >= state.validators.length) {

packages/beacon-node/src/api/impl/validator/index.ts

@@ -1359,7 +1359,7 @@ export function getValidatorApi(
       const filteredRegistrations = registrations.filter((registration) => {
         const {pubkey} = registration.message;
         const validatorIndex = headState.epochCtx.pubkey2index.get(pubkey);
-        if (validatorIndex === undefined) return false;
+        if (validatorIndex === null) return false;
 
         const validator = headState.validators.getReadonly(validatorIndex);
         const status = getValidatorStatus(validator, currentEpoch);

packages/beacon-node/src/chain/chain.ts

@@ -1,5 +1,6 @@
 import path from "node:path";
 import {CompositeTypeAny, TreeView, Type} from "@chainsafe/ssz";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {
   BeaconStateAllForks,
   CachedBeaconStateAllForks,
@@ -10,7 +11,6 @@ import {
   getEffectiveBalanceIncrementsZeroInactive,
   isCachedBeaconState,
   Index2PubkeyCache,
-  PubkeyIndexMap,
   EpochShuffling,
   computeEndSlotAtEpoch,
   computeAnchorCheckpoint,

packages/beacon-node/src/chain/historicalState/getHistoricalState.ts

@@ -1,9 +1,9 @@
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {
   BeaconStateAllForks,
   CachedBeaconStateAllForks,
   DataAvailableStatus,
   ExecutionPayloadStatus,
-  PubkeyIndexMap,
   createCachedBeaconState,
   stateTransition,
 } from "@lodestar/state-transition";

packages/beacon-node/src/chain/historicalState/worker.ts

@@ -1,13 +1,9 @@
 import worker from "node:worker_threads";
 import {Transfer, expose} from "@chainsafe/threads/worker";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {createBeaconConfig, chainConfigFromJson} from "@lodestar/config";
 import {getNodeLogger} from "@lodestar/logger/node";
-import {
-  EpochTransitionStep,
-  PubkeyIndexMap,
-  StateCloneSource,
-  StateHashTreeRootSource,
-} from "@lodestar/state-transition";
+import {EpochTransitionStep, StateCloneSource, StateHashTreeRootSource} from "@lodestar/state-transition";
 import {LevelDbController} from "@lodestar/db";
 import {RegistryMetricCreator, collectNodeJSMetrics} from "../../metrics/index.js";
 import {JobFnQueue} from "../../util/queue/fnQueue.js";

packages/beacon-node/src/chain/interface.ts

@@ -1,4 +1,5 @@
 import {CompositeTypeAny, TreeView, Type} from "@chainsafe/ssz";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {
   UintNum64,
   Root,
@@ -21,7 +22,6 @@ import {
   CachedBeaconStateAllForks,
   EpochShuffling,
   Index2PubkeyCache,
-  PubkeyIndexMap,
 } from "@lodestar/state-transition";
 import {BeaconConfig} from "@lodestar/config";
 import {Logger} from "@lodestar/utils";

packages/beacon-node/src/chain/rewards/attestationsRewards.ts

@@ -24,6 +24,7 @@ import {
   isInInactivityLeak,
 } from "@lodestar/state-transition";
 import {BeaconConfig} from "@lodestar/config";
+import {fromHex} from "@lodestar/utils";
 
 export type AttestationsRewards = routes.beacon.AttestationsRewards;
 type IdealAttestationsReward = routes.beacon.IdealAttestationsReward;
@@ -143,7 +144,7 @@ function computeTotalAttestationsRewardsAltair(
   const {flags} = transitionCache;
   const {epochCtx, config} = state;
   const validatorIndices = validatorIds
-    .map((id) => (typeof id === "number" ? id : epochCtx.pubkey2index.get(id)))
+    .map((id) => (typeof id === "number" ? id : epochCtx.pubkey2index.get(fromHex(id))))
     .filter((index) => index !== undefined); // Validator indices to include in the result
 
   const inactivityPenaltyDenominator = config.INACTIVITY_SCORE_BIAS * INACTIVITY_PENALTY_QUOTIENT_ALTAIR;

packages/beacon-node/test/perf/chain/stateCache/updateUnfinalizedPubkeys.test.ts

@@ -1,10 +1,11 @@
 import {itBench, setBenchOpts} from "@dapplion/benchmark";
-import {Map} from "immutable";
+import {Map as ImmutableMap} from "immutable";
 import {toBufferBE} from "bigint-buffer";
 import {digest} from "@chainsafe/as-sha256";
 import {SecretKey} from "@chainsafe/blst";
-import {ssz} from "@lodestar/types";
-import {type CachedBeaconStateAllForks, PubkeyIndexMap} from "@lodestar/state-transition";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
+import {ValidatorIndex, ssz} from "@lodestar/types";
+import {type CachedBeaconStateAllForks, toMemoryEfficientHexStr} from "@lodestar/state-transition";
 import {bytesToBigInt, intToBytes} from "@lodestar/utils";
 import {InMemoryCheckpointStateCache, BlockStateCacheImpl} from "../../../../src/chain/stateCache/index.js";
 import {BlockStateCache} from "../../../../src/chain/stateCache/types.js";
@@ -31,7 +32,7 @@ describe("updateUnfinalizedPubkeys perf tests", function () {
     itBench({
       id: `updateUnfinalizedPubkeys - updating ${numPubkeysToBeFinalized} pubkeys`,
       beforeEach: async () => {
-        baseState.epochCtx.unfinalizedPubkey2index = Map(unfinalizedPubkey2Index.map);
+        baseState.epochCtx.unfinalizedPubkey2index = ImmutableMap(unfinalizedPubkey2Index);
         baseState.epochCtx.pubkey2index = new PubkeyIndexMap();
         baseState.epochCtx.index2pubkey = [];
 
@@ -80,12 +81,14 @@ describe("updateUnfinalizedPubkeys perf tests", function () {
     });
   }
 
-  function generatePubkey2Index(startIndex: number, endIndex: number): PubkeyIndexMap {
-    const pubkey2Index = new PubkeyIndexMap();
+  type PubkeyHex = string;
+
+  function generatePubkey2Index(startIndex: number, endIndex: number): Map<PubkeyHex, ValidatorIndex> {
+    const pubkey2Index = new Map<string, number>();
     const pubkeys = generatePubkeys(endIndex - startIndex);
 
     for (let i = startIndex; i < endIndex; i++) {
-      pubkey2Index.set(pubkeys[i], i);
+      pubkey2Index.set(toMemoryEfficientHexStr(pubkeys[i]), i);
     }
 
     return pubkey2Index;

packages/beacon-node/test/unit/api/impl/beacon/state/utils.test.ts

@@ -117,7 +117,13 @@ describe("beacon state api utils", function () {
       // "validator id not in state"
       expect(getStateValidatorIndex(String(state.validators.length), state, pubkey2index).valid).toBe(false);
       // "validator pubkey not in state"
-      expect(getStateValidatorIndex("0xabcd", state, pubkey2index).valid).toBe(false);
+      expect(
+        getStateValidatorIndex(
+          "0xa99af0913a2834ef4959637e8d7c4e17f0b63adc587d36ab43510452db3102d0771a4554ea4118a33913827d5ee80b76",
+          state,
+          pubkey2index
+        ).valid
+      ).toBe(false);
     });
 
     it("should return valid: true on validator indices / pubkeys in the state", () => {

packages/beacon-node/test/utils/state.ts

@@ -1,10 +1,10 @@
 import {SecretKey} from "@chainsafe/blst";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {config as minimalConfig} from "@lodestar/config/default";
 import {
   BeaconStateAllForks,
   CachedBeaconStateAllForks,
   createCachedBeaconState,
-  PubkeyIndexMap,
   CachedBeaconStateBellatrix,
   BeaconStateBellatrix,
   CachedBeaconStateElectra,

packages/state-transition/package.json

@@ -66,6 +66,7 @@
     "@lodestar/config": "^1.22.0",
     "@lodestar/params": "^1.22.0",
     "@lodestar/types": "^1.22.0",
+    "@chainsafe/pubkey-index-map": "2.0.0",
     "@lodestar/utils": "^1.22.0",
     "bigint-buffer": "^1.1.5",
     "immutable": "^4.3.2"

packages/state-transition/src/block/processConsolidationRequest.ts

@@ -24,7 +24,7 @@ export function processConsolidationRequest(
   const sourceIndex = state.epochCtx.getValidatorIndex(sourcePubkey);
   const targetIndex = state.epochCtx.getValidatorIndex(targetPubkey);
 
-  if (sourceIndex === undefined || targetIndex === undefined) {
+  if (sourceIndex === null || targetIndex === null) {
     return;
   }
 

packages/state-transition/src/block/processDeposit.ts

@@ -65,7 +65,7 @@ export function applyDeposit(
   const {pubkey, withdrawalCredentials, amount} = deposit;
 
   const cachedIndex = epochCtx.getValidatorIndex(pubkey);
-  if (cachedIndex === undefined || !Number.isSafeInteger(cachedIndex) || cachedIndex >= validators.length) {
+  if (cachedIndex === null || !Number.isSafeInteger(cachedIndex) || cachedIndex >= validators.length) {
     if (isValidDepositSignature(config, pubkey, withdrawalCredentials, amount, deposit.signature)) {
       addValidatorToRegistry(fork, state, pubkey, withdrawalCredentials, amount);
     }

packages/state-transition/src/block/processWithdrawalRequest.ts

@@ -33,7 +33,7 @@ export function processWithdrawalRequest(
   // bail out if validator is not in beacon state
   // note that we don't need to check for 6110 unfinalized vals as they won't be eligible for withdraw/exit anyway
   const validatorIndex = pubkey2index.get(withdrawalRequest.validatorPubkey);
-  if (validatorIndex === undefined) {
+  if (validatorIndex === null) {
     return;
   }
 

packages/state-transition/src/cache/epochCache.ts

@@ -1,5 +1,6 @@
 import {PublicKey} from "@chainsafe/blst";
 import * as immutable from "immutable";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {
   BLSSignature,
   CommitteeIndex,
@@ -53,7 +54,6 @@ import {EffectiveBalanceIncrements, getEffectiveBalanceIncrementsWithLen} from "
 import {BeaconStateAllForks, BeaconStateAltair} from "./types.js";
 import {
   Index2PubkeyCache,
-  PubkeyIndexMap,
   UnfinalizedPubkeyIndexMap,
   syncPubkeys,
   toMemoryEfficientHexStr,
@@ -1023,9 +1023,9 @@ export class EpochCache {
     return this.index2pubkey[index];
   }
 
-  getValidatorIndex(pubkey: Uint8Array | PubkeyHex): ValidatorIndex | undefined {
+  getValidatorIndex(pubkey: Uint8Array): ValidatorIndex | null {
     if (this.isPostElectra()) {
-      return this.pubkey2index.get(pubkey) ?? this.unfinalizedPubkey2index.get(toMemoryEfficientHexStr(pubkey));
+      return this.pubkey2index.get(pubkey) ?? this.unfinalizedPubkey2index.get(toMemoryEfficientHexStr(pubkey)) ?? null;
     } else {
       return this.pubkey2index.get(pubkey);
     }
@@ -1059,17 +1059,20 @@ export class EpochCache {
    * Add finalized validator index and pubkey into finalized cache.
    * Since addFinalizedPubkey() primarily takes pubkeys from unfinalized cache, it can take pubkey hex string directly
    */
-  addFinalizedPubkey(index: ValidatorIndex, pubkey: PubkeyHex | Uint8Array, metrics?: EpochCacheMetrics): void {
+  addFinalizedPubkey(index: ValidatorIndex, pubkeyOrHex: PubkeyHex | Uint8Array, metrics?: EpochCacheMetrics): void {
+    const pubkey = typeof pubkeyOrHex === "string" ? fromHex(pubkeyOrHex) : pubkeyOrHex;
     const existingIndex = this.pubkey2index.get(pubkey);
 
-    if (existingIndex !== undefined) {
+    if (existingIndex !== null) {
       if (existingIndex === index) {
         // Repeated insert.
         metrics?.finalizedPubkeyDuplicateInsert.inc();
         return;
       } else {
         // attempt to insert the same pubkey with different index, should never happen.
-        throw Error("inserted existing pubkey into finalizedPubkey2index cache with a different index");
+        throw Error(
+          `inserted existing pubkey into finalizedPubkey2index cache with a different index, index=${index} priorIndex=${existingIndex}`
+        );
       }
     }
 

packages/state-transition/src/cache/pubkeyCache.ts

@@ -1,4 +1,5 @@
 import {PublicKey} from "@chainsafe/blst";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import * as immutable from "immutable";
 import {ValidatorIndex, phase0} from "@lodestar/types";
 
@@ -39,26 +40,6 @@ export function newUnfinalizedPubkeyIndexMap(): UnfinalizedPubkeyIndexMap {
   return immutable.Map<PubkeyHex, ValidatorIndex>();
 }
 
-export class PubkeyIndexMap {
-  // We don't really need the full pubkey. We could just use the first 20 bytes like an Ethereum address
-  readonly map = new Map<PubkeyHex, ValidatorIndex>();
-
-  get size(): number {
-    return this.map.size;
-  }
-
-  /**
-   * Must support reading with string for API support where pubkeys are already strings
-   */
-  get(key: Uint8Array | PubkeyHex): ValidatorIndex | undefined {
-    return this.map.get(toMemoryEfficientHexStr(key));
-  }
-
-  set(key: Uint8Array | PubkeyHex, value: ValidatorIndex): void {
-    this.map.set(toMemoryEfficientHexStr(key), value);
-  }
-}
-
 /**
  * Checks the pubkey indices against a state and adds missing pubkeys
  *

packages/state-transition/src/cache/syncCommitteeCache.ts

@@ -1,7 +1,7 @@
 import {CompositeViewDU} from "@chainsafe/ssz";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {ssz, ValidatorIndex} from "@lodestar/types";
 import {toPubkeyHex} from "@lodestar/utils";
-import {PubkeyIndexMap} from "./pubkeyCache.js";
 
 type SyncComitteeValidatorIndexMap = Map<ValidatorIndex, number[]>;
 
@@ -82,7 +82,7 @@ function computeSyncCommitteeIndices(
   const pubkeys = syncCommittee.pubkeys.getAllReadonly();
   for (const pubkey of pubkeys) {
     const validatorIndex = pubkey2index.get(pubkey);
-    if (validatorIndex === undefined) {
+    if (validatorIndex === null) {
       throw Error(`SyncCommittee pubkey is unknown ${toPubkeyHex(pubkey)}`);
     }
 

packages/state-transition/src/index.ts

@@ -41,11 +41,11 @@ export {
   EpochCacheError,
   EpochCacheErrorCode,
 } from "./cache/epochCache.js";
+export {toMemoryEfficientHexStr} from "./cache/pubkeyCache.js";
 export {type EpochTransitionCache, beforeProcessEpoch} from "./cache/epochTransitionCache.js";
 
 // Aux data-structures
 export {
-  PubkeyIndexMap,
   type Index2PubkeyCache,
   type UnfinalizedPubkeyIndexMap,
   newUnfinalizedPubkeyIndexMap,

packages/state-transition/test/perf/util.ts

@@ -1,5 +1,6 @@
 import {BitArray, fromHexString} from "@chainsafe/ssz";
 import {PublicKey, SecretKey} from "@chainsafe/blst";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {phase0, ssz, Slot, BeaconState} from "@lodestar/types";
 import {config} from "@lodestar/config/default";
 import {createBeaconConfig, createChainForkConfig} from "@lodestar/config";
@@ -17,7 +18,6 @@ import {
   interopSecretKey,
   computeEpochAtSlot,
   getActiveValidatorIndices,
-  PubkeyIndexMap,
   newFilledArray,
   createCachedBeaconState,
   computeCommitteeCount,

packages/state-transition/test/perf/util/loadState/loadState.test.ts

@@ -1,8 +1,9 @@
 import {itBench, setBenchOpts} from "@dapplion/benchmark";
 import {PublicKey} from "@chainsafe/blst";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {loadState} from "../../../../src/util/loadState/loadState.js";
 import {createCachedBeaconState} from "../../../../src/cache/stateCache.js";
-import {Index2PubkeyCache, PubkeyIndexMap} from "../../../../src/cache/pubkeyCache.js";
+import {Index2PubkeyCache} from "../../../../src/cache/pubkeyCache.js";
 import {generatePerfTestCachedStateAltair} from "../../util.js";
 
 /**

packages/state-transition/test/unit/cachedBeaconState.test.ts

@@ -1,11 +1,11 @@
 import {fromHexString} from "@chainsafe/ssz";
 import {describe, it, expect} from "vitest";
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
 import {ssz} from "@lodestar/types";
 import {toHexString} from "@lodestar/utils";
 import {config as defaultConfig} from "@lodestar/config/default";
 import {createBeaconConfig, createChainForkConfig} from "@lodestar/config";
 import {createCachedBeaconStateTest} from "../utils/state.js";
-import {PubkeyIndexMap} from "../../src/cache/pubkeyCache.js";
 import {createCachedBeaconState, loadCachedBeaconState} from "../../src/cache/stateCache.js";
 import {interopPubkeysCached} from "../utils/interop.js";
 import {modifyStateSameValidator, newStateWithValidators} from "../utils/capella.js";
@@ -83,7 +83,7 @@ describe("CachedBeaconState", () => {
 
     expect(state1.epochCtx.getValidatorIndex(pubkey1)).toBe(index1);
     expect(state2.epochCtx.getValidatorIndex(pubkey1)).toBe(index1);
-    expect(state1.epochCtx.getValidatorIndex(pubkey2)).toBe(undefined);
+    expect(state1.epochCtx.getValidatorIndex(pubkey2)).toBe(null);
     expect(state2.epochCtx.getValidatorIndex(pubkey2)).toBe(index2);
   });