From 0d7958a110444d5368a3a0f6548bffb749e8a88a Mon Sep 17 00:00:00 2001 From: Raphael Panic Date: Fri, 3 Apr 2026 19:00:47 +0200 Subject: [PATCH 1/4] Refactored Runtime zkprogram bucketing into reusable functions --- packages/module/src/method/runtimeMethod.ts | 2 +- .../module/src/runtime/MethodIdResolver.ts | 7 +- packages/module/src/runtime/Runtime.ts | 108 ++++++------------ .../prover/transaction/TransactionProvable.ts | 1 - 4 files changed, 44 insertions(+), 74 deletions(-) diff --git a/packages/module/src/method/runtimeMethod.ts b/packages/module/src/method/runtimeMethod.ts index ffa8145af..2375e3104 100644 --- a/packages/module/src/method/runtimeMethod.ts +++ b/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"; diff --git a/packages/module/src/runtime/MethodIdResolver.ts b/packages/module/src/runtime/MethodIdResolver.ts index e134511e4..754481a7f 100644 --- a/packages/module/src/runtime/MethodIdResolver.ts +++ b/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; diff --git a/packages/module/src/runtime/Runtime.ts b/packages/module/src/runtime/Runtime.ts index 811fcf380..a6be26e0a 100644 --- a/packages/module/src/runtime/Runtime.ts +++ b/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 { @@ -76,25 +85,10 @@ export class RuntimeZkProgrammable< return this.runtime.areProofsEnabled; } - public async zkProgramFactory(): Promise< - PlainZkProgram[] - > { - 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 = this.runtime; - - const MAXIMUM_METHODS_PER_ZK_PROGRAM = 8; - - const runtimeMethods = runtime.runtimeModuleNames.reduce( + return runtime.runtimeModuleNames.reduce( (allMethods, runtimeModuleName) => { runtime.isValidModuleName(runtime.definition, runtimeModuleName); @@ -166,63 +160,27 @@ export class RuntimeZkProgrammable< }, {} ); + } + + public async zkProgramFactory(): Promise< + PlainZkProgram[] + > { + 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 return (method as (...args: unknown[]) => Promise).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 diff --git a/packages/protocol/src/prover/transaction/TransactionProvable.ts b/packages/protocol/src/prover/transaction/TransactionProvable.ts index 3c733430f..caaf277be 100644 --- a/packages/protocol/src/prover/transaction/TransactionProvable.ts +++ b/packages/protocol/src/prover/transaction/TransactionProvable.ts @@ -1,4 +1,3 @@ -// eslint-disable-next-line max-classes-per-file import { CompilableModule, WithZkProgrammable } from "@proto-kit/common"; import { DynamicProof, Field, Proof, Signature, Struct, Void } from "o1js"; From bd9f7ff2df311bab9b6539c5a29410dbe830dfb8 Mon Sep 17 00:00:00 2001 From: Raphael Panic Date: Fri, 3 Apr 2026 19:01:44 +0200 Subject: [PATCH 2/4] Made transactions prove only combined with the same bucket --- .../src/protocol/production/flow/BlockFlow.ts | 70 +++++++++++++++++-- .../tasks/TransactionProvingTask.ts | 1 - 2 files changed, 65 insertions(+), 6 deletions(-) diff --git a/packages/sequencer/src/protocol/production/flow/BlockFlow.ts b/packages/sequencer/src/protocol/production/flow/BlockFlow.ts index 874a4f03d..30728138f 100644 --- a/packages/sequencer/src/protocol/production/flow/BlockFlow.ts +++ b/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, + @inject("Runtime") + private readonly runtime: Runtime, 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, diff --git a/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts b/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts index 228c12e86..981a08e4f 100644 --- a/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts +++ b/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts @@ -3,7 +3,6 @@ import { Protocol, ProtocolModulesRecord, StateServiceProvider, - DynamicRuntimeProof, TransactionProvable, TransactionProof, TransactionProverPublicInput, From 97dc9a31ff69fb1d64c9773540ba39e66807e723 Mon Sep 17 00:00:00 2001 From: Raphael Panic Date: Fri, 3 Apr 2026 19:01:53 +0200 Subject: [PATCH 3/4] Better batch production info logs --- .../sequencer/src/protocol/production/BatchProducerModule.ts | 2 ++ .../src/protocol/production/tracing/BatchTracingService.ts | 4 +--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/sequencer/src/protocol/production/BatchProducerModule.ts b/packages/sequencer/src/protocol/production/BatchProducerModule.ts index 1dad75eaa..c1efd0b31 100644 --- a/packages/sequencer/src/protocol/production/BatchProducerModule.ts +++ b/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; diff --git a/packages/sequencer/src/protocol/production/tracing/BatchTracingService.ts b/packages/sequencer/src/protocol/production/tracing/BatchTracingService.ts index 1bf3e10e6..d8667e55b 100644 --- a/packages/sequencer/src/protocol/production/tracing/BatchTracingService.ts +++ b/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( From 8cbfe7a0df7f18b37c1bb61ac6fce53ff7c32077 Mon Sep 17 00:00:00 2001 From: Raphael Panic Date: Fri, 3 Apr 2026 19:18:19 +0200 Subject: [PATCH 4/4] Linting --- .../src/protocol/production/tasks/TransactionProvingTask.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts b/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts index 981a08e4f..4ff333207 100644 --- a/packages/sequencer/src/protocol/production/tasks/TransactionProvingTask.ts +++ b/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();