New STProver: Protocol side

Author: rpanic

packages/common/src/index.ts

@@ -2,6 +2,7 @@ export * from "./config/ModuleContainer";
 export * from "./config/ConfigurableModule";
 export * from "./config/ChildContainerProvider";
 export * from "./config/ChildContainerCreatable";
+export * from "./config/injectAlias";
 export * from "./types";
 export * from "./zkProgrammable/ZkProgrammable";
 export * from "./zkProgrammable/ProvableMethodExecutionContext";

packages/common/src/trees/RollupMerkleTree.ts

@@ -222,6 +222,7 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
     public static WITNESS = RollupMerkleWitness;
 
     // private in interface
+    // TODO Cache this in some static variable so that we don't recompute it every time
     readonly zeroes: bigint[];
 
     readonly store: MerkleTreeStore;

packages/common/src/utils.ts

@@ -7,6 +7,7 @@ import {
 } from "o1js";
 
 import { TypedClass } from "./types";
+import _ from "lodash";
 
 export function requireTrue(
   condition: boolean,
@@ -68,6 +69,26 @@ export function reduceSequential<T, U>(
   );
 }
 
+export function yieldSequential<Source, State, Target>(
+  array: Source[],
+  callbackfn: (
+    previousValue: State,
+    currentValue: Source,
+    currentIndex: number,
+    array: Source[]
+  ) => Promise<[State, Target]>,
+  initialValue: State
+): Promise<[State, Target[]]> {
+  return reduceSequential<Source, [State, Target[]]>(
+    array,
+    async ([state, collectedTargets], curr, index, arr) => {
+      const [newState, addition] = await callbackfn(state, curr, index, arr);
+      return [newState, collectedTargets.concat(addition)];
+    },
+    [initialValue, []]
+  );
+}
+
 export function mapSequential<T, R>(
   array: T[],
   f: (element: T, index: number, array: T[]) => Promise<R>
@@ -194,3 +215,39 @@ export function safeParseJson<T>(json: string) {
   // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
   return JSON.parse(json) as T;
 }
+
+export function isFull<T>(t: Partial<T>): t is T {
+  return Object.values(t).findIndex((v) => v === undefined) === -1;
+}
+
+// TODO Restructure utils into separate package and multiple files
+
+export function padArray<T>(
+  batch: T[],
+  batchSize: number,
+  generator: (index: number) => T
+): T[] {
+  const slice = batch.slice();
+  const dummies = range(0, batchSize - (batch.length % batchSize)).map((i) =>
+    generator(i + batch.length)
+  );
+  slice.push(...dummies);
+  return slice;
+}
+
+export function batch<T>(
+  arr: T[],
+  batchSize: number,
+  dummy: (index: number) => T
+): T[][] {
+  const padded = padArray(arr, batchSize, dummy);
+
+  const partitioned = _.groupBy(
+    padded.map((v, i) => [v, i] as const),
+    ([v, i]) => Math.floor(i / batchSize)
+  );
+
+  const numBatches = Math.floor(arr.length / batchSize);
+
+  return range(0, numBatches).map((i) => partitioned[i].map((x) => x[0]));
+}

packages/common/src/zkProgrammable/provableMethod.ts

@@ -31,7 +31,9 @@ export function toProver(
     );
 
     if (zkProgram === undefined) {
-      throw new Error("Correct ZkProgram not found");
+      throw new Error(
+        `Correct ZkProgram not found (searching for method ${methodName})`
+      );
     }
 
     if (areProofsEnabled) {

packages/module/src/method/runtimeMethod.ts

@@ -37,7 +37,7 @@ const errors = {
 };
 
 export function toStateTransitionsHash(
-  stateTransitions: StateTransition<any>[]
+  stateTransitions: { toProvable: () => ProvableStateTransition }[]
 ) {
   const stateTransitionsHashList = new StateTransitionReductionList(
     ProvableStateTransition

packages/protocol/src/hooks/AccountStateHook.ts

@@ -1,10 +1,13 @@
 import { PublicKey, Struct, UInt64 } from "o1js";
 import { injectable } from "tsyringe";
+import { noop } from "@proto-kit/common";
 
-import { BlockProverExecutionData } from "../prover/block/BlockProvable";
 import { StateMap } from "../state/StateMap";
 import { protocolState } from "../state/protocol/ProtocolState";
-import { ProvableTransactionHook } from "../protocol/ProvableTransactionHook";
+import {
+  ProvableTransactionHook,
+  BeforeTransactionHookArguments,
+} from "../protocol/ProvableTransactionHook";
 import { assert } from "../state/assert/assert";
 
 export class AccountState extends Struct({
@@ -18,7 +21,7 @@ export class AccountStateHook extends ProvableTransactionHook {
     AccountState
   );
 
-  public async onTransaction({ transaction }: BlockProverExecutionData) {
+  public async onTransaction({ transaction }: BeforeTransactionHookArguments) {
     const sender = transaction.sender.value;
 
     const aso = await this.accountState.get(sender);
@@ -45,4 +48,8 @@ export class AccountStateHook extends ProvableTransactionHook {
 
     await this.accountState.set(sender, new AccountState({ nonce: newNonce }));
   }
+
+  public async onAfterTransaction() {
+    noop();
+  }
 }

packages/protocol/src/hooks/LastStateRootBlockHook.ts

@@ -1,26 +1,25 @@
-import { ProvableBlockHook } from "../protocol/ProvableBlockHook";
+import {
+  AfterBlockHookArguments,
+  ProvableBlockHook,
+} from "../protocol/ProvableBlockHook";
 import { NetworkState } from "../model/network/NetworkState";
-import { BlockProverState } from "../prover/block/BlockProver";
 
 export class LastStateRootBlockHook extends ProvableBlockHook<
   Record<string, never>
 > {
   public async afterBlock(
     networkState: NetworkState,
-    state: BlockProverState
+    { stateRoot }: AfterBlockHookArguments
   ): Promise<NetworkState> {
     return new NetworkState({
       block: networkState.block,
       previous: {
-        rootHash: state.stateRoot,
+        rootHash: stateRoot,
       },
     });
   }
 
-  public async beforeBlock(
-    networkState: NetworkState,
-    state: BlockProverState
-  ): Promise<NetworkState> {
+  public async beforeBlock(networkState: NetworkState): Promise<NetworkState> {
     return networkState;
   }
 }

packages/protocol/src/hooks/NoopBlockHook.ts

@@ -1,20 +1,23 @@
 import { NoConfig } from "@proto-kit/common";
 
-import { ProvableBlockHook } from "../protocol/ProvableBlockHook";
+import {
+  AfterBlockHookArguments,
+  BeforeBlockHookArguments,
+  ProvableBlockHook,
+} from "../protocol/ProvableBlockHook";
 import { NetworkState } from "../model/network/NetworkState";
-import { BlockProverState } from "../prover/block/BlockProver";
 
 export class NoopBlockHook extends ProvableBlockHook<NoConfig> {
   public async afterBlock(
     networkState: NetworkState,
-    state: BlockProverState
+    state: AfterBlockHookArguments
   ): Promise<NetworkState> {
     return networkState;
   }
 
   public async beforeBlock(
     networkState: NetworkState,
-    state: BlockProverState
+    state: BeforeBlockHookArguments
   ): Promise<NetworkState> {
     return networkState;
   }

packages/protocol/src/hooks/NoopTransactionHook.ts

@@ -1,10 +1,13 @@
 import { noop } from "@proto-kit/common";
 
 import { ProvableTransactionHook } from "../protocol/ProvableTransactionHook";
-import { BlockProverExecutionData } from "../prover/block/BlockProvable";
 
 export class NoopTransactionHook extends ProvableTransactionHook {
-  public async onTransaction(executionData: BlockProverExecutionData) {
+  public async onTransaction() {
+    noop();
+  }
+
+  public async onAfterTransaction() {
     noop();
   }
 }

packages/protocol/src/index.ts

@@ -15,8 +15,12 @@ export * from "./utils/ProvableHashList";
 export * from "./utils/PrefixedProvableHashList";
 export * from "./utils/MinaPrefixedProvableHashList";
 export * from "./utils/ProvableReductionHashList";
-export * from "./utils/StateTransitionReductionList";
 export * from "./utils/utils";
+export * from "./utils/FieldOptions";
+export * from "./prover/accumulators/StateTransitionReductionList";
+export * from "./prover/accumulators/AppliedBatchHashList";
+export * from "./prover/accumulators/WitnessedRootHashList";
+export * from "./prover/accumulators/TransactionHashList";
 export * from "./prover/block/BlockProver";
 export * from "./prover/block/BlockProvable";
 export * from "./prover/block/accummulators/RuntimeVerificationKeyTree";