feat: Enforce WitnessMap type in TS (#9)

chore: implement `From` between JS and Rust witness maps

src/abi.rs

@@ -1,3 +1,4 @@
+use acvm::{acir::native_types::Witness, FieldElement};
 use iter_extended::{btree_map, try_btree_map};
 use noirc_abi::{errors::InputParserError, input_parser::InputValue, Abi, MAIN_RETURN_NAME};
 use serde::Serialize;
@@ -8,7 +9,7 @@ use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
 
 mod temp;
 
-use crate::js_transforms::{js_map_to_witness_map, witness_map_to_js_map};
+use crate::JsWitnessMap;
 
 use self::temp::{input_value_from_json_type, JsonTypes};
 
@@ -17,7 +18,7 @@ pub fn abi_encode(
     abi: JsValue,
     inputs: JsValue,
     return_value: JsValue,
-) -> Result<js_sys::Map, JsValue> {
+) -> Result<JsWitnessMap, JsValue> {
     console_error_panic_hook::set_once();
     let abi: Abi = JsValueSerdeExt::into_serde(&abi).map_err(|err| err.to_string())?;
     let inputs: BTreeMap<String, JsonTypes> =
@@ -52,15 +53,15 @@ pub fn abi_encode(
 
     let witness_map = abi.encode(&parsed_inputs, return_value).map_err(|err| err.to_string())?;
 
-    Ok(witness_map_to_js_map(witness_map))
+    Ok(witness_map.into())
 }
 
 #[wasm_bindgen(js_name = abiDecode)]
-pub fn abi_decode(abi: JsValue, witness_map: js_sys::Map) -> Result<JsValue, JsValue> {
+pub fn abi_decode(abi: JsValue, witness_map: JsWitnessMap) -> Result<JsValue, JsValue> {
     console_error_panic_hook::set_once();
     let abi: Abi = JsValueSerdeExt::into_serde(&abi).map_err(|err| err.to_string())?;
 
-    let witness_map = js_map_to_witness_map(witness_map);
+    let witness_map: BTreeMap<Witness, FieldElement> = witness_map.into();
 
     let (inputs, return_value) = abi.decode(&witness_map).map_err(|err| err.to_string())?;
 

src/execute.rs

@@ -15,9 +15,9 @@ use std::collections::BTreeMap;
 
 use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
 
-use crate::js_transforms::{
-    field_element_to_js_string, js_map_to_witness_map, js_value_to_field_element,
-    witness_map_to_js_map,
+use crate::{
+    js_transforms::{field_element_to_js_string, js_value_to_field_element},
+    JsWitnessMap,
 };
 
 struct SimulatedBackend;
@@ -56,12 +56,12 @@ impl PartialWitnessGenerator for SimulatedBackend {
 #[wasm_bindgen(js_name = executeCircuit)]
 pub async fn execute_circuit(
     circuit: Vec<u8>,
-    initial_witness: js_sys::Map,
+    initial_witness: JsWitnessMap,
     oracle_resolver: js_sys::Function,
-) -> Result<js_sys::Map, JsValue> {
+) -> Result<JsWitnessMap, JsValue> {
     console_error_panic_hook::set_once();
     let circuit: Circuit = Circuit::read(&*circuit).expect("Failed to deserialize circuit");
-    let mut witness_map = js_map_to_witness_map(initial_witness);
+    let mut witness_map: BTreeMap<Witness, FieldElement> = initial_witness.into();
 
     let mut blocks = Blocks::default();
     let mut opcodes = circuit.opcodes;
@@ -102,7 +102,7 @@ pub async fn execute_circuit(
         }
     }
 
-    Ok(witness_map_to_js_map(witness_map))
+    Ok(witness_map.into())
 }
 
 fn insert_value(

src/js_transforms.rs

@@ -3,6 +3,33 @@ use js_sys::JsString;
 use std::collections::BTreeMap;
 use wasm_bindgen::JsValue;
 
+use crate::JsWitnessMap;
+
+impl From<BTreeMap<Witness, FieldElement>> for JsWitnessMap {
+    fn from(witness_map: BTreeMap<Witness, FieldElement>) -> Self {
+        let js_map = JsWitnessMap::new();
+        for (key, value) in witness_map {
+            js_map.set(
+                &js_sys::Number::from(key.witness_index()),
+                &field_element_to_js_string(&value),
+            );
+        }
+        js_map
+    }
+}
+
+impl From<JsWitnessMap> for BTreeMap<Witness, FieldElement> {
+    fn from(js_map: JsWitnessMap) -> Self {
+        let mut witness_map: BTreeMap<Witness, FieldElement> = BTreeMap::new();
+        js_map.for_each(&mut |value, key| {
+            let witness_index = Witness(key.as_f64().unwrap() as u32);
+            let witness_value = js_value_to_field_element(value).unwrap();
+            witness_map.insert(witness_index, witness_value);
+        });
+        witness_map
+    }
+}
+
 pub(crate) fn js_value_to_field_element(js_value: JsValue) -> Result<FieldElement, JsString> {
     let hex_str = js_value.as_string().ok_or("failed to parse field element from non-string")?;
 
@@ -18,24 +45,6 @@ pub(crate) fn field_element_to_js_string(field_element: &FieldElement) -> JsStri
     format!("0x{}", field_element.to_hex()).into()
 }
 
-pub(crate) fn js_map_to_witness_map(js_map: js_sys::Map) -> BTreeMap<Witness, FieldElement> {
-    let mut witness_map: BTreeMap<Witness, FieldElement> = BTreeMap::new();
-    js_map.for_each(&mut |value, key| {
-        let witness_index = Witness(key.as_f64().unwrap() as u32);
-        let witness_value = js_value_to_field_element(value).unwrap();
-        witness_map.insert(witness_index, witness_value);
-    });
-    witness_map
-}
-
-pub(crate) fn witness_map_to_js_map(witness_map: BTreeMap<Witness, FieldElement>) -> js_sys::Map {
-    let js_map = js_sys::Map::new();
-    for (key, value) in witness_map {
-        js_map.set(&js_sys::Number::from(key.witness_index()), &field_element_to_js_string(&value));
-    }
-    js_map
-}
-
 #[cfg(test)]
 mod test {
     use std::collections::BTreeMap;
@@ -44,7 +53,7 @@ mod test {
     use wasm_bindgen::JsValue;
     use wasm_bindgen_test::*;
 
-    use super::witness_map_to_js_map;
+    use crate::JsWitnessMap;
 
     #[wasm_bindgen_test]
     fn test_witness_map_to_js() {
@@ -54,7 +63,7 @@ mod test {
             (Witness(3), -FieldElement::one()),
         ]);
 
-        let js_map = witness_map_to_js_map(witness_map);
+        let js_map = JsWitnessMap::from(witness_map);
 
         assert_eq!(js_map.get(&JsValue::from("1")), JsValue::from_str("1"));
     }

src/lib.rs

@@ -1,8 +1,8 @@
-#![forbid(unsafe_code)]
 #![warn(unused_crate_dependencies, unused_extern_crates)]
 #![warn(unreachable_pub)]
 
 use gloo_utils::format::JsValueSerdeExt;
+use js_sys::Map;
 use log::Level;
 use serde::{Deserialize, Serialize};
 use std::str::FromStr;
@@ -45,3 +45,85 @@ pub fn build_info() -> JsValue {
     console_error_panic_hook::set_once();
     <JsValue as JsValueSerdeExt>::from_serde(&BUILD_INFO).unwrap()
 }
+
+#[wasm_bindgen(typescript_custom_section)]
+const TS_APPEND_CONTENT: &'static str = r#"
+// Map from witness index to hex string value of witness.
+export type WitnessMap = Map<number, string>;
+
+"#;
+
+// WitnessMap
+#[wasm_bindgen]
+extern "C" {
+    #[wasm_bindgen(extends = Map, js_name = "WitnessMap", typescript_type = "WitnessMap")]
+    #[derive(Clone, Debug, PartialEq, Eq)]
+    pub type JsWitnessMap;
+
+    /// The `clear()` method removes all elements from a Map object.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/clear)
+    #[wasm_bindgen(method, js_class = "Map")]
+    pub fn clear(this: &JsWitnessMap);
+
+    /// The `delete()` method removes the specified element from a Map object.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/delete)
+    #[wasm_bindgen(method, js_class = "Map")]
+    pub fn delete(this: &JsWitnessMap, key: &JsValue) -> bool;
+
+    /// The `forEach()` method executes a provided function once per each
+    /// key/value pair in the Map object, in insertion order.
+    /// Note that in Javascript land the `Key` and `Value` are reversed compared to normal expectations:
+    /// # Examples
+    /// ```
+    /// let js_map = Map::new();
+    /// js_map.for_each(&mut |value, key| {
+    ///     // Do something here...
+    /// })
+    /// ```
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/forEach)
+    #[wasm_bindgen(method, js_class = "Map", js_name = forEach)]
+    pub fn for_each(this: &JsWitnessMap, callback: &mut dyn FnMut(JsValue, JsValue));
+
+    /// The `get()` method returns a specified element from a Map object.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/get)
+    #[wasm_bindgen(method, js_class = "Map")]
+    pub fn get(this: &JsWitnessMap, key: &JsValue) -> JsValue;
+
+    /// The `has()` method returns a boolean indicating whether an element with
+    /// the specified key exists or not.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/has)
+    #[wasm_bindgen(method, js_class = "Map")]
+    pub fn has(this: &JsWitnessMap, key: &JsValue) -> bool;
+
+    /// The Map object holds key-value pairs. Any value (both objects and
+    /// primitive values) maybe used as either a key or a value.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map)
+    #[wasm_bindgen(constructor, js_class = "Map")]
+    pub fn new() -> JsWitnessMap;
+
+    /// The `set()` method adds or updates an element with a specified key
+    /// and value to a Map object.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/set)
+    #[wasm_bindgen(method, js_class = "Map")]
+    pub fn set(this: &JsWitnessMap, key: &JsValue, value: &JsValue) -> Map;
+
+    /// The value of size is an integer representing how many entries
+    /// the Map object has. A set accessor function for size is undefined;
+    /// you can not change this property.
+    ///
+    /// [MDN documentation](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Map/size)
+    #[wasm_bindgen(method, js_class = "Map", getter, structural)]
+    pub fn size(this: &JsWitnessMap) -> u32;
+}
+
+impl Default for JsWitnessMap {
+    fn default() -> Self {
+        Self::new()
+    }
+}

test/browser/abi_encode.test.ts

@@ -1,5 +1,9 @@
-import initACVMSimulator, { abiEncode, abiDecode } from "../../pkg/";
-import { DecodedInputs, WitnessMap } from "../types";
+import initACVMSimulator, {
+  abiEncode,
+  abiDecode,
+  WitnessMap,
+} from "../../pkg/";
+import { DecodedInputs } from "../types";
 
 test("recovers original inputs when abi encoding and decoding", async () => {
   await initACVMSimulator();

test/browser/execute_circuit.test.ts

@@ -2,8 +2,8 @@ import initACVMSimulator, {
   abiEncode,
   abiDecode,
   executeCircuit,
+  WitnessMap,
 } from "../../pkg/";
-import { WitnessMap } from "../types";
 
 test("successfully executes circuit and extracts return value", async () => {
   await initACVMSimulator();

test/node/abi_encode.test.ts

@@ -1,6 +1,6 @@
 import { expect, test } from "@jest/globals";
-import { abiEncode, abiDecode } from "../../pkg/";
-import { DecodedInputs, WitnessMap } from "../types";
+import { abiEncode, abiDecode, WitnessMap } from "../../pkg/";
+import { DecodedInputs } from "../types";
 
 test("recovers original inputs when abi encoding and decoding", () => {
   // TODO use ts-rs to get ABI type bindings.

test/node/execute_circuit.test.ts

@@ -1,6 +1,5 @@
 import { expect, test } from "@jest/globals";
-import { abiEncode, abiDecode, executeCircuit } from "../../pkg/";
-import { WitnessMap } from "../types";
+import { abiEncode, abiDecode, executeCircuit, WitnessMap } from "../../pkg/";
 
 test("successfully executes circuit and extracts return value", async () => {
   // Noir program which enforces that x != y and returns x + y.

test/types.ts

@@ -1,5 +1,2 @@
-// Map from witness index to hex string value of witness.
-export type WitnessMap = Map<number, string>;
-
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export type DecodedInputs = { inputs: Record<string, any>; return_value: any };