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

Issue:
[100 SigmaUSD] `avlTree` predefined function for AVL tree creation
Create `CreateAvlTree` IR node. Evaluation should return `Value::AvlTree`.

Reference Scala implementation - https://github.com/ScorexFoundation/sigmastate-interpreter/blob/c863a9b1a82589e47b15f76f3affdb30a475e740/sigmastate/src/main/scala/sigmastate/trees.scala#L197-L210
https://github.com/ScorexFoundation/sigmastate-interpreter/blob/c863a9b1a82589e47b15f76f3affdb30a475e740/sigmastate/src/main/scala/sigmastate/lang/SigmaPredef.scala#L293-L306



Discussion / hints:
will try this one
I have some issues in evaluation when I get byte array from digest argument. Its type is Vec<i8> but AvlTreeData accepts ADDigest which can be constructed from Vec<u8>.
I was trying to convert it but it turns out I'd have to use transmute which can be unsafe so somebody on Rust discord recommended me to use `bytemuck` crate.
Do you have any other ideas?
> I have some issues in evaluation when I get byte array from digest argument. Its type is Vec but AvlTreeData accepts ADDigest which can be constructed from Vec. I was trying to convert it but it turns out I'd have to use transmute which can be unsafe so somebody on Rust discord recommended me to use `bytemuck` crate. Do you have any other ideas?

This https://github.com/ergoplatform/sigma-rust/blob/8aea793155ae9259bb017d2430eb39ac11b7dc69/ergotree-interpreter/src/eval/savltree.rs#L83-L83 should work.
Sorry Github hid my type :P I got `Vec<i8>` and ADDigest needs `Vec<u8>` so try_from doesn't work.
> Sorry Github hid my type :P I got `Vec<i8>` and ADDigest needs `Vec<u8>` so try_from doesn't work.

Oh, we use https://github.com/ergoplatform/sigma-rust/blob/d31d096350dc3c351009fa27f556bba2746dbbdc/ergotree-ir/src/util.rs#L33-L33
for that. 
Ok I solved that.
I have another issue when testing serialization round-trip. It says `not yet implemented: SOption(SInt) not yet implemented`.

SOption is not yet implemented in `const_with_type` but I'm not quite sure how it works.
```rust
SType::SOption(tpe) => const_with_type(*tpe),
```
Would this be okay?

> Ok I solved that. I have another issue when testing serialization round-trip. It says `not yet implemented: SOption(SInt) not yet implemented`.
> 
> SOption is not yet implemented in `const_with_type` but I'm not quite sure how it works.
> 
> ```rust
> SType::SOption(tpe) => const_with_type(*tpe),
> ```
> 
> Would this be okay?

No, that will create SInt constant, but we need to wrap it in SOption. Check out `SType::SColl` in that code below. The same matching by inner type is needed for SOption. For Sint it would be something like:
```rust
SType::SInt => any::<Option<i32>>().prop_map_into().boxed(),
```
should work.
`thread 'mir::avl_tree::tests::ser_roundtrip' panicked at 'Test failed: serialization failed: NotSupported("Option"); minimal failing input: v = CreateAvlTree { flags: Const(Constant { tpe: SByte, v: Byte(0) }), digest: Const(Constant { tpe: SColl(SByte), v: Coll(NativeColl(CollByte([]))) }), key_length: Const(Constant { tpe: SInt, v: Int(0) }), value_length: Const(Constant { tpe: SOption(SInt), v: Opt(None) }) }`

How should Option be serialized?
https://github.com/ScorexFoundation/sigmastate-interpreter/blob/25251c1313b0131835f92099f02cef8a5d932b5e/sigmastate/src/main/scala/sigmastate/serialization/DataSerializer.scala#L26-L26 Scala's version doesn't implement Option too?
Check out how `ByIndex::default` is serialized:
https://github.com/ergoplatform/sigma-rust/blob/d31d096350dc3c351009fa27f556bba2746dbbdc/ergotree-ir/src/mir/coll_by_index.rs#L78-L78


Fix patch:
diff --git a/ergotree-interpreter/src/eval.rs b/ergotree-interpreter/src/eval.rs
index 75d369a3b..93cf09dcb 100644
--- a/ergotree-interpreter/src/eval.rs
+++ b/ergotree-interpreter/src/eval.rs
@@ -44,6 +44,7 @@ pub(crate) mod coll_slice;
 pub(crate) mod collection;
 pub(crate) mod cost_accum;
 pub(crate) mod costs;
+pub(crate) mod create_avl_tree;
 pub(crate) mod create_prove_dh_tuple;
 pub(crate) mod create_provedlog;
 pub(crate) mod decode_point;
diff --git a/ergotree-interpreter/src/eval/create_avl_tree.rs b/ergotree-interpreter/src/eval/create_avl_tree.rs
new file mode 100644
index 000000000..6104e4789
--- /dev/null
+++ b/ergotree-interpreter/src/eval/create_avl_tree.rs
@@ -0,0 +1,86 @@
+use super::Evaluable;
+use crate::eval::env::Env;
+use crate::eval::EvalContext;
+use crate::eval::EvalError;
+use ergotree_ir::chain::digest32::ADDigest;
+use ergotree_ir::mir::avl_tree_data::{AvlTreeData, AvlTreeFlags};
+use ergotree_ir::mir::constant::TryExtractInto;
+use ergotree_ir::mir::create_avl_tree::CreateAvlTree;
+use ergotree_ir::mir::value::Value;
+use ergotree_ir::util::AsVecU8;
+use std::convert::TryFrom;
+
+impl Evaluable for CreateAvlTree {
+    fn eval(&self, env: &Env, ctx: &mut EvalContext) -> Result<Value, EvalError> {
+        let flags_v = self.flags.eval(env, ctx)?.try_extract_into::<i8>()? as u8;
+        let digest_v = self.digest.eval(env, ctx)?.try_extract_into::<Vec<i8>>()?;
+        let key_length = self.key_length.eval(env, ctx)?.try_extract_into::<i32>()? as u32;
+        let value_length_opt = match self.value_length.clone() {
+            Some(expr) => Some(Box::new(
+                expr.eval(env, ctx)?.try_extract_into::<i32>()? as u32
+            )),
+            None => None,
+        };
+
+        let tree_flags = AvlTreeFlags::parse(flags_v);
+        let digest = ADDigest::try_from(digest_v.as_vec_u8()).map_err(map_eval_err)?;
+
+        Ok(Value::AvlTree(Box::new(AvlTreeData {
+            tree_flags,
+            digest,
+            key_length,
+            value_length_opt,
+        })))
+    }
+}
+
+fn map_eval_err<T: std::fmt::Debug>(e: T) -> EvalError {
+    EvalError::AvlTree(format!("{:?}", e))
+}
+
+#[allow(clippy::unwrap_used)]
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::eval::tests::eval_out_wo_ctx;
+
+    use ergotree_ir::chain::digest32::ADDigest;
+    use ergotree_ir::mir::{
+        avl_tree_data::{AvlTreeData, AvlTreeFlags},
+        expr::Expr,
+    };
+    use ergotree_ir::serialization::SigmaSerializable;
+    use scorex_crypto_avltree::authenticated_tree_ops::AuthenticatedTreeOps;
+    use scorex_crypto_avltree::batch_avl_prover::BatchAVLProver;
+    use scorex_crypto_avltree::batch_node::{AVLTree, Node, NodeHeader};
+
+    #[test]
+    fn eval_create_avl_tree() {
+        let prover = BatchAVLProver::new(
+            AVLTree::new(
+                |digest| Node::LabelOnly(NodeHeader::new(Some(*digest), None)),
+                1,
+                None,
+            ),
+            true,
+        );
+        let initial_digest =
+            ADDigest::sigma_parse_bytes(&prover.digest().unwrap().into_iter().collect::<Vec<_>>())
+                .unwrap();
+        let flags = AvlTreeFlags::new(false, false, false);
+        let expr: Expr = CreateAvlTree::new(
+            Expr::Const(flags.clone().into()),
+            Expr::Const(initial_digest.clone().into()),
+            1.into(),
+            None,
+        )
+        .unwrap()
+        .into();
+
+        let tree = eval_out_wo_ctx::<AvlTreeData>(&expr);
+        assert_eq!(tree.digest, initial_digest);
+        assert_eq!(tree.key_length, 1);
+        assert_eq!(tree.value_length_opt, None);
+        assert_eq!(tree.tree_flags, flags);
+    }
+}
diff --git a/ergotree-interpreter/src/eval/expr.rs b/ergotree-interpreter/src/eval/expr.rs
index e639220a5..3197fb9eb 100644
--- a/ergotree-interpreter/src/eval/expr.rs
+++ b/ergotree-interpreter/src/eval/expr.rs
@@ -79,6 +79,7 @@ impl Evaluable for Expr {
             Expr::ExtractBytes(op) => op.eval(env, ctx),
             Expr::ExtractBytesWithNoRef(op) => op.eval(env, ctx),
             Expr::TreeLookup(op) => op.eval(env, ctx),
+            Expr::CreateAvlTree(op) => op.eval(env, ctx),
         }
     }
 }
diff --git a/ergotree-ir/src/mir.rs b/ergotree-ir/src/mir.rs
index ab77336c2..1113dad3f 100644
--- a/ergotree-ir/src/mir.rs
+++ b/ergotree-ir/src/mir.rs
@@ -37,6 +37,8 @@ pub mod coll_slice;
 /// Collection of elements
 pub mod collection;
 pub mod constant;
+/// Creation of AVL tree
+pub mod create_avl_tree;
 /// Create proveDHTuple
 pub mod create_prove_dh_tuple;
 /// Create proveDlog from GroupElement(PK)
diff --git a/ergotree-ir/src/mir/constant.rs b/ergotree-ir/src/mir/constant.rs
index c85a29309..60845bc5c 100644
--- a/ergotree-ir/src/mir/constant.rs
+++ b/ergotree-ir/src/mir/constant.rs
@@ -2,6 +2,7 @@
 
 use crate::base16_str::Base16Str;
 use crate::bigint256::BigInt256;
+use crate::chain::digest32::ADDigest;
 use crate::chain::ergo_box::ErgoBox;
 use crate::mir::value::CollKind;
 use crate::serialization::SigmaSerializable;
@@ -24,6 +25,7 @@ mod constant_placeholder;
 pub use constant_placeholder::*;
 
 use super::avl_tree_data::AvlTreeData;
+use super::avl_tree_data::AvlTreeFlags;
 use super::value::NativeColl;
 use super::value::StoreWrapped;
 use super::value::Value;
@@ -393,6 +395,24 @@ impl From<AvlTreeData> for Constant {
     }
 }
 
+impl From<AvlTreeFlags> for Constant {
+    fn from(a: AvlTreeFlags) -> Self {
+        Constant {
+            tpe: SType::SByte,
+            v: Literal::Byte(a.serialize() as i8),
+        }
+    }
+}
+
+impl From<ADDigest> for Constant {
+    fn from(a: ADDigest) -> Self {
+        Constant {
+            tpe: SType::SColl(Box::new(SType::SByte)),
+            v: Literal::Coll(CollKind::NativeColl(NativeColl::CollByte(a.into()))),
+        }
+    }
+}
+
 #[allow(clippy::unwrap_used)]
 #[allow(clippy::from_over_into)]
 #[impl_for_tuples(2, 4)]
@@ -756,6 +776,14 @@ pub(crate) mod arbitrary {
             SType::SBox => any::<ErgoBox>().prop_map_into().boxed(),
             SType::SAvlTree => any::<AvlTreeData>().prop_map_into().boxed(),
             // SType::SOption(tpe) =>
+            SType::SOption(tpe) => match *tpe {
+                SType::SBoolean => any::<Option<bool>>().prop_map_into().boxed(),
+                SType::SByte => any::<Option<i8>>().prop_map_into().boxed(),
+                SType::SShort => any::<Option<i16>>().prop_map_into().boxed(),
+                SType::SInt => any::<Option<i32>>().prop_map_into().boxed(),
+                SType::SLong => any::<Option<i64>>().prop_map_into().boxed(),
+                _ => todo!(),
+            },
             SType::SColl(elem_tpe) => match *elem_tpe {
                 SType::SBoolean => vec(any::<bool>(), 0..400).prop_map_into().boxed(),
                 SType::SByte => vec(any::<u8>(), 0..400).prop_map_into().boxed(),
diff --git a/ergotree-ir/src/mir/create_avl_tree.rs b/ergotree-ir/src/mir/create_avl_tree.rs
new file mode 100644
index 000000000..820d4e52f
--- /dev/null
+++ b/ergotree-ir/src/mir/create_avl_tree.rs
@@ -0,0 +1,142 @@
+//! Create an AVL tree
+
+use super::expr::Expr;
+use crate::has_opcode::HasStaticOpCode;
+use crate::mir::expr::InvalidArgumentError;
+use crate::serialization::op_code::OpCode;
+use crate::serialization::sigma_byte_reader::SigmaByteRead;
+use crate::serialization::sigma_byte_writer::SigmaByteWrite;
+use crate::serialization::{SigmaParsingError, SigmaSerializable, SigmaSerializeResult};
+use crate::types::stype::SType;
+
+/// Creates an AVL tree
+#[derive(PartialEq, Eq, Debug, Clone)]
+pub struct CreateAvlTree {
+    /// Value of type Byte
+    pub flags: Box<Expr>,
+    /// Byte array with SColl(SByte) expr type
+    pub digest: Box<Expr>,
+    /// Value of type Int
+    pub key_length: Box<Expr>,
+    /// Optional value of type Int
+    pub value_length: Option<Box<Expr>>,
+}
+
+impl CreateAvlTree {
+    /// Creates new AVL Tree
+    pub fn new(
+        flags: Expr,
+        digest: Expr,
+        key_length: Expr,
+        value_length: Option<Box<Expr>>,
+    ) -> Result<Self, InvalidArgumentError> {
+        flags.check_post_eval_tpe(&SType::SByte)?;
+        digest.check_post_eval_tpe(&SType::SColl(Box::new(SType::SByte)))?;
+        key_length.check_post_eval_tpe(&SType::SInt)?;
+        if !value_length
+            .clone()
+            .map(|expr| expr.post_eval_tpe() == SType::SInt)
+            .unwrap_or(true)
+        {
+            return Err(InvalidArgumentError(format!(
+                "CreateAvlTree: expected value_length type to be Option[Int], got {0:?}",
+                value_length
+            )));
+        }
+
+        Ok(Self {
+            flags: flags.into(),
+            digest: digest.into(),
+            key_length: key_length.into(),
+            value_length,
+        })
+    }
+
+    /// Type
+    pub fn tpe(&self) -> SType {
+        SType::SAvlTree
+    }
+}
+
+impl HasStaticOpCode for CreateAvlTree {
+    const OP_CODE: OpCode = OpCode::AVL_TREE;
+}
+
+impl SigmaSerializable for CreateAvlTree {
+    fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SigmaParsingError> {
+        let flags = Expr::sigma_parse(r)?;
+        let digest = Expr::sigma_parse(r)?;
+        let key_length = Expr::sigma_parse(r)?;
+        let value_length = Option::<Box<Expr>>::sigma_parse(r)?;
+        Ok(Self::new(flags, digest, key_length, value_length)?)
+    }
+
+    fn sigma_serialize<W: SigmaByteWrite>(&self, w: &mut W) -> SigmaSerializeResult {
+        self.flags.sigma_serialize(w)?;
+        self.digest.sigma_serialize(w)?;
+        self.key_length.sigma_serialize(w)?;
+        self.value_length.sigma_serialize(w)
+    }
+}
+
+#[allow(clippy::unwrap_used)]
+#[cfg(feature = "arbitrary")]
+/// Arbitrary impl
+mod arbitrary {
+    use crate::mir::expr::arbitrary::ArbExprParams;
+
+    use super::*;
+    use proptest::prelude::*;
+    use proptest::result::Probability;
+
+    impl Arbitrary for CreateAvlTree {
+        type Strategy = BoxedStrategy<Self>;
+        type Parameters = ();
+
+        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+            (
+                any_with::<Expr>(ArbExprParams {
+                    tpe: SType::SByte,
+                    depth: 0,
+                }),
+                any_with::<Expr>(ArbExprParams {
+                    tpe: SType::SColl(Box::new(SType::SByte)),
+                    depth: 0,
+                }),
+                any_with::<Expr>(ArbExprParams {
+                    tpe: SType::SInt,
+                    depth: 0,
+                }),
+                any_with::<Option<Box<Expr>>>((
+                    Probability::default(),
+                    ArbExprParams {
+                        tpe: SType::SInt,
+                        depth: 0,
+                    },
+                )),
+            )
+                .prop_map(|(flags, digest, key_length, value_length)| {
+                    Self::new(flags, digest, key_length, value_length).unwrap()
+                })
+                .boxed()
+        }
+    }
+}
+
+#[cfg(test)]
+#[cfg(feature = "arbitrary")]
+#[allow(clippy::panic)]
+mod tests {
+    use super::*;
+    use crate::serialization::sigma_serialize_roundtrip;
+    use proptest::prelude::*;
+
+    proptest! {
+
+        #[test]
+        fn ser_roundtrip(v in any::<CreateAvlTree>()) {
+            let expr: Expr = v.into();
+            prop_assert_eq![sigma_serialize_roundtrip(&expr), expr];
+        }
+    }
+}
diff --git a/ergotree-ir/src/mir/expr.rs b/ergotree-ir/src/mir/expr.rs
index e12bc42e7..5b4e181b9 100644
--- a/ergotree-ir/src/mir/expr.rs
+++ b/ergotree-ir/src/mir/expr.rs
@@ -30,6 +30,7 @@ use super::constant::ConstantPlaceholder;
 use super::constant::Literal;
 use super::constant::TryExtractFrom;
 use super::constant::TryExtractFromError;
+use super::create_avl_tree::CreateAvlTree;
 use super::create_provedlog::CreateProveDlog;
 use super::decode_point::DecodePoint;
 use super::exponentiate::Exponentiate;
@@ -218,6 +219,8 @@ pub enum Expr {
     XorOf(XorOf),
     /// Perform a lookup by key in an AVL tree
     TreeLookup(TreeLookup),
+    /// Create an AVL tree
+    CreateAvlTree(CreateAvlTree),
 }
 
 impl Expr {
@@ -289,6 +292,7 @@ impl Expr {
             Expr::ExtractBytes(v) => v.tpe(),
             Expr::ExtractBytesWithNoRef(v) => v.tpe(),
             Expr::TreeLookup(v) => v.tpe(),
+            Expr::CreateAvlTree(v) => v.tpe(),
         }
     }
 
diff --git a/ergotree-ir/src/serialization/expr.rs b/ergotree-ir/src/serialization/expr.rs
index dd091e7de..044acd01c 100644
--- a/ergotree-ir/src/serialization/expr.rs
+++ b/ergotree-ir/src/serialization/expr.rs
@@ -31,6 +31,7 @@ use crate::mir::collection::coll_sigma_parse;
 use crate::mir::collection::coll_sigma_serialize;
 use crate::mir::constant::Constant;
 use crate::mir::constant::ConstantPlaceholder;
+use crate::mir::create_avl_tree::CreateAvlTree;
 use crate::mir::create_prove_dh_tuple::CreateProveDhTuple;
 use crate::mir::create_provedlog::CreateProveDlog;
 use crate::mir::decode_point::DecodePoint;
@@ -188,6 +189,7 @@ impl Expr {
                 Exponentiate::OP_CODE => Ok(Exponentiate::sigma_parse(r)?.into()),
                 XorOf::OP_CODE => Ok(XorOf::sigma_parse(r)?.into()),
                 TreeLookup::OP_CODE => Ok(TreeLookup::sigma_parse(r)?.into()),
+                CreateAvlTree::OP_CODE => Ok(CreateAvlTree::sigma_parse(r)?.into()),
                 o => Err(SigmaParsingError::NotImplementedOpCode(format!(
                     "{0}(shift {1})",
                     o.value(),
@@ -289,6 +291,7 @@ impl SigmaSerializable for Expr {
             Expr::ExtractBytes(op) => op.sigma_serialize_w_opcode(w),
             Expr::ExtractBytesWithNoRef(op) => op.sigma_serialize_w_opcode(w),
             Expr::TreeLookup(op) => op.sigma_serialize_w_opcode(w),
+            Expr::CreateAvlTree(op) => op.sigma_serialize_w_opcode(w),
         }
     }
 


Test patch:
diff --git a/ergotree-ir/proptest-regressions/mir/create_avl_tree.txt b/ergotree-ir/proptest-regressions/mir/create_avl_tree.txt
new file mode 100644
index 000000000..35c3d106b
--- /dev/null
+++ b/ergotree-ir/proptest-regressions/mir/create_avl_tree.txt
@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc e1a2c540ac8d9e77d201c91442af1a31ddb6e0becb8511fb0831820a6b174659 # shrinks to v = CreateAvlTree { flags: Const(Constant { tpe: SByte, v: Byte(0) }), digest: Const(Constant { tpe: SColl(SByte), v: Coll(NativeColl(CollByte([]))) }), key_length: Const(Constant { tpe: SInt, v: Int(0) }), value_length: Const(Constant { tpe: SOption(SInt), v: Opt(None) }) }