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

Issue:
[100 SigmaUSD] AvlTree.update
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#L2283-L2329


Discussion / hints:
gonna work on this too

Fix patch:
diff --git a/ergotree-interpreter/src/eval.rs b/ergotree-interpreter/src/eval.rs
index a1c037c70..5b740547d 100644
--- a/ergotree-interpreter/src/eval.rs
+++ b/ergotree-interpreter/src/eval.rs
@@ -211,6 +211,7 @@ fn smethod_eval_fn(method: &SMethod) -> Result<EvalFn, EvalError> {
             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,
+            savltree::UPDATE_METHOD_ID => self::savltree::UPDATE_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 76206bd11..f5891a558 100644
--- a/ergotree-interpreter/src/eval/savltree.rs
+++ b/ergotree-interpreter/src/eval/savltree.rs
@@ -209,6 +209,69 @@ pub(crate) static REMOVE_EVAL_FN: EvalFn =
         }
     };
 
+pub(crate) static UPDATE_EVAL_FN: EvalFn =
+    |_env, _ctx, obj, args| {
+        let mut avl_tree_data = obj.try_extract_into::<AvlTreeData>()?;
+
+        if !avl_tree_data.tree_flags.update_allowed() {
+            return Err(EvalError::AvlTree("Updates not allowed".into()));
+        }
+
+        let entries = {
+            let v = args.get(0).cloned().ok_or_else(|| {
+                EvalError::AvlTree("eval is missing first arg (entries)".to_string())
+            })?;
+            v.try_extract_into::<Vec<(Vec<u8>, 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, value) in entries {
+            if bv
+                .perform_one_operation(&Operation::Update(KeyValue {
+                    key: key.into(),
+                    value: value.into(),
+                }))
+                .is_err()
+            {
+                return Err(EvalError::AvlTree(format!(
+                    "Incorrect update 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))
 }
@@ -562,6 +625,83 @@ mod tests {
         }
     }
 
+    #[test]
+    fn eval_avl_update() {
+        let mut prover = populate_tree(vec![
+            (vec![1u8], 10u64.to_be_bytes().to_vec()),
+            (vec![2u8], 20u64.to_be_bytes().to_vec()),
+            (vec![3u8], 30u64.to_be_bytes().to_vec()),
+        ]);
+        let initial_digest =
+            ADDigest::sigma_parse_bytes(&prover.digest().unwrap().into_iter().collect::<Vec<_>>())
+                .unwrap();
+
+        let op1 = Operation::Update(KeyValue {
+            key: Bytes::from(vec![2u8]),
+            value: Bytes::from(40u64.to_be_bytes().to_vec()),
+        });
+        let op2 = Operation::Update(KeyValue {
+            key: Bytes::from(vec![3u8]),
+            value: Bytes::from(50u64.to_be_bytes().to_vec()),
+        });
+        prover.perform_one_operation(&op1).unwrap();
+        prover.perform_one_operation(&op2).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, true, false);
+        let obj = Expr::Const(
+            AvlTreeData {
+                digest: initial_digest,
+                tree_flags,
+                key_length: 1,
+                value_length_opt: None,
+            }
+            .into(),
+        );
+
+        let pair1 = Literal::Tup(mk_pair(2u8, 40u64).into());
+        let pair2 = Literal::Tup(mk_pair(3u8, 50u64).into());
+        let entries = Constant {
+            tpe: SType::SColl(Box::new(SType::STuple(STuple::pair(
+                SType::SColl(Box::new(SType::SByte)),
+                SType::SColl(Box::new(SType::SByte)),
+            )))),
+            v: Literal::Coll(CollKind::WrappedColl {
+                items: vec![pair1, pair2],
+                elem_tpe: SType::STuple(STuple::pair(
+                    SType::SColl(Box::new(SType::SByte)),
+                    SType::SColl(Box::new(SType::SByte)),
+                )),
+            }),
+        };
+        let expr: Expr = MethodCall::new(
+            obj,
+            savltree::UPDATE_METHOD.clone(),
+            vec![entries.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(
diff --git a/ergotree-ir/src/types/savltree.rs b/ergotree-ir/src/types/savltree.rs
index 1361c061c..1f3ab54f7 100644
--- a/ergotree-ir/src/types/savltree.rs
+++ b/ergotree-ir/src/types/savltree.rs
@@ -35,6 +35,8 @@ pub const INSERT_METHOD_ID: MethodId = MethodId(12);
 pub const UPDATE_DIGEST_METHOD_ID: MethodId = MethodId(15);
 /// AvlTree.remove property
 pub const REMOVE_METHOD_ID: MethodId = MethodId(13);
+/// AvlTree.update property
+pub const UPDATE_METHOD_ID: MethodId = MethodId(14);
 
 lazy_static! {
     /// AvlTree method descriptors
@@ -51,6 +53,7 @@ lazy_static! {
             &UPDATE_OPERATIONS_METHOD_DESC,
             &UPDATE_DIGEST_METHOD_DESC,
             &REMOVE_METHOD_DESC,
+            &UPDATE_METHOD_DESC,
         ]
     ;
 }
@@ -238,3 +241,30 @@ lazy_static! {
     pub static ref REMOVE_METHOD: SMethod =
         SMethod::new(STypeCompanion::AvlTree, REMOVE_METHOD_DESC.clone(),);
 }
+
+lazy_static! {
+    static ref UPDATE_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: UPDATE_METHOD_ID,
+        name: "update",
+        tpe: SFunc {
+            t_dom: vec![ SType::SAvlTree,
+                         SType::SColl(
+                            Box::new(
+                                SType::STuple(
+                                    STuple::pair(
+                                        SType::SColl(Box::new(SType::SByte)),
+                                        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.update
+    pub static ref UPDATE_METHOD: SMethod =
+        SMethod::new(STypeCompanion::AvlTree, UPDATE_METHOD_DESC.clone(),);
+}


Test patch: