Re-org aggregationBits utils logic

packages/beacon-state-transition/src/allForks/util/epochContext.ts

@@ -1,4 +1,4 @@
-import {ByteVector, hash, toHexString, BitList, List, isTreeBacked, TreeBacked} from "@chainsafe/ssz";
+import {ByteVector, hash, toHexString, BitList, List} from "@chainsafe/ssz";
 import bls, {CoordType, PublicKey} from "@chainsafe/bls";
 import {
   BLSSignature,
@@ -10,7 +10,6 @@ import {
   allForks,
   altair,
   Gwei,
-  ssz,
 } from "@chainsafe/lodestar-types";
 import {IBeaconConfig} from "@chainsafe/lodestar-config";
 import {
@@ -34,7 +33,7 @@ import {
   getSeed,
   getTotalActiveBalance,
   isAggregatorFromCommitteeLength,
-  zipIndexesInBitList,
+  zipIndexesCommitteeBits,
 } from "../../util";
 import {getNextSyncCommitteeIndices} from "../../altair/epoch/sync_committee";
 import {computeEpochShuffling, IEpochShuffling} from "./epochShuffling";
@@ -406,21 +405,8 @@ export class EpochContext {
   getIndexedAttestation(attestation: phase0.Attestation): phase0.IndexedAttestation {
     const {aggregationBits, data} = attestation;
     const committeeIndices = this.getBeaconCommittee(data.slot, data.index);
-    let attestingIndices: phase0.ValidatorIndex[];
-    if (isTreeBacked(attestation)) {
-      attestingIndices = zipIndexesInBitList(
-        committeeIndices,
-        (attestation.aggregationBits as unknown) as TreeBacked<BitList>,
-        ssz.phase0.CommitteeBits
-      );
-    } else {
-      attestingIndices = [];
-      for (const [i, index] of committeeIndices.entries()) {
-        if (aggregationBits[i]) {
-          attestingIndices.push(index);
-        }
-      }
-    }
+    const attestingIndices = zipIndexesCommitteeBits(committeeIndices, aggregationBits);
+
     // sort in-place
     attestingIndices.sort((a, b) => a - b);
     return {
@@ -432,9 +418,7 @@ export class EpochContext {
 
   getAttestingIndices(data: phase0.AttestationData, bits: BitList): ValidatorIndex[] {
     const committeeIndices = this.getBeaconCommittee(data.slot, data.index);
-    const validatorIndices = isTreeBacked(bits)
-      ? zipIndexesInBitList(committeeIndices, (bits as unknown) as TreeBacked<BitList>, ssz.phase0.CommitteeBits)
-      : committeeIndices.filter((_, index) => !!bits[index]);
+    const validatorIndices = zipIndexesCommitteeBits(committeeIndices, bits);
     return validatorIndices;
   }
 

packages/beacon-state-transition/src/altair/block/processSyncCommittee.ts

@@ -4,13 +4,13 @@ import {DOMAIN_SYNC_COMMITTEE} from "@chainsafe/lodestar-params";
 import {
   computeEpochAtSlot,
   computeSigningRoot,
-  extractParticipantIndices,
   getBlockRootAtSlot,
   getDomain,
   increaseBalance,
   ISignatureSet,
   SignatureSetType,
   verifySignatureSet,
+  zipIndexesSyncCommitteeBits,
 } from "../../util";
 import {CachedBeaconState} from "../../allForks/util";
 
@@ -88,5 +88,5 @@ function getParticipantIndices(
   syncAggregate: altair.SyncAggregate
 ): number[] {
   const committeeIndices = state.currSyncCommitteeIndexes;
-  return extractParticipantIndices(committeeIndices, syncAggregate);
+  return zipIndexesSyncCommitteeBits(committeeIndices, syncAggregate.syncCommitteeBits);
 }

packages/beacon-state-transition/src/phase0/epoch/processPendingAttestations.ts

@@ -1,7 +1,7 @@
 import {allForks, Epoch, phase0, ssz} from "@chainsafe/lodestar-types";
-import {BitList, List, readonlyValues, TreeBacked} from "@chainsafe/ssz";
+import {List, readonlyValues} from "@chainsafe/ssz";
 import {CachedBeaconState, IAttesterStatus} from "../../allForks/util";
-import {computeStartSlotAtEpoch, getBlockRootAtSlot, zipIndexesInBitList} from "../../util";
+import {computeStartSlotAtEpoch, getBlockRootAtSlot, zipIndexesCommitteeBits} from "../../util";
 
 export function statusProcessEpoch<T extends allForks.BeaconState>(
   state: CachedBeaconState<T>,
@@ -28,11 +28,7 @@ export function statusProcessEpoch<T extends allForks.BeaconState>(
     const attVotedTargetRoot = rootType.equals(attTarget.root, actualTargetBlockRoot);
     const attVotedHeadRoot = rootType.equals(attBeaconBlockRoot, getBlockRootAtSlot(state, attSlot));
     const committee = epochCtx.getBeaconCommittee(attSlot, committeeIndex);
-    const participants = zipIndexesInBitList(
-      committee,
-      aggregationBits as TreeBacked<BitList>,
-      ssz.phase0.CommitteeBits
-    );
+    const participants = zipIndexesCommitteeBits(committee, aggregationBits);
 
     if (epoch === prevEpoch) {
       for (const p of participants) {

packages/beacon-state-transition/src/util/aggregationBits.ts

@@ -1,4 +1,5 @@
-import {BitList, BitListType, BitVectorType, TreeBacked} from "@chainsafe/ssz";
+import {BitList, BitListType, BitVector, isTreeBacked, TreeBacked, Type} from "@chainsafe/ssz";
+import {ssz} from "@chainsafe/lodestar-types";
 
 const BITS_PER_BYTE = 8;
 /** Globally cache this information. @see getUint8ByteToBitBooleanArray */
@@ -30,6 +31,37 @@ function computeUint8ByteToBitBooleanArray(byte: number): boolean[] {
   });
 }
 
+/** zipIndexes for CommitteeBits. @see zipIndexes */
+export function zipIndexesCommitteeBits(indexes: number[], bits: TreeBacked<BitVector> | BitVector): number[] {
+  return zipIndexes(indexes, bits, ssz.phase0.CommitteeBits);
+}
+/** zipIndexes for SyncCommitteeBits. @see zipIndexes */
+export function zipIndexesSyncCommitteeBits(indexes: number[], bits: TreeBacked<BitVector> | BitVector): number[] {
+  return zipIndexes(indexes, bits, ssz.altair.SyncCommitteeBits);
+}
+
+/**
+ * Performant indexing of a BitList, both as struct or TreeBacked
+ * @see zipIndexesInBitListTreeBacked
+ */
+export function zipIndexes<BitArr extends BitList | BitVector>(
+  indexes: number[],
+  bitlist: TreeBacked<BitArr> | BitArr,
+  sszType: Type<BitArr>
+): number[] {
+  if (isTreeBacked<BitArr>(bitlist)) {
+    return zipIndexesTreeBacked(indexes, bitlist, sszType);
+  } else {
+    const attestingIndices = [];
+    for (let i = 0, len = indexes.length; i < len; i++) {
+      if (bitlist[i]) {
+        attestingIndices.push(indexes[i]);
+      }
+    }
+    return attestingIndices;
+  }
+}
+
 /**
  * Returns a new `indexes` array with only the indexes that participated in `bitlist`.
  * Participation of `indexes[i]` means that the bit at position `i` in `bitlist` is true.
@@ -39,19 +71,24 @@ function computeUint8ByteToBitBooleanArray(byte: number): boolean[] {
  * This function uses a precomputed array of booleans `Uint8 -> boolean[]` @see uint8ByteToBitBooleanArrays.
  * This approach is x15 times faster.
  */
-export function zipIndexesInBitList(
+export function zipIndexesTreeBacked<BitArr extends BitList | BitVector>(
   indexes: number[],
-  bitlist: TreeBacked<BitList>,
-  sszType: BitVectorType | BitListType
+  bits: TreeBacked<BitArr>,
+  sszType: Type<BitArr>
 ): number[] {
-  const attBytes = bitlistToUint8Array(bitlist as TreeBacked<BitList>, sszType);
+  const bytes = bitsToUint8Array(bits, sszType);
 
   const indexesSelected: number[] = [];
 
-  // Iterate over each byte of bitlist
-  for (let iByte = 0, byteLen = attBytes.length; iByte < byteLen; iByte++) {
+  // Iterate over each byte of bits
+  for (let iByte = 0, byteLen = bytes.length; iByte < byteLen; iByte++) {
+    // If it's exactly zero, there won't be any indexes, continue early
+    if (bytes[iByte] === 0) {
+      continue;
+    }
+
     // Get the precomputed boolean array for this byte
-    const booleansInByte = getUint8ByteToBitBooleanArray(attBytes[iByte]);
+    const booleansInByte = getUint8ByteToBitBooleanArray(bytes[iByte]);
     // For each bit in the byte check participation and add to indexesSelected array
     for (let iBit = 0; iBit < BITS_PER_BYTE; iBit++) {
       const participantIndex = indexes[iByte * BITS_PER_BYTE + iBit];
@@ -66,20 +103,21 @@ export function zipIndexesInBitList(
 
 /**
  * Efficiently extract the Uint8Array inside a `TreeBacked<BitList>` structure.
- * @see zipIndexesInBitList for reasoning and advantatges.
+ * @see zipIndexesInBitListTreeBacked for reasoning and advantatges.
  */
-export function bitlistToUint8Array(
-  aggregationBits: TreeBacked<BitList>,
-  sszType: BitVectorType | BitListType
+export function bitsToUint8Array<BitArr extends BitList | BitVector>(
+  bits: TreeBacked<BitArr>,
+  sszType: Type<BitArr>
 ): Uint8Array {
-  const tree = aggregationBits.tree;
-  const chunkCount = sszType.tree_getChunkCount(tree);
-  const chunkDepth = sszType.getChunkDepth();
+  const tree = bits.tree;
+  const treeType = (sszType as unknown) as BitListType;
+  const chunkCount = treeType.tree_getChunkCount(tree);
+  const chunkDepth = treeType.getChunkDepth();
   const nodeIterator = tree.iterateNodesAtDepth(chunkDepth, 0, chunkCount);
   const chunks: Uint8Array[] = [];
   for (const node of nodeIterator) {
     chunks.push(node.root);
   }
   // the last chunk has 32 bytes but we don't use all of them
-  return Buffer.concat(chunks).subarray(0, Math.ceil(aggregationBits.length / BITS_PER_BYTE));
+  return Buffer.concat(chunks).subarray(0, Math.ceil(bits.length / BITS_PER_BYTE));
 }

packages/beacon-state-transition/src/util/syncCommittee.ts

@@ -4,12 +4,9 @@ import {
   SYNC_COMMITTEE_SUBNET_COUNT,
   TARGET_AGGREGATORS_PER_SYNC_SUBCOMMITTEE,
 } from "@chainsafe/lodestar-params";
-import {altair} from "@chainsafe/lodestar-types";
-import {ssz} from "@chainsafe/lodestar-types";
 import {BLSSignature, phase0} from "@chainsafe/lodestar-types";
 import {bytesToInt, intDiv} from "@chainsafe/lodestar-utils";
-import {BitList, hash, isTreeBacked, TreeBacked} from "@chainsafe/ssz";
-import {zipIndexesInBitList} from "./aggregationBits";
+import {hash} from "@chainsafe/ssz";
 
 /**
  * TODO
@@ -28,24 +25,3 @@ export function isSyncCommitteeAggregator(selectionProof: BLSSignature): boolean
   );
   return bytesToInt(hash(selectionProof.valueOf() as Uint8Array).slice(0, 8)) % modulo === 0;
 }
-
-export function extractParticipantIndices(
-  committeeIndices: phase0.ValidatorIndex[],
-  syncAggregate: altair.SyncAggregate
-): phase0.ValidatorIndex[] {
-  if (isTreeBacked(syncAggregate)) {
-    return zipIndexesInBitList(
-      committeeIndices,
-      syncAggregate.syncCommitteeBits as TreeBacked<BitList>,
-      ssz.altair.SyncCommitteeBits
-    );
-  } else {
-    const participantIndices = [];
-    for (const [i, index] of committeeIndices.entries()) {
-      if (syncAggregate.syncCommitteeBits[i]) {
-        participantIndices.push(index);
-      }
-    }
-    return participantIndices;
-  }
-}

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

@@ -2,7 +2,7 @@ import {MAX_VALIDATORS_PER_COMMITTEE} from "@chainsafe/lodestar-params";
 import {ssz} from "@chainsafe/lodestar-types";
 import {BenchmarkRunner} from "@chainsafe/lodestar-utils/test_utils/benchmark";
 import {List, readonlyValues} from "@chainsafe/ssz";
-import {zipIndexesInBitList} from "../../../src";
+import {zipIndexesCommitteeBits} from "../../../src";
 
 export async function runAggregationBitsTest(): Promise<void> {
   const runner = new BenchmarkRunner("aggregationBits", {
@@ -25,7 +25,7 @@ export async function runAggregationBitsTest(): Promise<void> {
   await runner.run({
     id: "zipIndexesInBitList",
     run: () => {
-      zipIndexesInBitList(indexes, bitlistTree, ssz.phase0.CommitteeBits);
+      zipIndexesCommitteeBits(indexes, bitlistTree);
     },
   });
 

packages/beacon-state-transition/test/unit/util/aggregationBits.test.ts

@@ -1,8 +1,8 @@
-import {MAX_VALIDATORS_PER_COMMITTEE} from "@chainsafe/lodestar-params";
+import {MAX_VALIDATORS_PER_COMMITTEE, SYNC_COMMITTEE_SIZE} from "@chainsafe/lodestar-params";
 import {ssz} from "@chainsafe/lodestar-types";
 import {List} from "@chainsafe/ssz";
 import {expect} from "chai";
-import {getUint8ByteToBitBooleanArray, bitlistToUint8Array} from "../../../src";
+import {getUint8ByteToBitBooleanArray, bitsToUint8Array, zipIndexesSyncCommitteeBits} from "../../../src";
 
 const BITS_PER_BYTE = 8;
 
@@ -30,7 +30,7 @@ describe("aggregationBits", function () {
   for (const {name, data, numBytes} of testCases) {
     it(name, () => {
       const tree = ssz.phase0.CommitteeBits.createTreeBackedFromStruct(data as List<boolean>);
-      const aggregationBytes = bitlistToUint8Array(tree, ssz.phase0.CommitteeBits);
+      const aggregationBytes = bitsToUint8Array(tree, ssz.phase0.CommitteeBits);
       expect(aggregationBytes.length).to.be.equal(numBytes, "number of bytes is incorrect");
       const aggregationBits: boolean[] = [];
       for (let i = 0; i < tree.length; i++) {
@@ -50,6 +50,32 @@ describe("aggregationBits", function () {
   });
 });
 
+describe("zipIndexesSyncCommitteeBits", function () {
+  const committeeIndices = Array.from({length: SYNC_COMMITTEE_SIZE}, (_, i) => i * 2);
+  // 3 first bits are true
+  const syncCommitteeBits = Array.from({length: SYNC_COMMITTEE_SIZE}, (_, i) => {
+    return i < 3 ? true : false;
+  });
+
+  it("should extract from TreeBacked SyncAggregate", function () {
+    const syncAggregate = ssz.altair.SyncAggregate.defaultTreeBacked();
+    syncAggregate.syncCommitteeBits = syncCommitteeBits;
+    expect(zipIndexesSyncCommitteeBits(committeeIndices, syncAggregate.syncCommitteeBits)).to.be.deep.equal(
+      [0, 2, 4],
+      "Incorrect participant indices from TreeBacked SyncAggregate"
+    );
+  });
+
+  it("should extract from struct SyncAggregate", function () {
+    const syncAggregate = ssz.altair.SyncAggregate.defaultValue();
+    syncAggregate.syncCommitteeBits = syncCommitteeBits;
+    expect(zipIndexesSyncCommitteeBits(committeeIndices, syncAggregate.syncCommitteeBits)).to.be.deep.equal(
+      [0, 2, 4],
+      "Incorrect participant indices from TreeBacked SyncAggregate"
+    );
+  });
+});
+
 /**
  * Get aggregation bit (true/false) from an aggregation bytes array and validator index in committee.
  * Notice: If we want to access the bit in batch, using this method is not efficient, check the performance