@proto-kit/protocol 0.1.1-develop.1088 → 0.1.1-develop.1313

package/src/prover/statetransition/StateTransitionProver.ts

@@ -4,6 +4,9 @@ import {
   provableMethod,
   RollupMerkleTreeWitness,
   ZkProgrammable,
+  CompilableModule,
+  type ArtifactRecord,
+  type CompileRegistry,
 } from "@proto-kit/common";
 import { Field, Provable, SelfProof, ZkProgram } from "o1js";
 import { injectable } from "tsyringe";
@@ -27,15 +30,24 @@ import {
   StateTransitionProverPublicInput,
   StateTransitionProverPublicOutput,
 } from "./StateTransitionProvable";
-import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider";
-import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
 
 const errors = {
   propertyNotMatching: (property: string, step: string) =>
     `${property} not matching ${step}`,
 
-  merkleWitnessNotCorrect: (index: number, type: string) =>
-    `MerkleWitness not valid for StateTransition (${index}, type ${type})`,
+  merkleWitnessNotCorrect: (
+    index: number,
+    type: ProvableStateTransitionType
+  ) => {
+    let s = `MerkleWitness not valid for StateTransition (${index}, type unknown)`;
+    Provable.asProver(() => {
+      s = s.replace(
+        "unknown",
+        type.isNormal().toBoolean() ? "normal" : "protocol"
+      );
+    });
+    return s;
+  },
 
   noWitnessProviderSet: () =>
     new Error(
@@ -64,14 +76,13 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
   StateTransitionProverPublicOutput
 > {
   public constructor(
-    private readonly stateTransitionProver: StateTransitionProver,
-    public readonly witnessProviderReference: StateTransitionWitnessProviderReference
+    private readonly stateTransitionProver: StateTransitionProver
   ) {
     super();
   }
 
-  public get appChain(): AreProofsEnabled | undefined {
-    return this.stateTransitionProver.appChain;
+  public get areProofsEnabled(): AreProofsEnabled | undefined {
+    return this.stateTransitionProver.areProofsEnabled;
   }
 
   public zkProgramFactory(): PlainZkProgram<
@@ -86,7 +97,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
       publicOutput: StateTransitionProverPublicOutput,
 
       methods: {
-        proveBatch: {
+        runBatch: {
           privateInputs: [StateTransitionProvableBatch],
 
           async method(
@@ -115,7 +126,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
     });
 
     const methods = {
-      proveBatch: program.proveBatch.bind(program),
+      runBatch: program.runBatch.bind(program),
       merge: program.merge.bind(program),
     };
 
@@ -123,6 +134,7 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
 
     return [
       {
+        name: program.name,
         compile: program.compile.bind(program),
         verify: program.verify.bind(program),
         analyzeMethods: program.analyzeMethods.bind(program),
@@ -132,14 +144,6 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
     ];
   }
 
-  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
@@ -168,12 +172,19 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
 
     const transitions = transitionBatch.batch;
     const types = transitionBatch.transitionTypes;
+    const merkleWitness = transitionBatch.merkleWitnesses;
     for (
       let index = 0;
       index < constants.stateTransitionProverBatchSize;
       index++
     ) {
-      this.applyTransition(state, transitions[index], types[index], index);
+      this.applyTransition(
+        state,
+        transitions[index],
+        types[index],
+        merkleWitness[index],
+        index
+      );
     }
 
     return state;
@@ -187,13 +198,10 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
     state: StateTransitionProverExecutionState,
     transition: ProvableStateTransition,
     type: ProvableStateTransitionType,
+    merkleWitness: RollupMerkleTreeWitness,
     index = 0
   ) {
-    const witness = Provable.witness(RollupMerkleTreeWitness, () =>
-      this.witnessProvider.getWitness(transition.path)
-    );
-
-    const membershipValid = witness.checkMembership(
+    const membershipValid = merkleWitness.checkMembership(
       state.stateRoot,
       transition.path,
       transition.from.value
@@ -201,14 +209,9 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
 
     membershipValid
       .or(transition.from.isSome.not())
-      .assertTrue(
-        errors.merkleWitnessNotCorrect(
-          index,
-          type.isNormal().toBoolean() ? "normal" : "protocol"
-        )
-      );
+      .assertTrue(errors.merkleWitnessNotCorrect(index, type));
 
-    const newRoot = witness.calculateRoot(transition.to.value);
+    const newRoot = merkleWitness.calculateRoot(transition.to.value);
 
     state.stateRoot = Provable.if(
       transition.to.isSome,
@@ -339,19 +342,22 @@ export class StateTransitionProverProgrammable extends ZkProgrammable<
 @injectable()
 export class StateTransitionProver
   extends ProtocolModule
-  implements StateTransitionProvable, StateTransitionProverType
+  implements
+    StateTransitionProvable,
+    StateTransitionProverType,
+    CompilableModule
 {
   public zkProgrammable: StateTransitionProverProgrammable;
 
-  public constructor(
-    // Injected
-    public readonly witnessProviderReference: StateTransitionWitnessProviderReference
-  ) {
+  public constructor() {
     super();
-    this.zkProgrammable = new StateTransitionProverProgrammable(
-      this,
-      witnessProviderReference
-    );
+    this.zkProgrammable = new StateTransitionProverProgrammable(this);
+  }
+
+  public async compile(
+    registry: CompileRegistry
+  ): Promise<void | ArtifactRecord> {
+    return await this.zkProgrammable.compile(registry);
   }
 
   public runBatch(

package/src/settlement/ContractModule.ts

@@ -1,4 +1,11 @@
-import { ConfigurableModule, NoConfig, TypedClass } from "@proto-kit/common";
+import {
+  ArtifactRecord,
+  type CompilableModule,
+  CompileRegistry,
+  ConfigurableModule,
+  NoConfig,
+  TypedClass,
+} from "@proto-kit/common";
 import { SmartContract } from "o1js";
 
 export type SmartContractClassFromInterface<Type> = typeof SmartContract &
@@ -12,13 +19,13 @@ export type SmartContractClassFromInterface<Type> = typeof SmartContract &
  * of SmartContract and implements a certain interface as specified by the
  * ContractType generic.
  */
-export abstract class ContractModule<
-  ContractType,
-  // undefined = no args
-  Arguments = undefined,
-  Config = NoConfig,
-> extends ConfigurableModule<Config> {
-  public abstract contractFactory(
-    args: Arguments
-  ): SmartContractClassFromInterface<ContractType>;
+export abstract class ContractModule<ContractType, Config = NoConfig>
+  extends ConfigurableModule<Config>
+  implements CompilableModule
+{
+  public abstract contractFactory(): SmartContractClassFromInterface<ContractType>;
+
+  public abstract compile(
+    registry: CompileRegistry
+  ): Promise<ArtifactRecord | undefined>;
 }

package/src/settlement/SettlementContractModule.ts

@@ -5,17 +5,15 @@ import {
   ModulesRecord,
   TypedClass,
   noop,
+  StringKeyOf,
 } from "@proto-kit/common";
-import { PublicKey, SmartContract } from "o1js";
+import { Field, PublicKey, SmartContract } from "o1js";
 import { injectable } from "tsyringe";
 
 import { ProtocolEnvironment } from "../protocol/ProtocolEnvironment";
 import { ProtocolModule } from "../protocol/ProtocolModule";
 
-import {
-  ContractModule,
-  SmartContractClassFromInterface,
-} from "./ContractModule";
+import { ContractModule } from "./ContractModule";
 import { DispatchContractProtocolModule } from "./contracts/DispatchContractProtocolModule";
 import { DispatchContractType } from "./contracts/DispatchSmartContract";
 import {
@@ -23,21 +21,24 @@ import {
   SettlementContractProtocolModule,
 } from "./contracts/SettlementContractProtocolModule";
 import { SettlementContractType } from "./contracts/SettlementSmartContract";
+import { BridgeContractType } from "./contracts/BridgeContract";
+import {
+  BridgeContractConfig,
+  BridgeContractProtocolModule,
+} from "./contracts/BridgeContractProtocolModule";
+import { GetContracts } from "./modularity/types";
 
 export type SettlementModulesRecord = ModulesRecord<
-  TypedClass<ContractModule<unknown, unknown, unknown>>
+  TypedClass<ContractModule<unknown, unknown>>
 >;
 
 export type MandatorySettlementModulesRecord = {
   SettlementContract: TypedClass<
-    ContractModule<
-      SettlementContractType,
-      SmartContractClassFromInterface<DispatchContractType>,
-      SettlementContractConfig
-    >
+    ContractModule<SettlementContractType, SettlementContractConfig>
   >;
-  DispatchContract: TypedClass<
-    ContractModule<DispatchContractType, unknown, unknown>
+  DispatchContract: TypedClass<ContractModule<DispatchContractType, unknown>>;
+  BridgeContract: TypedClass<
+    ContractModule<BridgeContractType, BridgeContractConfig>
   >;
 };
 
@@ -66,17 +67,33 @@ export class SettlementContractModule<
     };
   }
 
-  public static fromDefaults() {
-    return SettlementContractModule.from({
+  public static mandatoryModules() {
+    return {
       SettlementContract: SettlementContractProtocolModule,
       DispatchContract: DispatchContractProtocolModule,
-    });
+      BridgeContract: BridgeContractProtocolModule,
+    } as const;
+  }
+
+  public static fromDefaults() {
+    return SettlementContractModule.from(
+      SettlementContractModule.mandatoryModules()
+    );
+  }
+
+  public static with<AdditionalModules extends SettlementModulesRecord>(
+    additionalModules: AdditionalModules
+  ) {
+    return SettlementContractModule.from({
+      ...SettlementContractModule.mandatoryModules(),
+      ...additionalModules,
+    } as const);
   }
 
   // ** For protocol module
   public protocol?: ProtocolEnvironment;
 
-  public get appChain(): AreProofsEnabled | undefined {
+  public get areProofsEnabled(): AreProofsEnabled | undefined {
     return this.protocol?.getAreProofsEnabled();
   }
   // **
@@ -89,25 +106,14 @@ export class SettlementContractModule<
     noop();
   }
 
-  public getContractClasses(): {
-    settlement: SmartContractClassFromInterface<SettlementContractType>;
-    dispatch: SmartContractClassFromInterface<DispatchContractType>;
-  } {
-    const settlementContractKey = "SettlementContract";
-    const dispatchContractKey = "DispatchContract";
-    this.assertIsValidModuleName(settlementContractKey);
-    this.assertIsValidModuleName(dispatchContractKey);
-
-    const settlementModule = this.resolve(settlementContractKey);
-    const dispatchModule = this.resolve(dispatchContractKey);
-
-    const dispatch = dispatchModule.contractFactory(undefined);
-    const settlement = settlementModule.contractFactory(dispatch);
-
-    return {
-      settlement,
-      dispatch,
-    };
+  public getContractClasses(): GetContracts<SettlementModules> {
+    const contracts =
+      // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+      (this.moduleNames as StringKeyOf<SettlementModules>[]).map((name) => {
+        const module = this.resolve(name);
+        return [name, module.contractFactory()];
+      });
+    return Object.fromEntries(contracts);
   }
 
   public createContracts(addresses: {
@@ -117,16 +123,23 @@ export class SettlementContractModule<
     settlement: SettlementContractType & SmartContract;
     dispatch: DispatchContractType & SmartContract;
   } {
-    const { dispatch, settlement } = this.getContractClasses();
+    const { DispatchContract, SettlementContract } = this.getContractClasses();
 
-    // eslint-disable-next-line new-cap
-    const dispatchInstance = new dispatch(addresses.dispatch);
-    // eslint-disable-next-line new-cap
-    const settlementInstance = new settlement(addresses.settlement);
+    const dispatchInstance = new DispatchContract(addresses.dispatch);
+    const settlementInstance = new SettlementContract(addresses.settlement);
 
     return {
       dispatch: dispatchInstance,
       settlement: settlementInstance,
     };
   }
+
+  public createBridgeContract(
+    address: PublicKey,
+    tokenId?: Field
+  ): BridgeContractType & SmartContract {
+    const { BridgeContract } = this.getContractClasses();
+
+    return new BridgeContract(address, tokenId);
+  }
 }

package/src/settlement/contracts/BridgeContract.ts

@@ -0,0 +1,237 @@
+import {
+  AccountUpdate,
+  Bool,
+  Field,
+  method,
+  Permissions,
+  Poseidon,
+  Provable,
+  PublicKey,
+  SmartContract,
+  State,
+  state,
+  Struct,
+  TokenContractV2,
+  TokenId,
+  VerificationKey,
+} from "o1js";
+import { noop, range, TypedClass } from "@proto-kit/common";
+
+import {
+  OUTGOING_MESSAGE_BATCH_SIZE,
+  OutgoingMessageArgumentBatch,
+} from "../messages/OutgoingMessageArgument";
+import { Path } from "../../model/Path";
+import { Withdrawal } from "../messages/Withdrawal";
+
+import type { SettlementContractType } from "./SettlementSmartContract";
+
+export type BridgeContractType = {
+  stateRoot: State<Field>;
+  outgoingMessageCursor: State<Field>;
+
+  rollupOutgoingMessages: (
+    batch: OutgoingMessageArgumentBatch
+  ) => Promise<Field>;
+  redeem: (additionUpdate: AccountUpdate) => Promise<void>;
+
+  deployProvable: (
+    args: VerificationKey | undefined,
+    signedSettlement: boolean,
+    permissions: Permissions,
+    settlementContractAddress: PublicKey
+  ) => Promise<AccountUpdate>;
+
+  updateStateRoot: (root: Field) => Promise<void>;
+};
+
+// Equal to WithdrawalKey
+export class OutgoingMessageKey extends Struct({
+  index: Field,
+  tokenId: Field,
+}) {}
+
+export abstract class BridgeContractBase extends TokenContractV2 {
+  public static args: {
+    SettlementContract:
+      | (TypedClass<SettlementContractType> & typeof SmartContract)
+      | undefined;
+    withdrawalStatePath: [string, string];
+  };
+
+  abstract settlementContractAddress: State<PublicKey>;
+
+  abstract stateRoot: State<Field>;
+
+  abstract outgoingMessageCursor: State<Field>;
+
+  /**
+   * Function to deploy the bridging contract in a provable way, so that it can be
+   * a provable process initiated by the settlement contract with a baked-in vk
+   *
+   * @returns Creates and returns an account update deploying the bridge contract
+   */
+  public async deployProvable(
+    verificationKey: VerificationKey | undefined,
+    signedSettlement: boolean,
+    permissions: Permissions,
+    settlementContractAddress: PublicKey
+  ) {
+    const accountUpdate = this.self;
+
+    if (!signedSettlement) {
+      if (verificationKey === undefined) {
+        throw new Error("Verification Key not provided, can't deploy");
+      }
+      accountUpdate.account.verificationKey.set(verificationKey);
+    }
+
+    accountUpdate.requireSignature();
+    this.account.permissions.set(permissions);
+
+    range(0, 8).forEach((i) => {
+      accountUpdate.update.appState[i] = {
+        isSome: Bool(true),
+        value: Field(0),
+      };
+    });
+
+    this.settlementContractAddress.set(settlementContractAddress);
+
+    accountUpdate.body.mayUseToken = {
+      // Set to true for custom tokens only
+      inheritFromParent: accountUpdate.tokenId.equals(TokenId.default).not(),
+      parentsOwnToken: Bool(false),
+    };
+
+    return accountUpdate;
+  }
+
+  public async approveBase(): Promise<void> {
+    noop();
+  }
+
+  public async updateStateRootBase(root: Field) {
+    this.stateRoot.set(root);
+
+    const settlementContractAddress =
+      this.settlementContractAddress.getAndRequireEquals();
+    const SettlementContractClass = BridgeContractBase.args.SettlementContract;
+    if (SettlementContractClass === undefined) {
+      throw new Error(
+        "Settlement Contract class hasn't been set yet, something is wrong with your module composition"
+      );
+    }
+    const settlementContract = new SettlementContractClass(
+      settlementContractAddress
+    );
+    const accountUpdate = settlementContract.assertStateRoot(root);
+    this.approve(accountUpdate);
+  }
+
+  public async rollupOutgoingMessagesBase(
+    batch: OutgoingMessageArgumentBatch
+  ): Promise<Field> {
+    let counter = this.outgoingMessageCursor.getAndRequireEquals();
+    const stateRoot = this.stateRoot.getAndRequireEquals();
+
+    const [withdrawalModule, withdrawalStateName] =
+      BridgeContractBase.args.withdrawalStatePath;
+    const mapPath = Path.fromProperty(withdrawalModule, withdrawalStateName);
+
+    // Count account creation fee to return later, so that the sender can fund
+    // those accounts with a separate AU
+    let accountCreationFeePaid = Field(0);
+
+    for (let i = 0; i < OUTGOING_MESSAGE_BATCH_SIZE; i++) {
+      const args = batch.arguments[i];
+
+      // Check witness
+      const path = Path.fromKey(mapPath, OutgoingMessageKey, {
+        index: counter,
+        tokenId: this.tokenId,
+      });
+
+      // Process message
+      const { address, amount } = args.value;
+      const isDummy = address.equals(this.address);
+
+      args.witness
+        .checkMembership(
+          stateRoot,
+          path,
+          Poseidon.hash(Withdrawal.toFields(args.value))
+        )
+        .or(isDummy)
+        .assertTrue("Provided Withdrawal witness not valid");
+
+      const tokenAu = this.internal.mint({ address, amount });
+      const isNewAccount = tokenAu.account.isNew.getAndRequireEquals();
+
+      accountCreationFeePaid = accountCreationFeePaid.add(
+        Provable.if(isNewAccount, Field(1e9), Field(0))
+      );
+
+      counter = counter.add(Provable.if(isDummy, Field(0), Field(1)));
+    }
+
+    this.outgoingMessageCursor.set(counter);
+
+    return accountCreationFeePaid;
+  }
+
+  protected async redeemBase(additionUpdate: AccountUpdate) {
+    additionUpdate.body.tokenId.assertEquals(
+      this.tokenId,
+      "Tokenid not same as this bridging contract's tokenId"
+    );
+    additionUpdate.body.balanceChange.sgn
+      .isPositive()
+      .assertTrue("Sign not correct");
+    const amount = additionUpdate.body.balanceChange.magnitude;
+
+    // Burn tokens
+    this.internal.burn({
+      address: additionUpdate.publicKey,
+      amount,
+    });
+
+    // Inherit from parent for custom tokens
+    additionUpdate.body.mayUseToken = {
+      inheritFromParent: this.tokenId.equals(TokenId.default).not(),
+      parentsOwnToken: Bool(false),
+    };
+
+    // Send mina
+    this.approve(additionUpdate);
+    this.balance.subInPlace(amount);
+  }
+}
+
+export class BridgeContract
+  extends BridgeContractBase
+  implements BridgeContractType
+{
+  @state(PublicKey) public settlementContractAddress = State<PublicKey>();
+
+  @state(Field) public stateRoot = State<Field>();
+
+  @state(Field) public outgoingMessageCursor = State<Field>();
+
+  @method
+  public async updateStateRoot(root: Field) {
+    return await this.updateStateRootBase(root);
+  }
+
+  @method.returns(Field)
+  public async rollupOutgoingMessages(
+    batch: OutgoingMessageArgumentBatch
+  ): Promise<Field> {
+    return await this.rollupOutgoingMessagesBase(batch);
+  }
+
+  @method
+  public async redeem(additionUpdate: AccountUpdate) {
+    return await this.redeemBase(additionUpdate);
+  }
+}

package/src/settlement/contracts/BridgeContractProtocolModule.ts

@@ -0,0 +1,43 @@
+import { injectable } from "tsyringe";
+import { CompileRegistry } from "@proto-kit/common";
+
+import { ContractModule } from "../ContractModule";
+
+import {
+  BridgeContract,
+  BridgeContractBase,
+  BridgeContractType,
+} from "./BridgeContract";
+
+export type BridgeContractConfig = {
+  withdrawalStatePath: `${string}.${string}`;
+  withdrawalEventName: string;
+};
+
+@injectable()
+export class BridgeContractProtocolModule extends ContractModule<
+  BridgeContractType,
+  BridgeContractConfig
+> {
+  public contractFactory() {
+    const { config } = this;
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const withdrawalStatePathSplit = config.withdrawalStatePath.split(".") as [
+      string,
+      string,
+    ];
+
+    BridgeContractBase.args = {
+      withdrawalStatePath: withdrawalStatePathSplit,
+      SettlementContract: BridgeContractBase.args?.SettlementContract,
+    };
+
+    return BridgeContract;
+  }
+
+  public async compile(registry: CompileRegistry) {
+    return {
+      BridgeContract: await registry.compile(BridgeContract),
+    };
+  }
+}