@proto-kit/protocol 0.1.1-develop.153

package/src/prover/statetransition/StateTransitionProver.ts

@@ -0,0 +1,270 @@
+import { Experimental, Field, Provable, SelfProof } from "snarkyjs";
+import { inject, injectable } from "tsyringe";
+
+import {
+  MerkleTreeUtils,
+  RollupMerkleWitness,
+} from "../../utils/merkletree/RollupMerkleTree.js";
+import {
+  DefaultProvableHashList,
+  ProvableHashList,
+} from "../../utils/ProvableHashList";
+import { ProvableStateTransition } from "../../model/StateTransition";
+import { StateTransitionProvableBatch } from "../../model/StateTransitionProvableBatch";
+import { constants } from "../../Constants";
+
+import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider.js";
+import {
+  AreProofsEnabled,
+  PlainZkProgram,
+  provableMethod,
+} from "@proto-kit/common";
+import {
+  StateTransitionProvable,
+  StateTransitionProverPublicInput,
+  StateTransitionProof,
+  StateTransitionProverPublicOutput,
+} from "./StateTransitionProvable";
+import { ProtocolModule } from "../../protocol/ProtocolModule";
+import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
+
+const errors = {
+  stateRootNotMatching: (step: string) => `StateRoots not matching ${step}`,
+
+  stateTransitionsHashNotMatching: (step: string) =>
+    `State transitions hash not matching ${step}`,
+
+  merkleWitnessNotCorrect: (index: number) =>
+    `MerkleWitness not valid for StateTransition (${index})`,
+
+  noWitnessProviderSet: () =>
+    new Error(
+      "WitnessProvider not set, set it before you use StateTransitionProvider"
+    ),
+
+  propertyNotMatching: (propertyName: string) => `${propertyName} not matching`,
+};
+
+interface StateTransitionProverExecutionState {
+  stateRoot: Field;
+  stateTransitionList: ProvableHashList<ProvableStateTransition>;
+}
+
+const StateTransitionSelfProofClass = SelfProof<
+  StateTransitionProverPublicInput,
+  StateTransitionProverPublicOutput
+>;
+
+/**
+ * StateTransitionProver is the prover that proves the application of some state
+ * transitions and checks and updates their merkle-tree entries
+ */
+@injectable()
+export class StateTransitionProver
+  extends ProtocolModule<
+    StateTransitionProverPublicInput,
+    StateTransitionProverPublicOutput
+  >
+  implements StateTransitionProvable
+{
+  public constructor(
+    // Injected
+    public readonly witnessProviderReference: StateTransitionWitnessProviderReference
+  ) {
+    super();
+  }
+
+  public zkProgramFactory(): PlainZkProgram<
+    StateTransitionProverPublicInput,
+    StateTransitionProverPublicOutput
+  > {
+    const instance = this;
+
+    const program = Experimental.ZkProgram({
+      publicInput: StateTransitionProverPublicInput,
+      publicOutput: StateTransitionProverPublicOutput,
+
+      methods: {
+        proveBatch: {
+          privateInputs: [StateTransitionProvableBatch],
+
+          method(
+            publicInput: StateTransitionProverPublicInput,
+            batch: StateTransitionProvableBatch
+          ) {
+            return instance.runBatch(publicInput, batch);
+          },
+        },
+
+        merge: {
+          privateInputs: [
+            StateTransitionSelfProofClass,
+            StateTransitionSelfProofClass,
+          ],
+
+          method(
+            publicInput: StateTransitionProverPublicInput,
+            proof1: StateTransitionProof,
+            proof2: StateTransitionProof
+          ) {
+            return instance.merge(publicInput, proof1, proof2);
+          },
+        },
+      },
+    });
+
+    const methods = {
+      proveBatch: program.proveBatch.bind(program),
+      merge: program.merge.bind(program),
+    };
+
+    const SelfProofClass = Experimental.ZkProgram.Proof(program);
+
+    return {
+      compile: program.compile.bind(program),
+      verify: program.verify.bind(program),
+      Proof: SelfProofClass,
+      methods,
+    };
+  }
+
+  private get witnessProvider(): StateTransitionWitnessProvider {
+    const provider = this.witnessProviderReference.getWitnessProvider();
+    if (provider === undefined) {
+      throw errors.noWitnessProviderSet();
+    }
+    return provider;
+  }
+
+  /**
+   * Applies the state transitions to the current stateRoot
+   * and returns the new prover state
+   */
+  public applyTransitions(
+    stateRoot: Field,
+    stateTransitionCommitmentFrom: Field,
+    transitionBatch: StateTransitionProvableBatch
+  ): StateTransitionProverExecutionState {
+    const state: StateTransitionProverExecutionState = {
+      stateRoot,
+
+      stateTransitionList: new DefaultProvableHashList(
+        ProvableStateTransition,
+        stateTransitionCommitmentFrom
+      ),
+    };
+
+    const transitions = transitionBatch.batch;
+    for (
+      let index = 0;
+      index < constants.stateTransitionProverBatchSize;
+      index++
+    ) {
+      this.applyTransition(state, transitions[index], index);
+    }
+
+    return state;
+  }
+
+  /**
+   * Applies a single state transition to the given state
+   * and mutates it in place
+   */
+  public applyTransition(
+    state: StateTransitionProverExecutionState,
+    transition: ProvableStateTransition,
+    index = 0
+  ) {
+    const treeWitness = Provable.witness(RollupMerkleWitness, () =>
+      this.witnessProvider.getWitness(transition.path)
+    );
+    Provable.log(
+      `ST (${index})`,
+      state.stateRoot,
+      " value: ",
+      transition.from.value
+    );
+    const membershipValid = MerkleTreeUtils.checkMembership(
+      treeWitness,
+      state.stateRoot,
+      transition.path,
+      transition.from.value
+    );
+    membershipValid
+      .or(transition.from.isSome.not())
+      .assertTrue(errors.merkleWitnessNotCorrect(index));
+
+    const newRoot = MerkleTreeUtils.computeRoot(
+      treeWitness,
+      transition.to.value
+    );
+    state.stateRoot = Provable.if(
+      transition.to.isSome,
+      newRoot,
+      state.stateRoot
+    );
+
+    state.stateTransitionList.pushIf(
+      transition,
+      transition.path.equals(Field(0)).not()
+    );
+  }
+
+  /**
+   * Applies a whole batch of StateTransitions at once
+   */
+  @provableMethod()
+  public runBatch(
+    publicInput: StateTransitionProverPublicInput,
+    batch: StateTransitionProvableBatch
+  ): StateTransitionProverPublicOutput {
+    Provable.log(publicInput);
+    const result = this.applyTransitions(
+      publicInput.stateRoot,
+      publicInput.stateTransitionsHash,
+      batch
+    );
+
+    const output = new StateTransitionProverPublicOutput({
+      stateRoot: result.stateRoot,
+      stateTransitionsHash: result.stateTransitionList.commitment,
+    });
+    Provable.log(output);
+    return output;
+  }
+
+  @provableMethod()
+  public merge(
+    publicInput: StateTransitionProverPublicInput,
+    proof1: StateTransitionProof,
+    proof2: StateTransitionProof
+  ): StateTransitionProverPublicOutput {
+    proof1.verify();
+    proof2.verify();
+
+    // Check state
+    publicInput.stateRoot.assertEquals(
+      proof1.publicInput.stateRoot,
+      errors.stateRootNotMatching("publicInput.from -> proof1.from")
+    );
+    proof1.publicOutput.stateRoot.assertEquals(
+      proof2.publicInput.stateRoot,
+      errors.stateRootNotMatching("proof1.to -> proof2.from")
+    );
+
+    // Check ST list
+    publicInput.stateTransitionsHash.assertEquals(
+      proof1.publicInput.stateTransitionsHash,
+      errors.stateTransitionsHashNotMatching("publicInput.from -> proof1.from")
+    );
+    proof1.publicOutput.stateTransitionsHash.assertEquals(
+      proof2.publicInput.stateTransitionsHash,
+      errors.stateTransitionsHashNotMatching("proof1.to -> proof2.from")
+    );
+
+    return new StateTransitionProverPublicInput({
+      stateRoot: proof2.publicOutput.stateRoot,
+      stateTransitionsHash: proof2.publicOutput.stateTransitionsHash,
+    });
+  }
+}

package/src/prover/statetransition/StateTransitionWitnessProvider.ts

@@ -0,0 +1,24 @@
+import type { Field } from "snarkyjs";
+import { injectable } from "tsyringe";
+
+import { RollupMerkleWitness } from "../../utils/merkletree/RollupMerkleTree.js";
+
+/**
+ * Interface for providing merkle witnesses to the state-transition prover
+ */
+export interface StateTransitionWitnessProvider {
+  /**
+   * Provides the merkle witness corresponding to the given key
+   * @param key Merkle-tree key
+   */
+  getWitness: (key: Field) => RollupMerkleWitness;
+}
+
+@injectable()
+export class NoOpStateTransitionWitnessProvider
+  implements StateTransitionWitnessProvider
+{
+  public getWitness(key: Field): RollupMerkleWitness {
+    return new RollupMerkleWitness({ path: [], isLeft: [] });
+  }
+}

package/src/prover/statetransition/StateTransitionWitnessProviderReference.ts

@@ -0,0 +1,17 @@
+import { injectable, Lifecycle, scoped } from "tsyringe";
+import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider";
+
+@injectable()
+@scoped(Lifecycle.ContainerScoped)
+export class StateTransitionWitnessProviderReference {
+
+  private witnessProvider?: StateTransitionWitnessProvider;
+
+  public setWitnessProvider(provider: StateTransitionWitnessProvider) {
+    this.witnessProvider = provider;
+  }
+
+  public getWitnessProvider(): StateTransitionWitnessProvider | undefined {
+    return this.witnessProvider;
+  }
+}

package/src/utils/PrefixedProvableHashList.ts

@@ -0,0 +1,21 @@
+import { Field, FlexibleProvablePure, Poseidon } from "snarkyjs";
+
+import { ProvableHashList } from "./ProvableHashList.js";
+import { stringToField } from "./utils";
+
+export class PrefixedProvableHashList<Value> extends ProvableHashList<Value> {
+  private readonly prefix: Field;
+
+  public constructor(
+    valueType: FlexibleProvablePure<Value>,
+    prefix: string,
+    internalCommitment: Field = Field(0)
+  ) {
+    super(valueType, internalCommitment);
+    this.prefix = stringToField(prefix);
+  }
+
+  protected hash(elements: Field[]): Field {
+    return Poseidon.hash([this.prefix, ...elements]);
+  }
+}

package/src/utils/ProvableHashList.ts

@@ -0,0 +1,50 @@
+import { Field, Poseidon, FlexibleProvablePure, Bool, Provable } from "snarkyjs";
+
+/**
+ * Utilities for creating a hash list from a given value type.
+ */
+export abstract class ProvableHashList<Value> {
+  public constructor(
+    private readonly valueType: FlexibleProvablePure<Value>,
+    public commitment: Field = Field(0)
+  ) {}
+
+  protected abstract hash(elements: Field[]): Field;
+
+  /**
+   * Converts the provided value to Field[] and appends it to
+   * the current hashlist.
+   *
+   * @param value - Value to be appended to the hash list
+   * @returns Current hash list.
+   */
+  public push(value: Value) {
+    this.commitment = this.hash([
+      this.commitment,
+      ...this.valueType.toFields(value),
+    ]);
+    return this;
+  }
+
+  public pushIf(value: Value, condition: Bool) {
+    const newCommitment = this.hash([
+      this.commitment,
+      ...this.valueType.toFields(value),
+    ]);
+    this.commitment = Provable.if(condition, newCommitment, this.commitment);
+    return this;
+  }
+
+  /**
+   * @returns Traling hash of the current hashlist.
+   */
+  public toField() {
+    return this.commitment;
+  }
+}
+
+export class DefaultProvableHashList<Value> extends ProvableHashList<Value> {
+  public hash(elements: Field[]): Field {
+    return Poseidon.hash(elements);
+  }
+}

package/src/utils/merkletree/InMemoryMerkleTreeStorage.ts

@@ -0,0 +1,99 @@
+// eslint-disable-next-line max-len
+/* eslint-disable @typescript-eslint/no-magic-numbers,@typescript-eslint/no-unnecessary-condition */
+import { log } from "@proto-kit/common";
+
+import { RollupMerkleTree } from "./RollupMerkleTree.js";
+import { AsyncMerkleTreeStore, MerkleTreeStore } from "./MerkleTreeStore";
+
+export class InMemoryMerkleTreeStorage implements MerkleTreeStore {
+  protected readonly nodes: {
+    [key: number]: {
+      [key: string]: bigint;
+    };
+  } = {};
+
+  public getNode(key: bigint, level: number): bigint | undefined {
+    return this.nodes[level]?.[key.toString()];
+  }
+
+  public setNode(key: bigint, level: number, value: bigint): void {
+    (this.nodes[level] ??= {})[key.toString()] = value;
+  }
+}
+
+export class CachedMerkleTreeStore extends InMemoryMerkleTreeStorage {
+  private writeCache: {
+    [key: number]: {
+      [key: string]: bigint;
+    };
+  } = {};
+
+  public constructor(private readonly parent: AsyncMerkleTreeStore) {
+    super();
+  }
+
+  public setNode(key: bigint, level: number, value: bigint) {
+    super.setNode(key, level, value);
+    (this.writeCache[level] ??= {})[key.toString()] = value;
+  }
+
+  public getWrittenNodes(): {
+    [key: number]: {
+      [key: string]: bigint;
+    };
+  } {
+    return this.writeCache;
+  }
+
+  public resetWrittenNodes() {
+    this.writeCache = {};
+  }
+
+  public async preloadKey(index: bigint): Promise<void> {
+    log.debug(`Preloading MT ${index}`);
+    // Algo from RollupMerkleTree.getWitness()
+    const { leafCount, height } = RollupMerkleTree;
+
+    if (index >= leafCount) {
+      index %= leafCount;
+    }
+
+    // eslint-disable-next-line no-warning-comments,max-len
+    // TODO Not practical at the moment. Improve pattern when implementing DB storage
+    for (let level = 0; level < height; level++) {
+      const key = index;
+
+      // eslint-disable-next-line no-await-in-loop
+      const value = await this.parent.getNode(key, level);
+      if (level === 0) {
+        log.debug(`Preloaded ${key} -> ${value ?? "-"}`);
+      }
+      if (value !== undefined) {
+        (this.nodes[level] ??= {})[key.toString()] = value;
+      }
+      index /= 2n;
+    }
+  }
+
+  public async mergeIntoParent(): Promise<void> {
+    // In case no state got set we can skip this step
+    if (Object.keys(this.writeCache).length === 0) {
+      return;
+    }
+
+    this.parent.openTransaction();
+    const { height } = RollupMerkleTree;
+    const nodes = this.getWrittenNodes();
+
+    const promises = Array.from({ length: height }).flatMap((ignored, level) =>
+      Object.entries(nodes[level]).map(async (entry) => {
+        await this.parent.setNode(BigInt(entry[0]), level, entry[1]);
+      })
+    );
+
+    await Promise.all(promises);
+
+    this.parent.commit();
+    this.resetWrittenNodes();
+  }
+}

package/src/utils/merkletree/MerkleTreeStore.ts

@@ -0,0 +1,15 @@
+export interface AsyncMerkleTreeStore {
+  openTransaction: () => void;
+
+  commit: () => void;
+
+  setNode: (key: bigint, level: number, value: bigint) => Promise<void>;
+
+  getNode: (key: bigint, level: number) => Promise<bigint | undefined>;
+}
+
+export interface MerkleTreeStore {
+  setNode: (key: bigint, level: number, value: bigint) => void;
+
+  getNode: (key: bigint, level: number) => bigint | undefined;
+}