simplify

Author: olehmisar

packages/contracts/sdk/LobService.ts

@@ -167,18 +167,12 @@ export class LobService {
       ...inputPublic,
     });
     const orderId = randomness; // TODO: is randomness a good order id?
-    const proofs = await Promise.all(
-      inputsShared.map(({ partyIndex, inputShared }) => {
-        return this.mpcProver.requestProveAsParty({
-          orderId,
-          inputShared,
-          partyIndex,
-          circuit: swapCircuit.circuit,
-          numPublicInputs: 8,
-          side,
-        });
-      }),
-    );
+    const proofs = await this.mpcProver.prove(inputsShared, {
+      orderId,
+      side,
+      circuit: swapCircuit.circuit,
+      numPublicInputs: 8,
+    });
     assert(uniq(proofs).length === 1, "proofs mismatch");
     const proof = proofs[0]!;
     return {

packages/contracts/sdk/mpc/MpcNetworkService.ts

@@ -5,40 +5,51 @@ import fs from "node:fs";
 import path from "node:path";
 import { z } from "zod";
 import { promiseWithResolvers } from "../utils.js";
-import { inWorkingDir, makeRunCommand } from "./utils.js";
-
-export type OrderId = string & { __brand: "OrderId" };
-export type PartyIndex = 0 | 1 | 2;
-/**
- * Deterministically determined based on the tokens being swapped
- */
-export type Side = "seller" | "buyer";
-
-type Order = {
-  side: Side;
-  id: OrderId;
-  inputShared: string;
-  result: ReturnType<typeof promiseWithResolvers<string>>;
-};
+import { inWorkingDir, makeRunCommand, splitInput } from "./utils.js";
 
 export class MpcProverService {
-  // TODO: split this service into per party service to manage storage easier
-  #storage: Record<PartyIndex, Map<OrderId, Order>> = {
-    0: new Map(),
-    1: new Map(),
-    2: new Map(),
+  readonly #parties = {
+    0: new MpcProverPartyService(0),
+    1: new MpcProverPartyService(1),
+    2: new MpcProverPartyService(2),
   };
+
+  async prove(
+    inputsShared: Awaited<ReturnType<typeof splitInput>>,
+    params: {
+      orderId: OrderId;
+      side: Side;
+      circuit: CompiledCircuit;
+      // TODO: infer number of public inputs
+      numPublicInputs: number;
+    },
+  ) {
+    return await Promise.all(
+      inputsShared.map(async ({ partyIndex, inputShared }) => {
+        return await this.#parties[partyIndex].requestProveAsParty({
+          ...params,
+          inputShared,
+        });
+      }),
+    );
+  }
+}
+
+class MpcProverPartyService {
+  #storage: Map<OrderId, Order> = new Map();
+
+  constructor(readonly partyIndex: PartyIndex) {}
+
   async requestProveAsParty(params: {
     orderId: OrderId;
     side: Side;
-    partyIndex: PartyIndex;
     inputShared: string;
     circuit: CompiledCircuit;
     // TODO: infer number of public inputs
     numPublicInputs: number;
   }) {
     // TODO(security): authorization
-    if (this.#storage[params.partyIndex].has(params.orderId)) {
+    if (this.#storage.has(params.orderId)) {
       throw new Error(`order already exists ${params.orderId}`);
     }
     const order: Order = {
@@ -47,10 +58,9 @@ export class MpcProverService {
       side: params.side,
       result: promiseWithResolvers(),
     };
-    this.#storage[params.partyIndex].set(params.orderId, order);
+    this.#storage.set(params.orderId, order);
 
     this.#tryExecuteOrder(params.orderId, {
-      partyIndex: params.partyIndex,
       circuit: params.circuit,
       numPublicInputs: params.numPublicInputs,
     });
@@ -62,21 +72,20 @@ export class MpcProverService {
     orderId: OrderId,
 
     params: {
-      partyIndex: PartyIndex;
       circuit: CompiledCircuit;
       numPublicInputs: number;
     },
   ) {
-    const order = this.#storage[params.partyIndex].get(orderId);
+    const order = this.#storage.get(orderId);
     if (!order) {
       throw new Error(
-        `order not found in party storage ${params.partyIndex}: ${orderId}`,
+        `order not found in party storage ${this.partyIndex}: ${orderId}`,
       );
     }
 
-    const otherOrders = Array.from(
-      this.#storage[params.partyIndex].values(),
-    ).filter((o) => o.id !== order.id && o.side !== order.side);
+    const otherOrders = Array.from(this.#storage.values()).filter(
+      (o) => o.id !== order.id && o.side !== order.side,
+    );
     if (otherOrders.length === 0) {
       return;
     }
@@ -87,13 +96,12 @@ export class MpcProverService {
         : ([otherOrder.inputShared, order.inputShared] as const);
     console.log(
       "executing orders",
-      params.partyIndex,
+      this.partyIndex,
       omit(order, ["inputShared", "result"]),
       omit(otherOrder, ["inputShared", "result"]),
     );
     try {
-      const { proof } = await this.proveAsParty({
-        partyIndex: params.partyIndex,
+      const { proof } = await this.#prove({
         circuit: params.circuit,
         input0Shared: inputsShared[0],
         input1Shared: inputsShared[1],
@@ -108,15 +116,14 @@ export class MpcProverService {
     }
   }
 
-  async proveAsParty(params: {
-    partyIndex: number;
+  async #prove(params: {
     circuit: CompiledCircuit;
     input0Shared: string;
     input1Shared: string;
     // TODO: infer number of public inputs
     numPublicInputs: number;
   }) {
-    console.log("proving as party", params.partyIndex);
+    console.log("proving as party", this.partyIndex);
     return await inWorkingDir(async (workingDir) => {
       for (const [traderIndex, inputShared] of [
         params.input0Shared,
@@ -125,7 +132,7 @@ export class MpcProverService {
         fs.writeFileSync(
           path.join(
             workingDir,
-            `Prover${traderIndex}.toml.${params.partyIndex}.shared`,
+            `Prover${traderIndex}.toml.${this.partyIndex}.shared`,
           ),
           ethers.getBytes(inputShared),
         );
@@ -136,7 +143,7 @@ export class MpcProverService {
 
       const runCommand = makeRunCommand(__dirname);
       await runCommand(
-        `./run-party.sh ${workingDir} ${circuitPath} ${params.partyIndex}`,
+        `./run-party.sh ${workingDir} ${circuitPath} ${this.partyIndex}`,
       );
 
       const publicInputs = z
@@ -152,7 +159,7 @@ export class MpcProverService {
         );
       const proofData = Uint8Array.from(
         fs.readFileSync(
-          path.join(workingDir, `proof.${params.partyIndex}.proof`),
+          path.join(workingDir, `proof.${this.partyIndex}.proof`),
         ),
       );
       // arcane magic
@@ -168,3 +175,17 @@ export class MpcProverService {
     });
   }
 }
+
+export type OrderId = string & { __brand: "OrderId" };
+export type PartyIndex = 0 | 1 | 2;
+/**
+ * Deterministically determined based on the tokens being swapped
+ */
+export type Side = "seller" | "buyer";
+
+type Order = {
+  side: Side;
+  id: OrderId;
+  inputShared: string;
+  result: ReturnType<typeof promiseWithResolvers<string>>;
+};