Fix dynamic proof consumption in the same transaction prover method

packages/module/src/method/runtimeMethod.ts

@@ -161,7 +161,7 @@ export function combineMethodName(
   runtimeModuleName: string,
   methodName: string
 ) {
-  return `${runtimeModuleName}.${methodName}`;
+  return `${runtimeModuleName}.${methodName}` as const;
 }
 
 export const runtimeMethodMetadataKey = "yab-method";

packages/module/src/runtime/MethodIdResolver.ts

@@ -4,6 +4,7 @@ import { Poseidon } from "o1js";
 import { inject, injectable } from "tsyringe";
 
 import {
+  combineMethodName,
   RuntimeMethodInvocationType,
   runtimeMethodTypeMetadataKey,
 } from "../method/runtimeMethod";
@@ -41,6 +42,10 @@ export class MethodIdResolver {
     }, {});
   }
 
+  public getAllRuntimeMethodNames() {
+    return Object.values(this.dictionary);
+  }
+
   /**
    * The purpose of this method is to provide a dictionary where
    * we can look up properties like methodId and invocationType
@@ -64,7 +69,7 @@ export class MethodIdResolver {
 
         if (type !== undefined) {
           return {
-            name: `${moduleName}.${method}`,
+            name: combineMethodName(moduleName, method),
             methodId: methodIdResolver.getMethodId(moduleName, method),
             type,
           } as const;

packages/module/src/runtime/Runtime.ts

@@ -21,6 +21,7 @@ import {
   RuntimeTransaction,
   NetworkState,
 } from "@proto-kit/protocol";
+import chunk from "lodash/chunk";
 
 import {
   combineMethodName,
@@ -65,6 +66,14 @@ const errors = {
     new Error(`Unable to find method with id ${methodKey}`),
 };
 
+type Methods = Record<
+  string,
+  {
+    privateInputs: any;
+    method: AsyncWrappedMethod;
+  }
+>;
+
 export class RuntimeZkProgrammable<
   Modules extends RuntimeModulesRecord,
 > extends ZkProgrammable<undefined, MethodPublicOutput> {
@@ -76,25 +85,10 @@ export class RuntimeZkProgrammable<
     return this.runtime.areProofsEnabled;
   }
 
-  public async zkProgramFactory(): Promise<
-    PlainZkProgram<undefined, MethodPublicOutput>[]
-  > {
-    type Methods = Record<
-      string,
-      {
-        privateInputs: any;
-        method: AsyncWrappedMethod;
-      }
-    >;
-    // We need to use explicit type annotations here,
-    // therefore we can't use destructuring
+  private extractRuntimeMethods() {
+    const { runtime } = this;
 
-    // eslint-disable-next-line prefer-destructuring
-    const runtime: Runtime<Modules> = this.runtime;
-
-    const MAXIMUM_METHODS_PER_ZK_PROGRAM = 8;
-
-    const runtimeMethods = runtime.runtimeModuleNames.reduce<Methods>(
+    return runtime.runtimeModuleNames.reduce<Methods>(
       (allMethods, runtimeModuleName) => {
         runtime.isValidModuleName(runtime.definition, runtimeModuleName);
 
@@ -166,63 +160,27 @@ export class RuntimeZkProgrammable<
       },
       {}
     );
+  }
+
+  public async zkProgramFactory(): Promise<
+    PlainZkProgram<undefined, MethodPublicOutput>[]
+  > {
+    const runtimeMethods = this.extractRuntimeMethods();
 
-    const sortedRuntimeMethods = Object.fromEntries(
-      Object.entries(runtimeMethods).sort()
+    const buckets = this.runtime.bucketRuntimeMethods(
+      Object.keys(runtimeMethods)
     );
 
-    const splitRuntimeMethods = () => {
-      const buckets: Array<
-        Record<
-          string,
-          {
-            privateInputs: any;
-            method: AsyncWrappedMethod;
-          }
-        >
-      > = [];
-      Object.entries(sortedRuntimeMethods).forEach(
-        async ([methodName, method]) => {
-          let methodAdded = false;
-          for (const bucket of buckets) {
-            if (buckets.length === 0) {
-              const record: Record<
-                string,
-                {
-                  privateInputs: any;
-                  method: AsyncWrappedMethod;
-                }
-              > = {};
-              record[methodName] = method;
-              buckets.push(record);
-              methodAdded = true;
-              break;
-            } else if (
-              Object.keys(bucket).length <=
-              MAXIMUM_METHODS_PER_ZK_PROGRAM - 1
-            ) {
-              bucket[methodName] = method;
-              methodAdded = true;
-              break;
-            }
-          }
-          if (!methodAdded) {
-            const record: Record<
-              string,
-              {
-                privateInputs: any;
-                method: AsyncWrappedMethod;
-              }
-            > = {};
-            record[methodName] = method;
-            buckets.push(record);
-          }
-        }
-      );
-      return buckets;
-    };
+    const splitRuntimeMethods = buckets.map((bucket) =>
+      Object.fromEntries(
+        bucket.map((methodName) => {
+          const method = runtimeMethods[methodName];
+          return [methodName, method] as const;
+        })
+      )
+    );
 
-    return splitRuntimeMethods().map((bucket, index) => {
+    return splitRuntimeMethods.map((bucket, index) => {
       const name = `RuntimeProgram-${index}`;
       const wrappedBucket = Object.fromEntries(
         Object.entries(bucket).map(([methodName, methodDef]) => [
@@ -365,6 +323,14 @@ export class Runtime<Modules extends RuntimeModulesRecord>
     return (method as (...args: unknown[]) => Promise<unknown>).bind(module);
   }
 
+  public bucketRuntimeMethods(methods: string[]) {
+    const MAXIMUM_METHODS_PER_ZK_PROGRAM = 8;
+
+    const sorted = methods.slice().sort();
+
+    return chunk(sorted, MAXIMUM_METHODS_PER_ZK_PROGRAM);
+  }
+
   /**
    * Add a name and other respective properties required by RuntimeModules,
    * that come from the current Runtime

packages/sequencer/src/protocol/production/BatchProducerModule.ts

@@ -157,12 +157,14 @@ export class BatchProducerModule extends SequencerModule {
       this.merkleStore
     );
 
+    log.info(`Tracing ${blocks.length} blocks...`);
     const trace = await this.batchTraceService.traceBatch(
       blocks.map((block) => block),
       merkleTreeStore,
       batchId
     );
 
+    log.info("Proving batch...");
     const proof = await this.batchFlow.executeBatch(trace, batchId);
 
     const fromNetworkState = blocks[0].block.networkState.before;

packages/sequencer/src/protocol/production/flow/BlockFlow.ts

@@ -5,8 +5,13 @@ import {
   TransactionProof,
 } from "@proto-kit/protocol";
 import { mapSequential } from "@proto-kit/common";
-// eslint-disable-next-line import/no-extraneous-dependencies
-import chunk from "lodash/chunk";
+import {
+  combineMethodName,
+  MethodIdResolver,
+  Runtime,
+  RuntimeModulesRecord,
+} from "@proto-kit/module";
+import { Memoize } from "typescript-memoize";
 
 import { TransactionProvingTask } from "../tasks/TransactionProvingTask";
 import { FlowCreator } from "../../../worker/flow/Flow";
@@ -24,9 +29,13 @@ export class BlockFlow {
     private readonly flowCreator: FlowCreator,
     @inject("Protocol")
     private readonly protocol: Protocol<MandatoryProtocolModulesRecord>,
+    @inject("Runtime")
+    private readonly runtime: Runtime<RuntimeModulesRecord>,
     private readonly runtimeFlow: TransactionFlow,
     private readonly transactionTask: TransactionProvingTask,
-    private readonly transactionMergeTask: TransactionReductionTask
+    private readonly transactionMergeTask: TransactionReductionTask,
+    @inject("MethodIdResolver")
+    private readonly methodIdResolver: MethodIdResolver
   ) {}
 
   private dummyProof: TransactionProof | undefined = undefined;
@@ -46,11 +55,62 @@ export class BlockFlow {
     return dummy;
   }
 
+  @Memoize()
+  private getMethodNameBuckets() {
+    return this.runtime.bucketRuntimeMethods(
+      this.methodIdResolver
+        .getAllRuntimeMethodNames()
+        .map(({ moduleName, methodName }) =>
+          combineMethodName(moduleName, methodName)
+        )
+    );
+  }
+
+  private findZkProgramIndex(trace: TransactionTrace) {
+    const methodName = this.methodIdResolver.getMethodNameFromId(
+      trace.runtime.tx.methodId.toBigInt()
+    )!;
+
+    return this.getMethodNameBuckets().findIndex((bucket) =>
+      bucket.includes(combineMethodName(methodName[0], methodName[1]))
+    );
+  }
+
+  private chunkTransactions(tracesInput: TransactionTrace[]) {
+    const chunks = [];
+    const traces = tracesInput.slice().reverse();
+
+    while (traces.length > 0) {
+      const first = traces.pop()!;
+      const firstZkProgramIndex = this.findZkProgramIndex(first);
+
+      const second = traces.at(-1);
+      let secondZkProgramIndex = -1;
+      if (second !== undefined) {
+        secondZkProgramIndex = this.findZkProgramIndex(second);
+      }
+
+      if (
+        second !== undefined &&
+        secondZkProgramIndex === firstZkProgramIndex
+      ) {
+        traces.pop();
+        chunks.push([first, second]);
+      } else {
+        chunks.push([first]);
+      }
+    }
+
+    return chunks;
+  }
+
   private async proveTransactions(height: string, traces: TransactionTrace[]) {
+    const traceChunks = this.chunkTransactions(traces);
+
     const flow = new ReductionTaskFlow(
       {
         name: `transaction-${height}`,
-        inputLength: Math.ceil(traces.length / 2),
+        inputLength: traceChunks.length,
         mappingTask: this.transactionTask,
         reductionTask: this.transactionMergeTask,
         mergableFunction: (a, b) =>
@@ -66,7 +126,7 @@ export class BlockFlow {
       this.flowCreator
     );
 
-    await mapSequential(chunk(traces, 2), async (traceChunk, index) => {
+    await mapSequential(traceChunks, async (traceChunk, index) => {
       await this.runtimeFlow.proveRuntimes(
         traceChunk,
         height,

packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts

@@ -129,7 +129,6 @@ export class TransactionProvingTask
       return await this.computeDummy();
     }
 
-    // eslint-disable-next-line @typescript-eslint/no-shadow
     const DynamicRuntimeProof =
       await this.runtime.zkProgrammable.dynamicProofType();
 

packages/sequencer/src/protocol/production/tracing/BatchTracingService.ts

@@ -1,4 +1,4 @@
-import { log, range, unzip, yieldSequential } from "@proto-kit/common";
+import { range, unzip, yieldSequential } from "@proto-kit/common";
 import {
   AppliedBatchHashList,
   MinaActionsHashList,
@@ -68,8 +68,6 @@ export class BatchTracingService {
 
   @trace("batch.trace.blocks")
   public async traceBlocks(blocks: BlockWithResult[]) {
-    log.debug(`Tracing ${blocks.length} blocks...`);
-
     const batchState = this.createBatchState(blocks[0]);
 
     const publicInput = this.blockTracingService.openBatch(