Add a new interface to remove keys from proof (#127)

* 🌱 Add a new interface to remove keys from proof

* ♻️ Fix some bugs

 - Also add a new unit test

* ♻️ Introduce the prepare_result_proof function

* ♻️ Introduce is_bitmap_valid for SMT

 - Also remove redundant code from calculate_root

* ♻️ Introduce function next_query in SMT

 - Also remove redundant line form calculate_root function

* ♻️ Remove BTreeMap and use Vec instead

* ♻️ Remove clone and use iter in loops

 - Also fix a clippy bug

* ♻️ refactor to address some comments

* ♻️ Edit for some naming suggestion

* ♻️ Edit to address some comments

sparse_merkle_tree.js

@@ -18,6 +18,7 @@ const {
     in_memory_smt_prove,
     in_memory_smt_verify,
     in_memory_smt_calculate_root,
+    in_memory_smt_remove_keys_from_proof,
 } = require("./bin-package/index.node");
 const { isInclusionProofForQueryKey } = require('./utils');
 
@@ -108,6 +109,18 @@ class SparseMerkleTree {
             });
         });
     }
+
+    async removeKeysFromProof(proof, removedKeys) {
+        return new Promise((resolve, reject) => {
+            in_memory_smt_remove_keys_from_proof.call(null, proof, removedKeys, (err, result) => {
+                if (err) {
+                    reject(err);
+                    return;
+                }
+                resolve(result);
+            });
+        });
+    }
 }
 
 module.exports = {

src/lib.rs

@@ -94,11 +94,13 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
 
     let in_memory_smt_new = InMemorySMT::js_new_with_arc_mutex::<InMemorySMT>;
     let in_memory_smt_calculate_root = InMemorySMT::js_calculate_root;
+    let remove_keys_proof = InMemorySMT::js_remove_keys_from_proof;
     cx.export_function("in_memory_smt_new", in_memory_smt_new)?;
     cx.export_function("in_memory_smt_update", InMemorySMT::js_update)?;
     cx.export_function("in_memory_smt_prove", InMemorySMT::js_prove)?;
     cx.export_function("in_memory_smt_verify", InMemorySMT::js_verify)?;
     cx.export_function("in_memory_smt_calculate_root", in_memory_smt_calculate_root)?;
+    cx.export_function("in_memory_smt_remove_keys_from_proof", remove_keys_proof)?;
 
     Ok(())
 }

src/sparse_merkle_tree/in_memory_smt.rs

@@ -16,6 +16,7 @@ use crate::types::{ArcMutex, Cache, KVPair, KeyLength, NestedVec};
 type SharedInMemorySMT = JsArcMutex<InMemorySMT>;
 type DatabaseParameters = (ArcMutex<InMemorySMT>, Vec<u8>, Root<JsFunction>);
 type VerifyParameters = (Vec<u8>, NestedVec, Proof, KeyLength, Root<JsFunction>);
+type RemovedKeysParameters = (Proof, NestedVec, Root<JsFunction>);
 
 struct JsFunctionContext<'a> {
     context: FunctionContext<'a>,
@@ -229,6 +230,30 @@ impl JsFunctionContext<'_> {
         Ok(proof)
     }
 
+    fn get_removed_keys_parameters(&mut self) -> NeonResult<RemovedKeysParameters> {
+        let proof = self.get_proof(0)?;
+
+        let removed_keys = self
+            .context
+            .argument::<JsArray>(1)?
+            .to_vec(&mut self.context)?;
+        let mut parsed_removed_keys = NestedVec::new();
+        for key in removed_keys.iter() {
+            let key = key
+                .downcast_or_throw::<JsTypedArray<u8>, _>(&mut self.context)?
+                .as_slice(&self.context)
+                .to_vec();
+            parsed_removed_keys.push(key);
+        }
+
+        let callback = self
+            .context
+            .argument::<JsFunction>(2)?
+            .root(&mut self.context);
+
+        Ok((proof, parsed_removed_keys, callback))
+    }
+
     fn get_verify_parameters(&mut self) -> NeonResult<VerifyParameters> {
         let state_root = self
             .context
@@ -367,4 +392,56 @@ impl InMemorySMT {
 
         Ok(js_context.context.undefined())
     }
+
+    pub fn js_remove_keys_from_proof(ctx: FunctionContext) -> JsResult<JsUndefined> {
+        let mut js_context = JsFunctionContext { context: ctx };
+
+        let (proof, parsed_removed_keys, callback) = js_context.get_removed_keys_parameters()?;
+        let channel = js_context.context.channel();
+        thread::spawn(move || {
+            let result = SparseMerkleTree::remove_keys_from_proof(
+                &proof,
+                &parsed_removed_keys
+                    .iter()
+                    .map(|x| x.as_slice())
+                    .collect::<Vec<_>>(),
+            );
+
+            channel.send(move |mut ctx| {
+                let callback = callback.into_inner(&mut ctx);
+                let this = ctx.undefined();
+                let args: Vec<Handle<JsValue>> = match result {
+                    Ok(val) => {
+                        let obj: Handle<JsObject> = ctx.empty_object();
+                        let sibling_hashes = ctx.empty_array();
+                        for (i, h) in val.sibling_hashes.iter().enumerate() {
+                            let val_res = JsBuffer::external(&mut ctx, h.to_vec());
+                            sibling_hashes.set(&mut ctx, i as u32, val_res)?;
+                        }
+                        obj.set(&mut ctx, "siblingHashes", sibling_hashes)?;
+                        let queries = ctx.empty_array();
+                        obj.set(&mut ctx, "queries", queries)?;
+                        for (i, v) in val.queries.iter().enumerate() {
+                            let obj = ctx.empty_object();
+                            let key = JsBuffer::external(&mut ctx, v.key_as_vec());
+                            obj.set(&mut ctx, "key", key)?;
+                            let value = JsBuffer::external(&mut ctx, v.value_as_vec());
+                            obj.set(&mut ctx, "value", value)?;
+                            let bitmap = JsBuffer::external(&mut ctx, v.bitmap.to_vec());
+                            obj.set(&mut ctx, "bitmap", bitmap)?;
+
+                            queries.set(&mut ctx, i as u32, obj)?;
+                        }
+                        vec![ctx.null().upcast(), obj.upcast()]
+                    },
+                    Err(err) => vec![ctx.error(err.to_string())?.upcast()],
+                };
+                callback.call(&mut ctx, this, args)?;
+
+                Ok(())
+            })
+        });
+
+        Ok(js_context.context.undefined())
+    }
 }