From c2cb067333a1a2fa4cc3cf082fe2d41f7eed11ef Mon Sep 17 00:00:00 2001
From: tyshkor <tyshko1@gmail.com>
Date: Tue, 28 Mar 2023 21:21:34 +0200
Subject: [PATCH] most changes

---
 .gitignore          |  1 +
 rln-wasm/src/lib.rs | 87 +++++++++++++++++++++++++++++++--------------
 2 files changed, 62 insertions(+), 26 deletions(-)

diff --git a/.gitignore b/.gitignore
index c51faaa..201b119 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,6 +8,7 @@ rln/pmtree_db
 # will have compiled files and executables
 debug/
 target/
+wabt/
 
 # These are backup files generated by rustfmt
 **/*.rs.bk
diff --git a/rln-wasm/src/lib.rs b/rln-wasm/src/lib.rs
index 27c8e25..ffbf81d 100644
--- a/rln-wasm/src/lib.rs
+++ b/rln-wasm/src/lib.rs
@@ -20,6 +20,64 @@ pub struct RLNWrapper {
     instance: RLN<'static>,
 }
 
+// Macro to call_with_error_msg methods with arbitrary amount of arguments,
+// First argument to the macro is context,
+// second is the actual method on `RLNWrapper`
+// rest are all other arguments to the method
+macro_rules! call_with_error_msg {
+    ($instance:expr, $method:ident, $error_msg:expr $(, $arg:expr)*) => {
+        {
+            let new_instance: &mut RLNWrapper = $instance.process();
+            if new_instance.instance.$method($($arg.process()),*).is_ok() {
+                Ok(())
+            } else {
+                Err($error_msg)
+            }
+
+        }
+    }
+}
+
+trait ProcessArg {
+    type ReturnType;
+    fn process(self) -> Self::ReturnType;
+}
+
+impl ProcessArg for usize {
+    type ReturnType = usize;
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
+impl<'a> ProcessArg for *const RLN<'a> {
+    type ReturnType = &'a RLN<'a>;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &*self }
+    }
+}
+
+trait ProcessArgRef {
+    type ReturnType;
+    fn process(self) -> Self::ReturnType;
+}
+
+impl ProcessArgRef for *mut RLNWrapper {
+    type ReturnType = &'static mut RLNWrapper;
+
+    fn process(self) -> Self::ReturnType {
+        unsafe { &mut *self }
+    }
+}
+
+impl<'a> ProcessArgRef for &'a [u8] {
+    type ReturnType = &'a [u8];
+
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = newRLN)]
 pub fn wasm_new(
@@ -51,12 +109,7 @@ pub fn wasm_get_serialized_rln_witness(
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = insertMember)]
 pub fn wasm_set_next_leaf(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper.instance.set_next_leaf(&input.to_vec()[..]).is_ok() {
-        Ok(())
-    } else {
-        Err("could not insert member into merkle tree".into())
-    }
+    call_with_error_msg!(ctx, set_next_leaf, "could not insert member into merkle tree".to_string(), &input.to_vec()[..])
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -66,31 +119,13 @@ pub fn wasm_set_leaves_from(
     index: usize,
     input: Uint8Array,
 ) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper
-        .instance
-        .set_leaves_from(index as usize, &input.to_vec()[..])
-        .is_ok()
-    {
-        Ok(())
-    } else {
-        Err("could not set multiple leaves".into())
-    }
+    call_with_error_msg!(ctx, set_leaves_from, "could not set multiple leaves".to_string(), index, &*input.to_vec())
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = initTreeWithLeaves)]
 pub fn wasm_init_tree_with_leaves(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper
-        .instance
-        .init_tree_with_leaves(&input.to_vec()[..])
-        .is_ok()
-    {
-        Ok(())
-    } else {
-        Err("could not init merkle tree".into())
-    }
+    call_with_error_msg!(ctx, init_tree_with_leaves, "could not init merkle tree".to_string(), &*input.to_vec())
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]