First working e2e, fixed dummies in STProver, code cleanup

Author: rpanic

package-lock.json

@@ -200,7 +200,7 @@ export function createMerkleTree(height: number): AbstractMerkleTreeClass {
 
     public static dummy() {
       return new RollupMerkleWitness({
-        isLeft: Array<Bool>(height - 1).fill(Bool(false)),
+        isLeft: Array<Bool>(height - 1).fill(Bool(true)),
         path: Array<Field>(height - 1).fill(Field(0)),
       });
     }

packages/common/src/trees/RollupMerkleTree.ts

@@ -5,9 +5,9 @@ import {
   DynamicProof,
   Proof,
 } from "o1js";
+import _ from "lodash";
 
 import { TypedClass } from "./types";
-import _ from "lodash";
 
 export function requireTrue(
   condition: boolean,
@@ -216,20 +216,24 @@ export function safeParseJson<T>(json: string) {
   return JSON.parse(json) as T;
 }
 
-export function isFull<T>(t: Partial<T>): t is T {
+export type Nullable<T> = {
+  [Key in keyof T]: T[Key] | undefined;
+};
+
+export function isFull<T>(t: Nullable<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[],
+  array: 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)
+  const slice = array.slice();
+  const dummies = range(0, batchSize - (array.length % batchSize)).map((i) =>
+    generator(i + array.length)
   );
   slice.push(...dummies);
   return slice;
@@ -247,7 +251,7 @@ export function batch<T>(
     ([v, i]) => Math.floor(i / batchSize)
   );
 
-  const numBatches = Math.floor(arr.length / batchSize);
+  const numBatches = Math.ceil(arr.length / batchSize);
 
   return range(0, numBatches).map((i) => partitioned[i].map((x) => x[0]));
 }

packages/common/src/utils.ts

@@ -6,11 +6,12 @@ import {
 } from "@proto-kit/module";
 import { inject, injectable } from "tsyringe";
 import {
+  BeforeTransactionHookArguments,
   ProvableTransactionHook,
-  BlockProverExecutionData,
   PublicKeyOption,
 } from "@proto-kit/protocol";
 import { Field, Provable, PublicKey } from "o1js";
+import { noop } from "@proto-kit/common";
 
 import { UInt64 } from "../math/UInt64";
 import { Balance, TokenId } from "../runtime/Balances";
@@ -126,7 +127,7 @@ export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHo
    * @param executionData
    */
   public async onTransaction(
-    executionData: BlockProverExecutionData
+    executionData: BeforeTransactionHookArguments
   ): Promise<void> {
     const feeConfig = Provable.witness(MethodFeeConfigData, () =>
       this.feeAnalyzer.getFeeConfig(
@@ -157,4 +158,8 @@ export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHo
       UInt64.Unsafe.fromField(fee.value)
     );
   }
+
+  public async onAfterTransaction(): Promise<void> {
+    noop();
+  }
 }

packages/library/src/hooks/TransactionFeeHook.ts

@@ -1,7 +1,6 @@
 import { Bool, Field, FlexibleProvablePure, Poseidon } from "o1js";
 import { container } from "tsyringe";
 import {
-  StateTransition,
   ProvableStateTransition,
   MethodPublicOutput,
   RuntimeMethodExecutionContext,

packages/module/src/method/runtimeMethod.ts

@@ -11,6 +11,7 @@ export * from "./model/transaction/RuntimeTransaction";
 export * from "./model/transaction/ValueOption";
 export * from "./model/MethodPublicOutput";
 export * from "./model/RuntimeLike";
+export * from "./model/AppliedStateTransitionBatch";
 export * from "./utils/ProvableHashList";
 export * from "./utils/PrefixedProvableHashList";
 export * from "./utils/MinaPrefixedProvableHashList";

packages/protocol/src/index.ts

@@ -23,6 +23,12 @@ export class ProvableStateTransition extends Struct({
       to: Option.none().toProvable(),
     });
   }
+
+  public static isDummy(stateTransition: ProvableStateTransition) {
+    return stateTransition.path
+      .equals(0)
+      .and(stateTransition.from.isSome.not());
+  }
 }
 
 /**

packages/protocol/src/model/StateTransition.ts

@@ -20,8 +20,6 @@ export class StateTransitionType {
  */
 export class ProvableStateTransitionType extends Struct({
   type: Field,
-  // TODO Remove accumulate or remove the Bool array args
-  accumulate: Bool,
 }) {
   public static get nothing(): ProvableStateTransitionType {
     return this.from(StateTransitionType.nothing);
@@ -35,10 +33,9 @@ export class ProvableStateTransitionType extends Struct({
     return this.from(StateTransitionType.closeAndThrowAway);
   }
 
-  private static from(constant: number, accumulate = false) {
+  private static from(constant: number) {
     return new ProvableStateTransitionType({
       type: Field(constant),
-      accumulate: Bool(accumulate),
     });
   }
 
@@ -48,6 +45,10 @@ export class ProvableStateTransitionType extends Struct({
     // 0^2 == 0 && 1^2 == 1
     return type.mul(type).equals(type);
   }
+
+  public isNothing() {
+    return this.type.equals(ProvableStateTransitionType.nothing.type);
+  }
 }
 
 export class MerkleWitnessBatch extends Struct({
@@ -57,17 +58,12 @@ export class MerkleWitnessBatch extends Struct({
   ),
 }) {}
 
-// export class STProverBoolArray extends Struct({
-//   values: Provable.Array(Bool, constants.stateTransitionProverBatchSize),
-// }) {}
-
-// TODO Name
-export class ProvableStateTransitionInformation extends Struct({
+export class ProvableStateTransitionEntry extends Struct({
   stateTransition: ProvableStateTransition,
   type: ProvableStateTransitionType,
   witnessRoot: Bool,
 }) {
-  public static dummy(): ProvableStateTransitionInformation {
+  public static dummy(): ProvableStateTransitionEntry {
     return {
       stateTransition: ProvableStateTransition.dummy(),
       type: ProvableStateTransitionType.nothing,
@@ -80,47 +76,53 @@ export class ProvableStateTransitionInformation extends Struct({
  * A Batch of StateTransitions to be consumed by the StateTransitionProver
  * to prove multiple STs at once
  *
- * bases: Describes the state root on which the ST will be applied on
- * If it is zero, this means that this ST should connect with the previous one
- * If it is one, this means that the batch should be closed
+ * The batch is formed as an array fo ProvableSTEntries, which have a type and
+ * witnessesRoot flag attached to them.
  */
 export class StateTransitionProvableBatch extends Struct({
   batch: Provable.Array(
-    ProvableStateTransitionInformation,
+    ProvableStateTransitionEntry,
     constants.stateTransitionProverBatchSize
   ),
 }) {
-  // TODO Test
   public static fromBatches(
     batches: {
       stateTransitions: ProvableStateTransition[];
       applied: Bool;
       witnessRoot: Bool;
     }[]
   ): StateTransitionProvableBatch[] {
-    const flattened = batches.flatMap((stBatch, i) =>
-      stBatch.stateTransitions.map<ProvableStateTransitionInformation>(
-        (stateTransition, j, sts) => {
-          return {
-            stateTransition,
-            type:
-              // eslint-disable-next-line no-nested-ternary
-              j === sts.length - 1
-                ? stBatch.applied.toBoolean()
-                  ? ProvableStateTransitionType.closeAndApply
-                  : ProvableStateTransitionType.closeAndThrowAway
-                : ProvableStateTransitionType.nothing,
-            witnessRoot:
-              j === sts.length - 1 ? stBatch.witnessRoot : Bool(false),
-          };
-        }
-      )
-    );
+    const flattened: ProvableStateTransitionEntry[] = [];
+
+    for (const stBatch of batches) {
+      const entries =
+        stBatch.stateTransitions.map<ProvableStateTransitionEntry>(
+          (stateTransition, j, sts) => {
+            return {
+              stateTransition,
+              type:
+                // eslint-disable-next-line no-nested-ternary
+                j === sts.length - 1
+                  ? stBatch.applied.toBoolean()
+                    ? ProvableStateTransitionType.closeAndApply
+                    : ProvableStateTransitionType.closeAndThrowAway
+                  : ProvableStateTransitionType.nothing,
+              witnessRoot: Bool(false),
+            };
+          }
+        );
+
+      flattened.push(...entries);
+
+      if (stBatch.witnessRoot.toBoolean() && flattened.length > 0) {
+        flattened.at(-1)!.witnessRoot = Bool(true);
+      }
+    }
 
     const values = batch(
       flattened,
       constants.stateTransitionProverBatchSize,
-      () => ProvableStateTransitionInformation.dummy()
+      () => ProvableStateTransitionEntry.dummy()
     );
 
     return values.map((stBatch) => {

packages/protocol/src/model/StateTransitionProvableBatch.ts

@@ -25,11 +25,19 @@ export class WitnessedRootHashList extends DefaultProvableHashList<WitnessedRoot
     super(WitnessedRoot, commitment);
   }
 
-  // To be used by the STProver
-  // public pushWitness() {}
-
-  // To be used by the BlockProver
-  public validateWitnessedRoot(
+  /**
+   * To be used by the BlockProver or for tracing
+   *
+   * The main purpose of this method compared to the simple push methods
+   * is for deduplicating witnessed roots. We need to do this because the
+   * STProver can only witness once per batch, therefore if multiple witness
+   * points fall back to the same ST (because any batches in between were empty),
+   * this has to be detected and compensated for.
+   * This function does this using the preimage of the current list state.
+   *
+   * @param preimage The preimage to the **current** state of the list.
+   */
+  public witnessRoot(
     witnessedRoot: WitnessedRoot,
     preimage: Field,
     condition: Bool