Repository: ergoplatform/sigma-rust (language: rust)

Issue:
[100 SigmaUSD] AvlTree.remove
Implement `SMethod`(in `types::savltree`) and evaluation (in `eval::savltree`).
Reference Scala implementation - https://github.com/ScorexFoundation/sigmastate-interpreter/blob/c863a9b1a82589e47b15f76f3affdb30a475e740/sigmastate/src/main/scala/sigmastate/types.scala#L2331-L2374



Discussion / hints:
I will work on this

Fix patch:
diff --git a/ergotree-interpreter/src/eval.rs b/ergotree-interpreter/src/eval.rs
index a6a032e53..a1c037c70 100644
--- a/ergotree-interpreter/src/eval.rs
+++ b/ergotree-interpreter/src/eval.rs
@@ -210,6 +210,7 @@ fn smethod_eval_fn(method: &SMethod) -> Result<EvalFn, EvalError> {
             savltree::IS_REMOVE_ALLOWED_METHOD_ID => self::savltree::IS_REMOVE_ALLOWED_EVAL_FN,
             savltree::UPDATE_OPERATIONS_METHOD_ID => self::savltree::UPDATE_OPERATIONS_EVAL_FN,
             savltree::INSERT_METHOD_ID => self::savltree::INSERT_EVAL_FN,
+            savltree::REMOVE_METHOD_ID => self::savltree::REMOVE_EVAL_FN,
             method_id => {
                 return Err(EvalError::NotFound(format!(
                     "Eval fn: unknown method id in SAvlTree: {:?}",
diff --git a/ergotree-interpreter/src/eval/savltree.rs b/ergotree-interpreter/src/eval/savltree.rs
index 1d7518c91..76206bd11 100644
--- a/ergotree-interpreter/src/eval/savltree.rs
+++ b/ergotree-interpreter/src/eval/savltree.rs
@@ -149,6 +149,66 @@ pub(crate) static INSERT_EVAL_FN: EvalFn =
         }
     };
 
+pub(crate) static REMOVE_EVAL_FN: EvalFn =
+    |_env, _ctx, obj, args| {
+        let mut avl_tree_data = obj.try_extract_into::<AvlTreeData>()?;
+
+        if !avl_tree_data.tree_flags.remove_allowed() {
+            return Err(EvalError::AvlTree("Removals not allowed".into()));
+        }
+
+        let keys = {
+            let v = args.get(0).cloned().ok_or_else(|| {
+                EvalError::AvlTree("eval is missing first arg (keys)".to_string())
+            })?;
+            v.try_extract_into::<Vec<Vec<u8>>>()?
+        };
+
+        let proof = {
+            let v = args.get(1).cloned().ok_or_else(|| {
+                EvalError::AvlTree("eval is missing second arg (proof)".to_string())
+            })?;
+            Bytes::from(v.try_extract_into::<Vec<u8>>()?)
+        };
+
+        let starting_digest = Bytes::from(avl_tree_data.digest.0.to_vec());
+        let mut bv = BatchAVLVerifier::new(
+            &starting_digest,
+            &proof,
+            AVLTree::new(
+                |digest| Node::LabelOnly(NodeHeader::new(Some(*digest), None)),
+                avl_tree_data.key_length as usize,
+                avl_tree_data
+                    .value_length_opt
+                    .as_ref()
+                    .map(|v| **v as usize),
+            ),
+            None,
+            None,
+        )
+        .map_err(map_eval_err)?;
+        for key in keys {
+            if bv
+                .perform_one_operation(&Operation::Remove(Bytes::from(key)))
+                .is_err()
+            {
+                return Err(EvalError::AvlTree(format!(
+                    "Incorrect remove for {:?}",
+                    avl_tree_data
+                )));
+            }
+        }
+        if let Some(new_digest) = bv.digest() {
+            let digest = ADDigest::sigma_parse_bytes(&new_digest)?;
+            avl_tree_data.digest = digest;
+            Ok(Value::Opt(Box::new(Some(Value::AvlTree(
+                avl_tree_data.into(),
+            )))))
+        } else {
+            Err(EvalError::AvlTree("Cannot update digest".into()))
+        }
+    };
+
 fn map_eval_err<T: std::fmt::Debug>(e: T) -> EvalError {
     EvalError::AvlTree(format!("{:?}", e))
 }
@@ -444,6 +504,86 @@ mod tests {
         }
     }
 
+    #[test]
+    fn eval_avl_remove() {
+        let mut prover = populate_tree(vec![(vec![1u8], 10u64.to_be_bytes().to_vec())]);
+        let initial_digest =
+            ADDigest::sigma_parse_bytes(&prover.digest().unwrap().into_iter().collect::<Vec<_>>())
+                .unwrap();
+
+        let key1 = Bytes::from(vec![1u8]);
+        let op1 = Operation::Remove(key1);
+        prover.perform_one_operation(&op1).unwrap();
+        let final_digest =
+            ADDigest::sigma_parse_bytes(&prover.digest().unwrap().into_iter().collect::<Vec<_>>())
+                .unwrap();
+        let proof: Constant = prover
+            .generate_proof()
+            .into_iter()
+            .collect::<Vec<_>>()
+            .into();
+
+        let tree_flags = AvlTreeFlags::new(false, false, true);
+        let obj = Expr::Const(
+            AvlTreeData {
+                digest: initial_digest,
+                tree_flags,
+                key_length: 1,
+                value_length_opt: None,
+            }
+            .into(),
+        );
+
+        let key1 = Literal::try_from(vec![1u8]).unwrap();
+        let keys = Constant {
+            tpe: SType::SColl(Box::new(SType::SColl(Box::new(SType::SByte)))),
+            v: Literal::Coll(CollKind::WrappedColl {
+                items: vec![key1],
+                elem_tpe: SType::SColl(Box::new(SType::SByte)),
+            }),
+        };
+        let expr: Expr = MethodCall::new(
+            obj,
+            savltree::REMOVE_METHOD.clone(),
+            vec![keys.into(), proof.into()],
+        )
+        .unwrap()
+        .into();
+
+        let res = eval_out_wo_ctx::<Value>(&expr);
+        if let Value::Opt(opt) = res {
+            if let Some(Value::AvlTree(avl)) = *opt {
+                assert_eq!(avl.digest, final_digest);
+            } else {
+                unreachable!();
+            }
+        } else {
+            unreachable!();
+        }
+    }
+
+    fn populate_tree(entries: Vec<(Vec<u8>, Vec<u8>)>) -> BatchAVLProver {
+        let mut prover = BatchAVLProver::new(
+            AVLTree::new(
+                |digest| Node::LabelOnly(NodeHeader::new(Some(*digest), None)),
+                1,
+                None,
+            ),
+            true,
+        );
+
+        for (key, value) in entries {
+            let op = Operation::Insert(KeyValue {
+                key: Bytes::from(key),
+                value: Bytes::from(value),
+            });
+            prover.perform_one_operation(&op).unwrap();
+        }
+
+        prover.generate_proof();
+        prover
+    }
+
     fn mk_pair(x: u8, y: u64) -> [Literal; 2] {
         [
             Literal::try_from(vec![x]).unwrap(),
diff --git a/ergotree-ir/src/types/savltree.rs b/ergotree-ir/src/types/savltree.rs
index 56cafdfbb..1361c061c 100644
--- a/ergotree-ir/src/types/savltree.rs
+++ b/ergotree-ir/src/types/savltree.rs
@@ -33,6 +33,8 @@ pub const UPDATE_OPERATIONS_METHOD_ID: MethodId = MethodId(8);
 pub const INSERT_METHOD_ID: MethodId = MethodId(12);
 /// AvlTree.updateDigest property
 pub const UPDATE_DIGEST_METHOD_ID: MethodId = MethodId(15);
+/// AvlTree.remove property
+pub const REMOVE_METHOD_ID: MethodId = MethodId(13);
 
 lazy_static! {
     /// AvlTree method descriptors
@@ -48,6 +50,7 @@ lazy_static! {
             &IS_REMOVE_ALLOWED_METHOD_DESC,
             &UPDATE_OPERATIONS_METHOD_DESC,
             &UPDATE_DIGEST_METHOD_DESC,
+            &REMOVE_METHOD_DESC,
         ]
     ;
 }
@@ -213,3 +216,25 @@ lazy_static! {
     pub static ref INSERT_METHOD: SMethod =
         SMethod::new(STypeCompanion::AvlTree, INSERT_METHOD_DESC.clone(),);
 }
+
+lazy_static! {
+    static ref REMOVE_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: REMOVE_METHOD_ID,
+        name: "remove",
+        tpe: SFunc {
+            t_dom: vec![ SType::SAvlTree,
+                         SType::SColl(
+                            Box::new(
+                                SType::SColl(Box::new(SType::SByte))
+                            )
+                         ),
+                         SType::SColl(Box::new(SType::SByte)),
+                       ],
+            t_range: SType::SOption(Box::new(SType::SAvlTree)).into(),
+            tpe_params: vec![],
+        },
+    };
+    /// AvlTree.remove
+    pub static ref REMOVE_METHOD: SMethod =
+        SMethod::new(STypeCompanion::AvlTree, REMOVE_METHOD_DESC.clone(),);
+}


Test patch: