Optimize fork-choice iterators

packages/fork-choice/src/forkChoice/forkChoice.ts

@@ -140,7 +140,7 @@ export class ForkChoice implements IForkChoice {
     if (block.slot > ancestorSlot) {
       // Search for a slot that is lte the target slot.
       // We check for lower slots to account for skip slots.
-      for (const node of this.protoArray.iterateNodes(blockRootHex)) {
+      for (const node of this.protoArray.iterateAncestorNodes(blockRootHex)) {
         if (node.slot <= ancestorSlot) {
           return fromHexString(node.blockRoot);
         }
@@ -559,26 +559,40 @@ export class ForkChoice implements IForkChoice {
     this.protoArray.pruneThreshold = threshold;
   }
 
+  *iterateAncestorBlocks(blockRoot: phase0.Root): IterableIterator<IBlockSummary> {
+    for (const block of this.protoArray.iterateAncestorNodes(toHexString(blockRoot))) {
+      yield toBlockSummary(block);
+    }
+  }
+
   /**
    * Iterates backwards through block summaries, starting from a block root.
    * Return only the non-finalized blocks.
    */
-  iterateBlockSummaries(blockRoot: phase0.Root): IBlockSummary[] {
-    const blocks = this.protoArray.iterateNodes(toHexString(blockRoot)).map(toBlockSummary);
+  getAllAncestorBlocks(blockRoot: phase0.Root): IBlockSummary[] {
+    const blocks = this.protoArray.getAllAncestorNodes(toHexString(blockRoot)).map(toBlockSummary);
     // the last node is the previous finalized one, it's there to check onBlock finalized checkpoint only.
     return blocks.slice(0, blocks.length - 1);
   }
 
   /**
-   * The same to iterateBlockSummaries but this gets non-ancestor nodes instead of ancestor nodes.
+   * The same to getAllAncestorBlocks but this gets non-ancestor nodes instead of ancestor nodes.
    */
-  iterateNonAncestors(blockRoot: phase0.Root): IBlockSummary[] {
-    return this.protoArray.iterateNonAncestorNodes(toHexString(blockRoot)).map(toBlockSummary);
+  getAllNonAncestorBlocks(blockRoot: phase0.Root): IBlockSummary[] {
+    return this.protoArray.getAllNonAncestorNodes(toHexString(blockRoot)).map(toBlockSummary);
   }
 
   getCanonicalBlockSummaryAtSlot(slot: Slot): IBlockSummary | null {
-    const head = this.getHeadRoot();
-    return this.iterateBlockSummaries(head).find((summary) => summary.slot === slot) || null;
+    if (slot >= this.head.slot) {
+      return this.head;
+    }
+
+    for (const block of this.protoArray.iterateAncestorNodes(toHexString(this.head.blockRoot))) {
+      if (block.slot === slot) {
+        return toBlockSummary(block);
+      }
+    }
+    return null;
   }
 
   forwardIterateBlockSummaries(): IBlockSummary[] {
@@ -591,7 +605,33 @@ export class ForkChoice implements IForkChoice {
   }
 
   getBlockSummariesAtSlot(slot: Slot): IBlockSummary[] {
-    return this.protoArray.nodes.filter((node) => node.slot === slot).map(toBlockSummary);
+    const nodes = this.protoArray.nodes;
+    const blocksAtSlot: IBlockSummary[] = [];
+    for (let i = 0, len = nodes.length; i < len; i++) {
+      const node = nodes[i];
+      if (node.slot === slot) {
+        blocksAtSlot.push(toBlockSummary(node));
+      }
+    }
+    return blocksAtSlot;
+  }
+
+  getCommonAncestorDistance(prevBlock: IBlockSummary, newBlock: IBlockSummary): number | null {
+    const prevNode = this.protoArray.getNode(toHexString(prevBlock.blockRoot));
+    const newNode = this.protoArray.getNode(toHexString(newBlock.blockRoot));
+    if (!prevNode) throw Error(`No node if forkChoice for blockRoot ${prevBlock.blockRoot}`);
+    if (!newNode) throw Error(`No node if forkChoice for blockRoot ${newBlock.blockRoot}`);
+
+    const commonAncestor = this.protoArray.getCommonAncestor(prevNode, newNode);
+    // No common ancestor, should never happen. Return null to not throw
+    if (!commonAncestor) return null;
+
+    // If common node is one of both nodes, then they are direct descendants, return null
+    if (commonAncestor.blockRoot === prevNode.blockRoot || commonAncestor.blockRoot === newNode.blockRoot) {
+      return null;
+    }
+
+    return newNode.slot - commonAncestor.slot;
   }
 
   private updateJustified(justifiedCheckpoint: phase0.Checkpoint, justifiedBalances: number[]): void {

packages/fork-choice/src/forkChoice/interface.ts

@@ -107,20 +107,26 @@ export interface IForkChoice {
   prune(finalizedRoot: phase0.Root): IBlockSummary[];
   setPruneThreshold(threshold: number): void;
   /**
-   * Iterates backwards through block summaries, starting from a block root
+   * Iterates backwards through ancestor block summaries, starting from a block root
    */
-  iterateBlockSummaries(blockRoot: phase0.Root): IBlockSummary[];
+  iterateAncestorBlocks(blockRoot: phase0.Root): IterableIterator<IBlockSummary>;
   /**
-   * The same to iterateBlockSummaries but this gets non-ancestor nodes instead of ancestor nodes.
+   * Returns all ancestor blocks backwards, starting from a block root
    */
-  iterateNonAncestors(blockRoot: phase0.Root): IBlockSummary[];
+  getAllAncestorBlocks(blockRoot: phase0.Root): IBlockSummary[];
+  /**
+   * The same to getAllAncestorBlocks but this gets non-ancestor nodes instead of ancestor nodes.
+   */
+  getAllNonAncestorBlocks(blockRoot: phase0.Root): IBlockSummary[];
   getCanonicalBlockSummaryAtSlot(slot: Slot): IBlockSummary | null;
   /**
    * Iterates forwards through block summaries, exact order is not guaranteed
    */
   forwardIterateBlockSummaries(): IBlockSummary[];
   getBlockSummariesByParentRoot(parentRoot: phase0.Root): IBlockSummary[];
   getBlockSummariesAtSlot(slot: Slot): IBlockSummary[];
+  /** Returns the distance of common ancestor of nodes to newNode. Returns null if newNode is descendant of prevNode */
+  getCommonAncestorDistance(prevBlock: IBlockSummary, newBlock: IBlockSummary): number | null;
 }
 
 export interface ILatestMessage {

packages/fork-choice/src/protoArray/protoArray.ts

@@ -291,7 +291,7 @@ export class ProtoArray {
     }
 
     // Remove the this.indices key/values for all the to-be-deleted nodes
-    for (const nodeIndex of Array.from({length: finalizedIndex}, (_, i) => i)) {
+    for (let nodeIndex = 0; nodeIndex < finalizedIndex; nodeIndex++) {
       const node = this.nodes[nodeIndex];
       if (node === undefined) {
         throw new ProtoArrayError({
@@ -319,7 +319,8 @@ export class ProtoArray {
     }
 
     // Iterate through all the existing nodes and adjust their indices to match the new layout of this.nodes
-    for (const node of this.nodes) {
+    for (let i = 0, len = this.nodes.length; i < len; i++) {
+      const node = this.nodes[i];
       const parentIndex = node.parent;
       if (parentIndex !== undefined) {
         // If node.parent is less than finalizedIndex, set it to undefined
@@ -491,7 +492,37 @@ export class ProtoArray {
   /**
    * Iterate from a block root backwards over nodes
    */
-  iterateNodes(blockRoot: RootHex): IProtoNode[] {
+  *iterateAncestorNodes(blockRoot: RootHex): IterableIterator<IProtoNode> {
+    const startIndex = this.indices.get(blockRoot);
+    if (startIndex === undefined) {
+      return;
+    }
+
+    const node = this.nodes[startIndex];
+    if (node === undefined) {
+      throw new ProtoArrayError({
+        code: ProtoArrayErrorCode.INVALID_NODE_INDEX,
+        index: startIndex,
+      });
+    }
+
+    yield* this.iterateAncestorNodesFromNode(node);
+  }
+
+  /**
+   * Iterate from a block root backwards over nodes
+   */
+  *iterateAncestorNodesFromNode(node: IProtoNode): IterableIterator<IProtoNode> {
+    while (node.parent !== undefined) {
+      node = this.getNodeFromIndex(node.parent);
+      yield node;
+    }
+  }
+
+  /**
+   * get all nodes from a block root backwards
+   */
+  getAllAncestorNodes(blockRoot: RootHex): IProtoNode[] {
     const startIndex = this.indices.get(blockRoot);
     if (startIndex === undefined) {
       return [];
@@ -506,14 +537,7 @@ export class ProtoArray {
     }
     const result: IProtoNode[] = [node];
     while (node.parent !== undefined) {
-      const nodeParent = node.parent;
-      node = this.nodes[nodeParent];
-      if (node === undefined) {
-        throw new ProtoArrayError({
-          code: ProtoArrayErrorCode.INVALID_NODE_INDEX,
-          index: nodeParent,
-        });
-      }
+      node = this.getNodeFromIndex(node.parent);
       result.push(node);
     }
     return result;
@@ -524,7 +548,7 @@ export class ProtoArray {
    * iterateNodes is to find ancestor nodes of a blockRoot.
    * this is to find non-ancestor nodes of a blockRoot.
    */
-  iterateNonAncestorNodes(blockRoot: RootHex): IProtoNode[] {
+  getAllNonAncestorNodes(blockRoot: RootHex): IProtoNode[] {
     const startIndex = this.indices.get(blockRoot);
     if (startIndex === undefined) {
       return [];
@@ -541,13 +565,7 @@ export class ProtoArray {
     let nodeIndex = startIndex;
     while (node.parent !== undefined) {
       const parentIndex = node.parent;
-      node = this.nodes[parentIndex];
-      if (node === undefined) {
-        throw new ProtoArrayError({
-          code: ProtoArrayErrorCode.INVALID_NODE_INDEX,
-          index: parentIndex,
-        });
-      }
+      node = this.getNodeFromIndex(parentIndex);
       // nodes between nodeIndex and parentIndex means non-ancestor nodes
       result.push(...this.getNodesBetween(nodeIndex, parentIndex));
       nodeIndex = parentIndex;
@@ -556,44 +574,10 @@ export class ProtoArray {
     return result;
   }
 
-  nodesAtSlot(slot: Slot): IProtoNode[] {
-    const result: IProtoNode[] = [];
-    for (const node of this.nodes) {
-      if (node.slot === slot) {
-        result.push(node);
-      }
-    }
-    return result;
-  }
-
   hasBlock(blockRoot: RootHex): boolean {
     return this.indices.has(blockRoot);
   }
 
-  getNodeByIndex(blockIndex: number): IProtoNode | undefined {
-    const node = this.nodes[blockIndex];
-    if (!node) {
-      return undefined;
-    }
-
-    return node;
-  }
-
-  getNodesBetween(upperIndex: number, lowerIndex: number): IProtoNode[] {
-    const result = [];
-    for (let index = upperIndex - 1; index > lowerIndex; index--) {
-      const node = this.nodes[index];
-      if (node === undefined) {
-        throw new ProtoArrayError({
-          code: ProtoArrayErrorCode.INVALID_NODE_INDEX,
-          index,
-        });
-      }
-      result.push(node);
-    }
-    return result;
-  }
-
   getNode(blockRoot: RootHex): IProtoNode | undefined {
     const blockIndex = this.indices.get(blockRoot);
     if (blockIndex === undefined) {
@@ -622,7 +606,12 @@ export class ProtoArray {
     if (!ancestorNode) {
       return false;
     }
-    for (const node of this.iterateNodes(descendantRoot)) {
+
+    if (ancestorRoot === descendantRoot) {
+      return true;
+    }
+
+    for (const node of this.iterateAncestorNodes(descendantRoot)) {
       if (node.slot < ancestorNode.slot) {
         return false;
       }
@@ -633,7 +622,68 @@ export class ProtoArray {
     return false;
   }
 
+  getCommonAncestor(prevNode: IProtoNode, newNode: IProtoNode): IProtoNode | null {
+    const isPrevNodeLower = prevNode.slot <= newNode.slot;
+    let lowNode = isPrevNodeLower ? prevNode : newNode;
+    let highNode = isPrevNodeLower ? newNode : prevNode;
+
+    highNode = this.getAncestorAtSlot(highNode, lowNode.slot);
+
+    // Now lowNode and highNode are at the same slot
+    while (lowNode.parent !== undefined && highNode.parent !== undefined) {
+      if (lowNode.blockRoot === highNode.blockRoot) {
+        return lowNode;
+      }
+
+      lowNode = this.getNodeFromIndex(lowNode.parent);
+      highNode = this.getNodeFromIndex(highNode.parent);
+    }
+
+    return null;
+  }
+
   length(): number {
     return this.indices.size;
   }
+
+  private getAncestorAtSlot(node: IProtoNode, slot: number): IProtoNode {
+    for (const parentNode of this.iterateAncestorNodesFromNode(node)) {
+      if (parentNode.slot === slot) {
+        return parentNode;
+      }
+    }
+    throw Error("slot less than finalized block");
+  }
+
+  private getNodeFromIndex(index: number): IProtoNode {
+    const node = this.nodes[index];
+    if (node === undefined) {
+      throw new ProtoArrayError({code: ProtoArrayErrorCode.INVALID_NODE_INDEX, index});
+    }
+    return node;
+  }
+
+  private getNodeByIndex(blockIndex: number): IProtoNode | undefined {
+    const node = this.nodes[blockIndex];
+    if (!node) {
+      return undefined;
+    }
+
+    return node;
+  }
+
+  private getNodesBetween(upperIndex: number, lowerIndex: number): IProtoNode[] {
+    const result = [];
+    for (let index = upperIndex - 1; index > lowerIndex; index--) {
+      const node = this.nodes[index];
+      if (node === undefined) {
+        throw new ProtoArrayError({
+          code: ProtoArrayErrorCode.INVALID_NODE_INDEX,
+          index,
+        });
+      }
+      result.push(node);
+    }
+    return result;
+  }
 }

packages/fork-choice/test/unit/forkChoice/forkChoice.test.ts

@@ -44,7 +44,7 @@ describe("Forkchoice", function () {
     bestJustifiedCheckpoint: {root: fromHexString(finalizedRoot), epoch: genesisEpoch},
   };
 
-  it("iterateBlockSummaries", function () {
+  it("getAllAncestorBlocks", function () {
     protoArr.onBlock(block);
     const forkchoice = new ForkChoice({
       config,
@@ -53,8 +53,8 @@ describe("Forkchoice", function () {
       queuedAttestations: new Set(),
       justifiedBalances: [],
     });
-    const summaries = forkchoice.iterateBlockSummaries(fromHexString(finalizedDesc));
-    // there are 2 blocks in protoArray but iterateBlockSummaries should only return non-finalized blocks
+    const summaries = forkchoice.getAllAncestorBlocks(fromHexString(finalizedDesc));
+    // there are 2 blocks in protoArray but getAllAncestorBlocks should only return non-finalized blocks
     expect(summaries.length).to.be.equals(1, "should not return the finalized block");
     expect(summaries[0]).to.be.deep.equals(toBlockSummary(block), "the block summary is not correct");
   });

packages/lodestar/src/chain/archiver/archiveBlocks.ts

@@ -20,10 +20,10 @@ export async function archiveBlocks(
   finalized: phase0.Checkpoint
 ): Promise<void> {
   // Use fork choice to determine the blocks to archive and delete
-  const allCanonicalSummaries = forkChoice.iterateBlockSummaries(finalized.root);
-  // 1st block in iterateBlockSummaries() is the finalized block itself
+  const allCanonicalSummaries = forkChoice.getAllAncestorBlocks(finalized.root);
+  // 1st block in getAllAncestorBlocks() is the finalized block itself
   // we move it to blockArchive but forkchoice still have it to check next onBlock calls
-  // the next iterateBlockSummaries call does not return this block
+  // the next getAllAncestorBlocks call does not return this block
   let i = 0;
   // this number of blocks per chunk is tested in e2e test blockArchive.test.ts
   const BATCH_SIZE = 1000;
@@ -65,7 +65,7 @@ export async function archiveBlocks(
 
   // deleteNonCanonicalBlocks
   // loop through forkchoice single time
-  const nonCanonicalSummaries = forkChoice.iterateNonAncestors(finalized.root);
+  const nonCanonicalSummaries = forkChoice.getAllNonAncestorBlocks(finalized.root);
   if (nonCanonicalSummaries && nonCanonicalSummaries.length > 0) {
     await db.block.batchDelete(nonCanonicalSummaries.map((summary) => summary.blockRoot));
     logger.verbose("deleteNonCanonicalBlocks", {