Add AvlTree.insertOrUpdate method (6.0)

Author: sethdusek

ergotree-interpreter/src/eval.rs

@@ -232,6 +232,7 @@ fn smethod_eval_fn(method: &SMethod) -> Result<EvalFn, EvalError> {
             savltree::CONTAINS_METHOD_ID => self::savltree::CONTAINS_EVAL_FN,
             savltree::REMOVE_METHOD_ID => self::savltree::REMOVE_EVAL_FN,
             savltree::UPDATE_METHOD_ID => self::savltree::UPDATE_EVAL_FN,
+            savltree::INSERT_OR_UPDATE_METHOD_ID => self::savltree::INSERT_OR_UPDATE_EVAL_FN,
             method_id => {
                 return Err(EvalError::NotFound(format!(
                     "Eval fn: unknown method id in SAvlTree: {:?}",

ergotree-interpreter/src/eval/savltree.rs

@@ -428,6 +428,65 @@ pub(crate) static UPDATE_EVAL_FN: EvalFn =
         })
     };
 
+pub(crate) static INSERT_OR_UPDATE_EVAL_FN: EvalFn = |_mc, _env, _ctx, obj, args| {
+    let mut avl_tree_data = obj.try_extract_into::<AvlTreeData>()?;
+
+    if !avl_tree_data.tree_flags.insert_allowed() || !avl_tree_data.tree_flags.update_allowed() {
+        return Ok(Value::Opt(None));
+    }
+
+    let entries = {
+        let v = args
+            .first()
+            .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::InsertOrUpdate(KeyValue {
+                key: key.into(),
+                value: value.into(),
+            }))
+            .is_err()
+        {
+            break;
+        }
+    }
+    Ok(if let Some(new_digest) = bv.digest() {
+        let digest = ADDigest::scorex_parse_bytes(&new_digest)?;
+        avl_tree_data.digest = digest;
+        Value::Opt(Some(Box::new(Value::AvlTree(avl_tree_data.into()))))
+    } else {
+        Value::Opt(None)
+    })
+};
+
 fn map_eval_err<T: core::fmt::Debug>(e: T) -> EvalError {
     EvalError::AvlTree(format!("{:?}", e))
 }
@@ -708,6 +767,75 @@ mod tests {
             Value::Opt(None)
         );
     }
+
+    #[test]
+    fn eval_avl_insert_or_update() {
+        let mut prover = BatchAVLProver::new(
+            AVLTree::new(
+                |digest| Node::LabelOnly(NodeHeader::new(Some(*digest), None)),
+                1,
+                None,
+            ),
+            true,
+        );
+        let initial_digest = ADDigest::scorex_parse_bytes(&prover.digest().unwrap()).unwrap();
+        let key1 = Bytes::from(vec![1u8]);
+        let op = Operation::InsertOrUpdate(KeyValue {
+            key: key1,
+            value: Bytes::from(10u64.to_be_bytes().to_vec()),
+        });
+        prover.perform_one_operation(&op).unwrap();
+        prover.perform_one_operation(&op).unwrap();
+        let final_digest = ADDigest::scorex_parse_bytes(&prover.digest().unwrap()).unwrap();
+        let proof: Constant = prover
+            .generate_proof()
+            .into_iter()
+            .collect::<Vec<_>>()
+            .into();
+
+        let tree_flags = AvlTreeFlags::new(true, 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(1u8, 10u64).into());
+        let entries = Constant {
+            tpe: SType::SColl(Arc::new(SType::STuple(STuple::pair(
+                SType::SColl(Arc::new(SType::SByte)),
+                SType::SColl(Arc::new(SType::SByte)),
+            )))),
+            v: Literal::Coll(CollKind::WrappedColl {
+                items: Arc::new([pair1.clone(), pair1.clone()]),
+                elem_tpe: SType::STuple(STuple::pair(
+                    SType::SColl(Arc::new(SType::SByte)),
+                    SType::SColl(Arc::new(SType::SByte)),
+                )),
+            }),
+        };
+        let expr: Expr = MethodCall::new(
+            obj.clone(),
+            savltree::INSERT_OR_UPDATE_METHOD.clone(),
+            vec![entries.clone().into(), proof.clone().into()],
+        )
+        .unwrap()
+        .into();
+
+        let res = eval_out_wo_ctx::<Value>(&expr);
+        if let Value::Opt(opt) = res {
+            if let Some(Value::AvlTree(avl)) = opt.as_deref() {
+                assert_eq!(avl.digest, final_digest);
+            } else {
+                unreachable!();
+            }
+        } else {
+            unreachable!();
+        }
+    }
     proptest! {
         #[test]
         fn eval_avl_digest(v in any::<AvlTreeData>()) {

ergotree-ir/src/types/savltree.rs

@@ -48,6 +48,8 @@ pub const REMOVE_METHOD_ID: MethodId = MethodId(14);
 pub const UPDATE_METHOD_ID: MethodId = MethodId(13);
 /// AvlTree.updateDigest property
 pub const UPDATE_DIGEST_METHOD_ID: MethodId = MethodId(15);
+/// AvlTree.insertOrUpdate property
+pub const INSERT_OR_UPDATE_METHOD_ID: MethodId = MethodId(16);
 
 lazy_static! {
     /// AvlTree method descriptors
@@ -68,6 +70,7 @@ lazy_static! {
             REMOVE_METHOD_DESC.clone(),
             UPDATE_METHOD_DESC.clone(),
             UPDATE_DIGEST_METHOD_DESC.clone(),
+            INSERT_OR_UPDATE_METHOD_DESC.clone(),
         ]
     ;
 }
@@ -370,6 +373,35 @@ lazy_static! {
         SMethod::new(STypeCompanion::AvlTree, UPDATE_DIGEST_METHOD_DESC.clone(),);
 }
 
+lazy_static! {
+    static ref INSERT_OR_UPDATE_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: INSERT_OR_UPDATE_METHOD_ID,
+        name: "insertOrUpdate",
+        tpe: SFunc {
+            t_dom: vec![ SType::SAvlTree,
+                         SType::SColl(
+                           Arc::new(
+                               SType::STuple(
+                                   STuple::pair(
+                                       SType::SColl(Arc::new(SType::SByte)),
+                                       SType::SColl(Arc::new(SType::SByte))
+                                   )
+                               )
+                           )
+                         ),
+                         SType::SColl(Arc::new(SType::SByte)),
+                       ],
+            t_range: SType::SOption(Arc::new(SType::SAvlTree)).into(),
+            tpe_params: vec![],
+        },
+        explicit_type_args: vec![],
+        min_version: ErgoTreeVersion::V3
+    };
+    /// AvlTree.insertOrUpdate
+    pub static ref INSERT_OR_UPDATE_METHOD: SMethod =
+        SMethod::new(STypeCompanion::AvlTree, INSERT_OR_UPDATE_METHOD_DESC.clone(),);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;