maximally split everything

packages/contracts/sdk/LobService.ts

@@ -1,8 +1,9 @@
-import { type AsyncOrSync } from "ts-essentials";
+import { uniq } from "lodash";
+import { assert, type AsyncOrSync } from "ts-essentials";
 import { type PoolERC20 } from "../typechain-types";
 import { NoteInputStruct } from "../typechain-types/contracts/PoolERC20";
-import { MpcProverService } from "./mpc/MpcNetworkService.js";
-import { splitInput } from "./mpc/utils";
+import { MpcProverService, type Side } from "./mpc/MpcNetworkService.js";
+import { splitInput, splitInput2 } from "./mpc/utils.js";
 import { type ITreesService } from "./RemoteTreesService.js";
 import {
   CompleteWaAddress,
@@ -125,4 +126,108 @@ export class LobService {
     const receipt = await tx.wait();
     console.log("swap gas used", receipt?.gasUsed);
   }
+
+  async requestSwap(params: {
+    secretKey: string;
+    note: Erc20Note;
+    sellAmount: TokenAmount;
+    buyAmount: TokenAmount;
+  }) {
+    const swapCircuit = (await this.circuits).swap;
+    const randomness = await getRandomness();
+
+    const changeNote = await Erc20Note.from({
+      owner: await CompleteWaAddress.fromSecretKey(params.secretKey),
+      amount: params.note.amount.sub(params.sellAmount),
+      randomness,
+    });
+    const swapNote = await Erc20Note.from({
+      owner: await CompleteWaAddress.fromSecretKey(params.secretKey),
+      amount: params.buyAmount,
+      randomness,
+    });
+
+    const order = {
+      sell_amount: await params.sellAmount.toNoir(),
+      buy_amount: await params.buyAmount.toNoir(),
+      randomness,
+    };
+
+    // deterministic side
+    const side: Side =
+      params.sellAmount.token.toLowerCase() <
+      params.buyAmount.token.toLowerCase()
+        ? "seller"
+        : "buyer";
+    const input = {
+      [`${side}_secret_key`]: params.secretKey,
+      [`${side}_note`]: await this.poolErc20.toNoteConsumptionInputs(
+        params.secretKey,
+        params.note,
+      ),
+      [`${side}_order`]: order,
+      [`${side}_randomness`]: randomness,
+    };
+    console.log("side", side, randomness);
+    // only one trading party need to provide public inputs
+    const inputPublic =
+      side === "seller"
+        ? {
+            tree_roots: await this.trees.getTreeRoots(),
+          }
+        : undefined;
+    const inputsShared = await splitInput2(swapCircuit.circuit, {
+      // merge public inputs into first input because it does not matter how public inputs are passed
+      ...input,
+      ...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,
+        });
+      }),
+    );
+    console.log("got proofs", proofs.length);
+    assert(uniq(proofs).length === 1, "proofs mismatch");
+    const proof = proofs[0]!;
+    return {
+      proof,
+      side,
+      changeNote: await changeNote.toSolidityNoteInput(),
+      swapNote: await swapNote.toSolidityNoteInput(),
+      nullifier: (
+        await params.note.computeNullifier(params.secretKey)
+      ).toString(),
+    };
+  }
+
+  async commitSwap(sellerSwap: SwapResult, buyerSwap: SwapResult) {
+    assert(
+      sellerSwap.proof === buyerSwap.proof,
+      "seller & buyer proof mismatch",
+    );
+    const proof = sellerSwap.proof;
+
+    const tx = await this.contract.swap(
+      proof,
+      [
+        sellerSwap.changeNote,
+        buyerSwap.swapNote,
+        buyerSwap.changeNote,
+        sellerSwap.swapNote,
+      ],
+      [sellerSwap.nullifier, buyerSwap.nullifier],
+    );
+    const receipt = await tx.wait();
+    console.log("swap gas used", receipt?.gasUsed);
+  }
 }
+
+type SwapResult = Awaited<ReturnType<LobService["requestSwap"]>>;

packages/contracts/sdk/mpc/MpcNetworkService.ts

@@ -1,13 +1,112 @@
 import type { CompiledCircuit } from "@noir-lang/noir_js";
 import { ethers } from "ethers";
-import { range } from "lodash";
+import { omit, range } from "lodash";
 import fs from "node:fs";
 import path from "node:path";
 import { assert } from "ts-essentials";
 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;
+export type Side = "seller" | "buyer";
+
+type Order = {
+  side: Side;
+  id: OrderId;
+  inputShared: string;
+  result: ReturnType<typeof promiseWithResolvers<string>>;
+};
+
 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(),
+  };
+  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)) {
+      throw new Error(`order already exists ${params.orderId}`);
+    }
+    const order: Order = {
+      id: params.orderId,
+      inputShared: params.inputShared,
+      side: params.side,
+      result: promiseWithResolvers(),
+    };
+    this.#storage[params.partyIndex].set(params.orderId, order);
+
+    this.#tryExecuteOrder(params.orderId, {
+      partyIndex: params.partyIndex,
+      circuit: params.circuit,
+      numPublicInputs: params.numPublicInputs,
+    });
+
+    return await order.result.promise;
+  }
+
+  async #tryExecuteOrder(
+    orderId: OrderId,
+
+    params: {
+      partyIndex: PartyIndex;
+      circuit: CompiledCircuit;
+      numPublicInputs: number;
+    },
+  ) {
+    const order = this.#storage[params.partyIndex].get(orderId);
+    if (!order) {
+      throw new Error(
+        `order not found in party storage ${params.partyIndex}: ${orderId}`,
+      );
+    }
+
+    const otherOrders = Array.from(
+      this.#storage[params.partyIndex].values(),
+    ).filter((o) => o.id !== order.id && o.side !== order.side);
+    if (otherOrders.length === 0) {
+      return;
+    }
+    const otherOrder = otherOrders[0]!;
+    const inputsShared =
+      order.side === "seller"
+        ? ([order.inputShared, otherOrder.inputShared] as const)
+        : ([otherOrder.inputShared, order.inputShared] as const);
+    console.log(
+      "executing orders",
+      params.partyIndex,
+      omit(order, "inputShared"),
+      omit(otherOrder, "inputShared"),
+    );
+    try {
+      const { proof } = await this.proveAsParty({
+        partyIndex: params.partyIndex,
+        circuit: params.circuit,
+        input0Shared: inputsShared[0],
+        input1Shared: inputsShared[1],
+        numPublicInputs: params.numPublicInputs,
+      });
+      console.log("got proof", orderId, params.partyIndex, proof.length);
+      const proofHex = ethers.hexlify(proof);
+      order.result.resolve(proofHex);
+      otherOrder.result.resolve(proofHex);
+    } catch (error) {
+      order.result.reject(error);
+      otherOrder.result.reject(error);
+    }
+  }
+
   async proveAsParty(params: {
     partyIndex: number;
     circuit: CompiledCircuit;
@@ -16,6 +115,7 @@ export class MpcProverService {
     // TODO: infer number of public inputs
     numPublicInputs: number;
   }) {
+    console.log("proving as party", params.partyIndex);
     return await inWorkingDir(async (workingDir) => {
       for (const [traderIndex, inputShared] of [
         params.input0Shared,

packages/contracts/sdk/mpc/utils.ts

@@ -4,7 +4,11 @@ import { range } from "lodash";
 import fs from "node:fs";
 import path from "node:path";
 import toml from "smol-toml";
+import type { PartyIndex } from "./MpcNetworkService.js";
 
+/**
+ * @deprecated use {@link splitInput2} instead
+ */
 export async function splitInput(circuit: CompiledCircuit, input: InputMap) {
   return await inWorkingDir(async (workingDir) => {
     const proverPath = path.join(workingDir, "ProverX.toml");
@@ -21,6 +25,14 @@ export async function splitInput(circuit: CompiledCircuit, input: InputMap) {
   });
 }
 
+export async function splitInput2(circuit: CompiledCircuit, input: InputMap) {
+  const shared = await splitInput(circuit, input);
+  return Array.from(shared.entries()).map(([partyIndex, inputShared]) => ({
+    partyIndex: partyIndex as PartyIndex,
+    inputShared,
+  }));
+}
+
 export async function inWorkingDir<T>(f: (workingDir: string) => Promise<T>) {
   const id = crypto.randomUUID();
   const workingDir = path.join(__dirname, "work-dirs", id);

packages/contracts/sdk/utils.ts

@@ -78,3 +78,16 @@ export async function prove(
   proof = proof.slice(4); // remove length
   return { proof, witness, returnValue, publicInputs };
 }
+
+export function promiseWithResolvers<T>(): {
+  promise: Promise<T>;
+  resolve: (value: T) => void;
+  reject: (reason: unknown) => void;
+} {
+  const ret: any = {};
+  ret.promise = new Promise((resolve, reject) => {
+    ret.resolve = resolve;
+    ret.reject = reject;
+  });
+  return ret;
+}

packages/contracts/test/PoolERC20.test.ts

@@ -408,4 +408,59 @@ describe("PoolERC20", () => {
     expect(await sdk.poolErc20.balanceOf(usdc, bobSecretKey)).to.equal(70n);
     expect(await sdk.poolErc20.balanceOf(btc, bobSecretKey)).to.equal(8n);
   });
+
+  it("swaps mpc", async () => {
+    const { note: aliceNote } = await sdk.poolErc20.shield({
+      account: alice,
+      token: usdc,
+      amount: 100n,
+      secretKey: aliceSecretKey,
+    });
+    const { note: bobNote } = await sdk.poolErc20.shield({
+      account: bob,
+      token: btc,
+      amount: 10n,
+      secretKey: bobSecretKey,
+    });
+
+    await backendSdk.rollup.rollup();
+
+    const sellerAmount = await TokenAmount.from({
+      token: await usdc.getAddress(),
+      amount: 70n,
+    });
+    const buyerAmount = await TokenAmount.from({
+      token: await btc.getAddress(),
+      amount: 2n,
+    });
+
+    const swapAlicePromise = sdk.lob.requestSwap({
+      secretKey: aliceSecretKey,
+      note: aliceNote,
+      sellAmount: sellerAmount,
+      buyAmount: buyerAmount,
+    });
+    const swapBobPromise = sdk.lob.requestSwap({
+      secretKey: bobSecretKey,
+      note: bobNote,
+      sellAmount: buyerAmount,
+      buyAmount: sellerAmount,
+    });
+    const [swapAlice, swapBob] = await Promise.all([
+      swapAlicePromise,
+      swapBobPromise,
+    ]);
+    const args =
+      swapAlice.side === "seller"
+        ? ([swapAlice, swapBob] as const)
+        : ([swapBob, swapAlice] as const);
+    await sdk.lob.commitSwap(...args);
+
+    await backendSdk.rollup.rollup();
+
+    expect(await sdk.poolErc20.balanceOf(usdc, aliceSecretKey)).to.equal(30n);
+    expect(await sdk.poolErc20.balanceOf(btc, aliceSecretKey)).to.equal(2n);
+    expect(await sdk.poolErc20.balanceOf(usdc, bobSecretKey)).to.equal(70n);
+    expect(await sdk.poolErc20.balanceOf(btc, bobSecretKey)).to.equal(8n);
+  });
 });