gausslabs · RiverRuby · Feb 23, 2024 · Feb 23, 2024 · Feb 24, 2024 · Feb 24, 2024
diff --git a/.gitignore b/.gitignore
@@ -1 +1 @@
-/target
+target
diff --git a/README.md b/README.md
@@ -3,13 +3,12 @@
 The library contains the following components:
 
 - `src`: Rust library for multi-party PSI using BFV
-- `pkg`: JS-TS-WASM package 
+- `pkg`: JS-TS-WASM package
 
-### Build 
+### Build
 
 The rust library is used to build the JS-TS-WASM package using `wasm-pack` targeting `web` [guide](https://developer.mozilla.org/en-US/docs/WebAssembly/Rust_to_Wasm). When compiling to `web` the output can natively be included on a web page, and doesn't require any further postprocessing. The output is included as an ES module. For more information check [`wasm-bindgen` guide](https://rustwasm.github.io/docs/wasm-bindgen/reference/deployment.html)
 
-
 ```bash
 wasm-pack build --target web
 ```
@@ -20,4 +19,4 @@ To test the rust library, run:
 
 ```bash
 cargo test --release
-```
+```
diff --git a/pkg/LICENSE b/pkg/LICENSE
diff --git a/pkg/README.md b/pkg/README.md
diff --git a/pkg/index.html b/pkg/index.html
@@ -1,69 +1,99 @@
-
-<!doctype html>
+<!DOCTYPE html>
 <html lang="en-US">
   <head>
     <meta charset="utf-8" />
     <title>mp-psi example</title>
   </head>
   <body>
     <script type="module">
-      import init, { state0_bindgen, state1_bindgen, state2_bindgen, state3_bindgen, state4_bindgen } from "./mp_psi.js";
+      import init, {
+        gen_keys_js,
+        round1_js,
+        round2_js,
+        round3_js,
+      } from "./mp_psi.js";
 
       function randomBitVector(hammingWeight, size) {
         let bitVector = new Array(size).fill(0);
 
         for (let i = 0; i < hammingWeight; i++) {
-            let sampleIndex;
-            do {
-                sampleIndex = Math.floor(Math.random() * size);
-            } while (bitVector[sampleIndex] === 1);
+          let sampleIndex;
+          do {
+            sampleIndex = Math.floor(Math.random() * size);
+          } while (bitVector[sampleIndex] === 1);
 
-            bitVector[sampleIndex] = 1;
+          bitVector[sampleIndex] = 1;
         }
         return bitVector;
       }
 
       function plainPsi(bitVector0, bitVector1) {
         if (bitVector0.length !== bitVector1.length) {
-            throw new Error("Both bit vectors must be of the same length");
+          throw new Error("Both bit vectors must be of the same length");
         }
 
         return bitVector0.map((element, index) => element * bitVector1[index]);
       }
 
       init().then(() => {
-        console.time('state0 Time');
-        const state0 = state0_bindgen();
-        console.timeEnd('state0 Time');
-
+        console.log("each time includes both parties");
+
+        console.time("gen_keys time");
+        const gen_keys_output_a = gen_keys_js();
+        const gen_keys_output_b = gen_keys_js();
+        console.timeEnd("gen_keys time");
+
+        console.log(gen_keys_output_a);
+
         const hammingWeight = 1000;
         const size = 2048 * 3;
-        const bit_vector_b = randomBitVector(hammingWeight, size);
-
-        console.time('state1 Time');
-        const state1 = state1_bindgen(state0.message_a_to_b, bit_vector_b);
-        console.timeEnd('state1 Time');
 
         const bit_vector_a = randomBitVector(hammingWeight, size);
+        const bit_vector_b = randomBitVector(hammingWeight, size);
+
+        console.time("round 1 time");
+        const round1_output_a = round1_js(
+          gen_keys_output_a,
+          gen_keys_output_b.message_round1,
+          bit_vector_a
+        );
+        const round1_output_b = round1_js(
+          gen_keys_output_b,
+          gen_keys_output_a.message_round1,
+          bit_vector_b
+        );
+        console.timeEnd("round 1 time");
 
-        console.time('state2 Time');
-        const state2 = state2_bindgen(state0.private_output_a, state0.public_output_a, state1.message_b_to_a, bit_vector_a);
-        console.timeEnd('state2 Time');
+        console.time("round 2 time");
+        const round2_output_a = round2_js(
+          gen_keys_output_a,
+          round1_output_a,
+          round1_output_b.message_round2,
+          true
+        );
+        const round2_output_b = round2_js(
+          gen_keys_output_b,
+          round1_output_b,
+          round1_output_a.message_round2,
+          false
+        );
+        console.timeEnd("round 2 time");
 
-        console.time('state3 Time');
-        const state3 = state3_bindgen(state1.private_output_b, state1.public_output_b, state2.message_a_to_b);
-        console.timeEnd('state3 Time');       
-
-        console.time('state4 Time');
-        const psi_output_a = state4_bindgen(state2.public_output_a, state3.message_b_to_a);
-        console.timeEnd('state4 Time');
+        console.time("round 3 time");
+        const psi_output_a = round3_js(
+          round2_output_a,
+          round2_output_b.message_round3
+        );
+        const psi_output_b = round3_js(
+          round2_output_b,
+          round2_output_a.message_round3
+        );
+        console.timeEnd("round 3 time");
 
-        const psi_output_b = state3.psi_output;
-        console.log("psi_output_a", psi_output_a)
-        console.log("psi_output_b", psi_output_b)
+        console.log("psi_output_a", psi_output_a);
+        console.log("psi_output_b", psi_output_b);
         const expected_psi_output = plainPsi(bit_vector_a, bit_vector_b);
-        console.log("expected_psi_output", expected_psi_output)
-        console.log("psi_output", psi_output_a)
+        console.log("expected_psi_output", expected_psi_output);
       });
     </script>
   </body>

diff --git a/pkg/mp_psi.d.ts b/pkg/mp_psi.d.ts
@@ -3,44 +3,37 @@
 /**
 * @returns {any}
 */
-export function state0_bindgen(): any;
+export function gen_keys_js(): any;
 /**
-* @param {any} message_a_to_b
+* @param {any} gen_keys_output
+* @param {any} other_message_round1
 * @param {Uint32Array} bit_vector
 * @returns {any}
 */
-export function state1_bindgen(message_a_to_b: any, bit_vector: Uint32Array): any;
+export function round1_js(gen_keys_output: any, other_message_round1: any, bit_vector: Uint32Array): any;
 /**
-* @param {any} private_output_a_state0
-* @param {any} public_output_a_state0
-* @param {any} message_b_to_a
-* @param {Uint32Array} bit_vector
-* @returns {any}
-*/
-export function state2_bindgen(private_output_a_state0: any, public_output_a_state0: any, message_b_to_a: any, bit_vector: Uint32Array): any;
-/**
-* @param {any} private_output_b_state1
-* @param {any} public_output_b_state1
-* @param {any} message_a_to_b
+* @param {any} gen_keys_output
+* @param {any} round1_output
+* @param {any} other_message_round2
+* @param {boolean} is_a
 * @returns {any}
 */
-export function state3_bindgen(private_output_b_state1: any, public_output_b_state1: any, message_a_to_b: any): any;
+export function round2_js(gen_keys_output: any, round1_output: any, other_message_round2: any, is_a: boolean): any;
 /**
-* @param {any} public_output_a_state2
-* @param {any} message_b_to_a
+* @param {any} round2_output
+* @param {any} other_message
 * @returns {Uint32Array}
 */
-export function state4_bindgen(public_output_a_state2: any, message_b_to_a: any): Uint32Array;
+export function round3_js(round2_output: any, other_message: any): Uint32Array;
 
 export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembly.Module;
 
 export interface InitOutput {
   readonly memory: WebAssembly.Memory;
-  readonly state0_bindgen: () => number;
-  readonly state1_bindgen: (a: number, b: number, c: number) => number;
-  readonly state2_bindgen: (a: number, b: number, c: number, d: number, e: number) => number;
-  readonly state3_bindgen: (a: number, b: number, c: number) => number;
-  readonly state4_bindgen: (a: number, b: number, c: number) => void;
+  readonly gen_keys_js: () => number;
+  readonly round1_js: (a: number, b: number, c: number, d: number) => number;
+  readonly round2_js: (a: number, b: number, c: number, d: number) => number;
+  readonly round3_js: (a: number, b: number, c: number) => void;
   readonly __wbindgen_malloc: (a: number, b: number) => number;
   readonly __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
   readonly __wbindgen_add_to_stack_pointer: (a: number) => number;