@proto-kit/protocol 0.1.1-develop.191 → 0.1.1-develop.2024

package/src/settlement/ContractArgsRegistry.ts

@@ -0,0 +1,60 @@
+import { injectable, singleton } from "tsyringe";
+import merge from "lodash/merge";
+
+export interface StaticInitializationContract<Args> {
+  getInitializationArgs(): Args;
+}
+
+export type NaiveObjectSchema<Obj> = {
+  [Key in keyof Obj]: undefined extends Obj[Key] ? "Optional" : "Required";
+};
+
+/*
+interface Test {
+  one: string;
+  two?: string;
+}
+
+const x: NaiveObjectSchema<Test> = {
+  one: "Required",
+  two: "Optional",
+};
+*/
+
+@injectable()
+@singleton()
+export class ContractArgsRegistry {
+  args: Record<string, any> = {};
+
+  public addArgs<Type>(name: string, addition: Partial<Type>) {
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const args: Partial<Type> = this.args[name] ?? {};
+    this.args[name] = merge(args, addition);
+  }
+
+  public resetArgs(name: string) {
+    delete this.args[name];
+  }
+
+  public getArgs<Type>(name: string, schema: NaiveObjectSchema<Type>): Type {
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const args = this.args[name] ?? {};
+
+    const missing = Object.entries<"Optional" | "Required">(schema).filter(
+      ([key, type]) => {
+        // We filter only if the key is required and isn't present
+        return type === "Required" && args[key] === undefined;
+      }
+    );
+
+    if (missing.length > 0) {
+      const missingKeys = missing.map(([key]) => key);
+      throw new Error(
+        `Contract args for ${name} not all present, ${missingKeys} are missing`
+      );
+    } else {
+      // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+      return args as Type;
+    }
+  }
+}

package/src/settlement/ContractModule.ts

@@ -0,0 +1,31 @@
+import {
+  ArtifactRecord,
+  type CompilableModule,
+  CompileRegistry,
+  ConfigurableModule,
+  NoConfig,
+  TypedClass,
+} from "@proto-kit/common";
+import { SmartContract } from "o1js";
+
+export type SmartContractClassFromInterface<Type> = typeof SmartContract &
+  TypedClass<Type>;
+
+/**
+ * This module type is used to define a contract module that can be used to
+ * construct and inject smart contract instances.
+ * It defines a method contractFactory, whose arguments can be configured via
+ * the Argument generic. It returns a smart contract class that is a subclass
+ * of SmartContract and implements a certain interface as specified by the
+ * ContractType generic.
+ */
+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

@@ -0,0 +1,150 @@
+import {
+  AreProofsEnabled,
+  ChildContainerProvider,
+  ModuleContainer,
+  ModulesRecord,
+  TypedClass,
+  noop,
+  StringKeyOf,
+} from "@proto-kit/common";
+import { Field, PublicKey, SmartContract } from "o1js";
+import { injectable } from "tsyringe";
+
+import { ProtocolEnvironment } from "../protocol/ProtocolEnvironment";
+import { ProtocolModule } from "../protocol/ProtocolModule";
+
+import { ContractModule } from "./ContractModule";
+import { DispatchContractProtocolModule } from "./contracts/DispatchContractProtocolModule";
+import { DispatchContractType } from "./contracts/DispatchSmartContract";
+import { BridgingSettlementContractModule } from "./contracts/BridgingSettlementContractModule";
+import { BridgeContractType } from "./contracts/BridgeContract";
+import {
+  BridgeContractConfig,
+  BridgeContractProtocolModule,
+} from "./contracts/BridgeContractProtocolModule";
+import { GetContracts, InferContractType } from "./modularity/types";
+import { BridgingSettlementContractType } from "./contracts/settlement/BridgingSettlementContract";
+import { SettlementContractType } from "./contracts/settlement/SettlementBase";
+import {
+  SettlementContractConfig,
+  SettlementSmartContractModule,
+} from "./contracts/SettlementSmartContractModule";
+
+export type SettlementModulesRecord = ModulesRecord<
+  TypedClass<ContractModule<unknown, unknown>>
+>;
+
+export type MandatorySettlementModulesRecord = {
+  SettlementContract: TypedClass<
+    ContractModule<SettlementContractType, SettlementContractConfig>
+  >;
+};
+
+export type BridgingSettlementModulesRecord = {
+  SettlementContract: TypedClass<
+    ContractModule<BridgingSettlementContractType, SettlementContractConfig>
+  >;
+  DispatchContract: TypedClass<ContractModule<DispatchContractType, unknown>>;
+  BridgeContract: TypedClass<
+    ContractModule<BridgeContractType, BridgeContractConfig>
+  >;
+};
+
+@injectable()
+export class SettlementContractModule<
+    SettlementModules extends SettlementModulesRecord,
+  >
+  extends ModuleContainer<SettlementModules>
+  implements ProtocolModule<unknown>
+{
+  public constructor(definition: SettlementModules) {
+    super(definition);
+  }
+
+  public static from<SettlementModules extends SettlementModulesRecord>(
+    modules: SettlementModules
+  ): TypedClass<SettlementContractModule<SettlementModules>> {
+    return class ScopedSettlementContractModule extends SettlementContractModule<SettlementModules> {
+      public constructor() {
+        super(modules);
+      }
+    };
+  }
+
+  public static settlementOnly() {
+    return {
+      SettlementContract: SettlementSmartContractModule,
+    } as const;
+  }
+
+  public static settlementAndBridging() {
+    return {
+      SettlementContract: BridgingSettlementContractModule,
+      DispatchContract: DispatchContractProtocolModule,
+      BridgeContract: BridgeContractProtocolModule,
+    } as const;
+  }
+
+  // ** For protocol module
+  public protocol?: ProtocolEnvironment;
+
+  public get areProofsEnabled(): AreProofsEnabled | undefined {
+    return this.protocol?.getAreProofsEnabled();
+  }
+  // **
+
+  public create(childContainerProvider: ChildContainerProvider) {
+    super.create(childContainerProvider);
+  }
+
+  public async start() {
+    noop();
+  }
+
+  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 createContract<ContractName extends StringKeyOf<SettlementModules>>(
+    contractName: ContractName,
+    address: PublicKey,
+    tokenId?: Field
+  ): InferContractType<SettlementModules[ContractName]> {
+    const module = this.resolve(contractName);
+    const ContractClass = module.contractFactory();
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return new ContractClass(address, tokenId) as InferContractType<
+      SettlementModules[ContractName]
+    >;
+  }
+
+  public createContracts<
+    ContractName extends keyof SettlementModules,
+  >(addresses: {
+    [Key in ContractName]: PublicKey;
+  }): {
+    [Key in ContractName]: SmartContract &
+      InferContractType<SettlementModules[Key]>;
+  } {
+    const classes = this.getContractClasses();
+
+    const obj: Record<string, SmartContract> = {};
+    // eslint-disable-next-line guard-for-in
+    for (const key in addresses) {
+      const ContractClass = classes[key];
+      obj[key] = new ContractClass(addresses[key]);
+    }
+
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return obj as {
+      [Key in keyof SettlementModules]: SmartContract &
+        InferContractType<SettlementModules[Key]>;
+    };
+  }
+}

package/src/settlement/contracts/BridgeContract.ts

@@ -0,0 +1,385 @@
+import {
+  AccountUpdate,
+  Bool,
+  Experimental,
+  Field,
+  method,
+  Permissions,
+  Poseidon,
+  Provable,
+  PublicKey,
+  SmartContract,
+  State,
+  state,
+  Struct,
+  TokenContract,
+  TokenId,
+  VerificationKey,
+} from "o1js";
+import { noop, range, TypedClass } from "@proto-kit/common";
+import { container, injectable, singleton } from "tsyringe";
+
+import {
+  OUTGOING_MESSAGE_BATCH_SIZE,
+  OutgoingMessageArgumentBatch,
+  createMessageStruct,
+  OutgoingMessageArgument,
+} from "../messages/OutgoingMessageArgument";
+import { Path } from "../../model/Path";
+import { OutgoingMessageProcessor } from "../modularity/OutgoingMessageProcessor";
+import { PROTOKIT_FIELD_PREFIXES } from "../../hashing/protokit-prefixes";
+import {
+  ContractArgsRegistry,
+  NaiveObjectSchema,
+  StaticInitializationContract,
+} from "../ContractArgsRegistry";
+
+import type { BridgingSettlementContractType } from "./settlement/BridgingSettlementContract";
+
+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,
+}) {}
+
+@injectable()
+@singleton()
+export class BridgeContractContext {
+  public data: {
+    messageInputs: any[][];
+  } = { messageInputs: [] };
+}
+
+export interface BridgeContractArgs {
+  SettlementContract: TypedClass<BridgingSettlementContractType> &
+    typeof SmartContract;
+  messageProcessors: OutgoingMessageProcessor<unknown>[];
+  batchSize?: number;
+}
+
+export const BridgeContractArgsSchema: NaiveObjectSchema<BridgeContractArgs> = {
+  batchSize: "Optional",
+  SettlementContract: "Required",
+  messageProcessors: "Required",
+};
+
+export abstract class BridgeContractBase
+  extends TokenContract
+  implements StaticInitializationContract<BridgeContractArgs>
+{
+  public constructor(address: PublicKey, tokenId?: Field) {
+    super(address, tokenId);
+  }
+
+  getInitializationArgs(): BridgeContractArgs {
+    return container
+      .resolve(ContractArgsRegistry)
+      .getArgs("BridgeContract", BridgeContractArgsSchema);
+  }
+
+  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) {
+    // It's fine for us to only store the actual root since we only have to
+    // witness values, not update/insert
+    this.stateRoot.set(root);
+
+    const args = this.getInitializationArgs();
+
+    const settlementContractAddress =
+      this.settlementContractAddress.getAndRequireEquals();
+    const SettlementContractClass = args.SettlementContract;
+    const settlementContract = new SettlementContractClass(
+      settlementContractAddress
+    );
+    const accountUpdate = settlementContract.assertStateRoot(root);
+    this.approve(accountUpdate);
+  }
+
+  private batchSize() {
+    return (
+      this.getInitializationArgs().batchSize ?? OUTGOING_MESSAGE_BATCH_SIZE
+    );
+  }
+
+  private executeProcessors(batchIndex: number, args: OutgoingMessageArgument) {
+    const { messageProcessors } = this.getInitializationArgs();
+    return messageProcessors.map((processor, j) => {
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      const value = Experimental.memoizeWitness(processor.type, () => {
+        return container.resolve(BridgeContractContext).data.messageInputs[
+          batchIndex
+        ][j];
+      });
+
+      const MessageType = createMessageStruct(processor.type);
+      const message = new MessageType({
+        messageType: args.messageType,
+        value,
+      });
+      return {
+        messageType: args.messageType,
+        result: processor.processMessage(value, {
+          bridgeContract: {
+            publicKey: this.address,
+            tokenId: this.tokenId,
+          },
+        }),
+        hash: Poseidon.hash(MessageType.toFields(message)),
+      };
+    });
+  }
+
+  public processMessage(
+    batchIndex: number,
+    args: OutgoingMessageArgument,
+    isDummy: Bool
+  ) {
+    const results = this.executeProcessors(batchIndex, args);
+
+    const maxAccountUpdates = Math.max(
+      0,
+      ...results.map(({ result: { accountUpdates } }) => accountUpdates.length)
+    );
+    const AccountUpdateArray = Provable.Array(AccountUpdate, maxAccountUpdates);
+
+    const dummyMessageType =
+      PROTOKIT_FIELD_PREFIXES.OUTGOING_MESSAGE_DUMMY_TYPE;
+    const argMessageType = Provable.if(
+      isDummy,
+      dummyMessageType,
+      args.messageType
+    );
+    const dummyAU = AccountUpdate.dummy();
+
+    const chosen = results
+      .map((a) => {
+        a.result.accountUpdates = a.result.accountUpdates.concat(
+          ...Array<AccountUpdate>(
+            maxAccountUpdates - a.result.accountUpdates.length
+          ).fill(dummyAU)
+        );
+        return a;
+      })
+      .concat({
+        messageType: dummyMessageType,
+        hash: Field(0),
+        result: {
+          accountUpdates: Array<AccountUpdate>(maxAccountUpdates).fill(dummyAU),
+          status: Bool(true),
+          statusMessage: undefined,
+        },
+      })
+      .reduce((a, b) => {
+        const isA = a.messageType.equals(argMessageType);
+
+        const messageType = Provable.if(isA, a.messageType, b.messageType);
+        const hash = Provable.if(isA, a.hash, b.hash);
+        const accountUpdates = Provable.if(
+          isA,
+          AccountUpdateArray,
+          a.result.accountUpdates,
+          b.result.accountUpdates
+        );
+        const status = Provable.if(isA, a.result.status, b.result.status);
+        let statusMessage: string | undefined = undefined;
+        Provable.asProver(() => {
+          if (isA.toBoolean()) {
+            statusMessage = a.result.statusMessage;
+          } else {
+            statusMessage = b.result.statusMessage;
+          }
+        });
+
+        return {
+          messageType,
+          hash,
+          result: { accountUpdates, status, statusMessage },
+        };
+      });
+
+    // If no processor picks up our message type, the reduce above returns
+    // the first candidate. This statement asserts that this is not the case
+    chosen.messageType.assertEquals(
+      argMessageType,
+      "No processor found for message type"
+    );
+
+    return chosen;
+  }
+
+  public async rollupOutgoingMessagesBase(batch: OutgoingMessageArgumentBatch) {
+    let counter = this.outgoingMessageCursor.getAndRequireEquals();
+    const stateRoot = this.stateRoot.getAndRequireEquals();
+
+    // 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 < this.batchSize(); i++) {
+      const args = batch.arguments[i];
+
+      const isDummy = batch.isDummys[i];
+
+      const message = this.processMessage(i, args, isDummy);
+
+      // Check witness
+      const path = Path.fromKey(
+        PROTOKIT_FIELD_PREFIXES.OUTGOING_MESSAGE_BASE_PATH,
+        OutgoingMessageKey,
+        {
+          index: counter,
+          tokenId: this.tokenId,
+        }
+      );
+
+      args.witness
+        .checkMembership(stateRoot, path, message.hash)
+        .or(isDummy)
+        .assertTrue("Provided Withdrawal witness not valid");
+
+      message.result.status.assertTrue(message.result.statusMessage);
+
+      message.result.accountUpdates.forEach((accountUpdate) => {
+        Provable.log("Approving account update", accountUpdate.label);
+        this.approve(accountUpdate);
+      });
+      counter = counter.add(Provable.if(isDummy, Field(0), Field(1)));
+
+      // Track new accounts to be able to know how much new accounts to fund
+      const newAccounts = message.result.accountUpdates
+        .map((accountUpdate) => {
+          const isNew = accountUpdate.account.isNew.getAndRequireEquals();
+          return Provable.if(isNew, Field(1), Field(0));
+        })
+        .reduce((a, b) => a.add(b));
+      accountCreationFeePaid = accountCreationFeePaid.add(newAccounts);
+    }
+
+    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) {
+    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,49 @@
+import { injectable, injectAll } from "tsyringe";
+import { CompileRegistry } from "@proto-kit/common";
+
+import { ContractModule } from "../ContractModule";
+import { OutgoingMessageProcessor } from "../modularity/OutgoingMessageProcessor";
+import { ContractArgsRegistry } from "../ContractArgsRegistry";
+
+import {
+  BridgeContract,
+  BridgeContractArgs,
+  BridgeContractType,
+} from "./BridgeContract";
+
+export type BridgeContractConfig = {
+  outgoingBatchSize?: number;
+};
+
+@injectable()
+export class BridgeContractProtocolModule extends ContractModule<
+  BridgeContractType,
+  BridgeContractConfig
+> {
+  public constructor(
+    @injectAll("OutgoingMessageProcessor", { isOptional: true })
+    private readonly messageProcessors: OutgoingMessageProcessor<unknown>[],
+    private readonly contractArgsRegistry: ContractArgsRegistry
+  ) {
+    super();
+  }
+
+  public contractFactory() {
+    const { config } = this;
+
+    this.contractArgsRegistry.addArgs<BridgeContractArgs>("BridgeContract", {
+      messageProcessors: this.messageProcessors,
+      batchSize: config.outgoingBatchSize,
+    });
+
+    return BridgeContract;
+  }
+
+  public async compile(registry: CompileRegistry) {
+    return {
+      BridgeContract: await registry.forceProverExists(
+        async () => await registry.compile(BridgeContract)
+      ),
+    };
+  }
+}