Class DehydratedBlockRegistry

packages/core/src/defer/registry.ts:34–117

 * functions post hydration.
 */
export class DehydratedBlockRegistry {
  private registry = new Map<string, DehydratedDeferBlock>();
  private cleanupFns = new Map<string, Function[]>();
  private jsActionMap: Map<string, Set<Element>> = inject(JSACTION_BLOCK_ELEMENT_MAP);
  private contract: EventContractDetails = inject(JSACTION_EVENT_CONTRACT);

  add(blockId: string, info: DehydratedDeferBlock) {
    this.registry.set(blockId, info);
    // It's possible that hydration is queued that's waiting for the
    // resolution of a lazy loaded route. In this case, we ensure
    // the callback function is called to continue the hydration process
    // for the queued block set.
    if (this.awaitingCallbacks.has(blockId)) {
      const awaitingCallbacks = this.awaitingCallbacks.get(blockId)!;
      for (const cb of awaitingCallbacks) {
        cb();
      }
    }
  }

  get(blockId: string): DehydratedDeferBlock | null {
    return this.registry.get(blockId) ?? null;
  }

  has(blockId: string): boolean {
    return this.registry.has(blockId);
  }

  cleanup(hydratedBlocks: string[]) {
    removeListenersFromBlocks(hydratedBlocks, this.jsActionMap);
    for (let blockId of hydratedBlocks) {
      this.registry.delete(blockId);
      this.jsActionMap.delete(blockId);
      this.invokeTriggerCleanupFns(blockId);
      this.hydrating.delete(blockId);
      this.awaitingCallbacks.delete(blockId);
    }
    if (this.size === 0) {
      this.contract.instance?.cleanUp();
    }
  }

  get size(): number {
    return this.registry.size;
  }

  // we have to leave the lowest block Id in the registry
  // unless that block has no children
  addCleanupFn(blockId: string, fn: Function) {
    let cleanupFunctions: Function[] = [];
    if (this.cleanupFns.has(blockId)) {
      cleanupFunctions = this.cleanupFns.get(blockId)!;
    }
    cleanupFunctions.push(fn);
    this.cleanupFns.set(blockId, cleanupFunctions);
  }

  invokeTriggerCleanupFns(blockId: string) {