Add SColl.get, startsWith, endsWith, reverse

ergoplatform · Jan 20, 2025 · ed6eb52 · ed6eb52
1 parent 84a905e
commit ed6eb52
Show file tree

Hide file tree

Showing 3 changed files with 245 additions and 0 deletions.
diff --git a/ergotree-interpreter/src/eval.rs b/ergotree-interpreter/src/eval.rs
@@ -272,6 +272,10 @@ fn smethod_eval_fn(method: &SMethod) -> Result<EvalFn, EvalError> {
             scoll::PATCH_METHOD_ID => self::scoll::PATCH_EVAL_FN,
             scoll::UPDATED_METHOD_ID => self::scoll::UPDATED_EVAL_FN,
             scoll::UPDATE_MANY_METHOD_ID => self::scoll::UPDATE_MANY_EVAL_FN,
+            scoll::REVERSE_METHOD_ID => self::scoll::REVERSE_EVAL_FN,
+            scoll::STARTS_WITH_METHOD_ID => self::scoll::STARTS_WITH_EVAL_FN,
+            scoll::ENDS_WITH_METHOD_ID => self::scoll::ENDS_WITH_EVAL_FN,
+            scoll::GET_METHOD_ID => self::scoll::GET_EVAL_FN,
             method_id => {
                 return Err(EvalError::NotFound(format!(
                     "Eval fn: unknown method id in SCollection: {:?}",

diff --git a/ergotree-interpreter/src/eval/scoll.rs b/ergotree-interpreter/src/eval/scoll.rs
@@ -1,6 +1,7 @@
 use crate::eval::EvalError;
 use crate::eval::Evaluable;
 
+use alloc::boxed::Box;
 use alloc::string::ToString;
 use alloc::vec::Vec;
 use ergotree_ir::mir::constant::TryExtractInto;
@@ -334,6 +335,75 @@ pub(crate) static UPDATE_MANY_EVAL_FN: EvalFn =
         Ok(Value::Coll(CollKind::from_collection(input_tpe, &res[..])?))
     };
 
+pub(crate) static REVERSE_EVAL_FN: EvalFn = |_mc, _env, _ctx, obj, _args| {
+    let Value::Coll(coll) = obj else {
+        return Err(EvalError::UnexpectedValue(format!(
+            "Reverse: expected Coll, found {obj:?}"
+        )));
+    };
+    Ok(Value::from(coll.reverse()))
+};
+
+pub(crate) static STARTS_WITH_EVAL_FN: EvalFn = |_mc, _env, _ctx, obj, args| {
+    let Value::Coll(coll) = obj else {
+        return Err(EvalError::UnexpectedValue(format!(
+            "endsWith: expected Coll, found {obj:?}"
+        )));
+    };
+    let Some(Value::Coll(prefix)) = args.first() else {
+        return Err(EvalError::UnexpectedValue(format!(
+            "startsWith: expected Coll argument, found {:?}",
+            args.first(),
+        )));
+    };
+    if prefix.elem_tpe() != coll.elem_tpe() {
+        return Err(EvalError::UnexpectedValue(format!(
+            "startsWith: expected prefix to be of type {:?}, found {:?}",
+            coll.elem_tpe(),
+            prefix.elem_tpe()
+        )));
+    }
+    Ok(Value::from(coll.starts_with(prefix)))
+};
+
+pub(crate) static ENDS_WITH_EVAL_FN: EvalFn = |_mc, _env, _ctx, obj, args| {
+    let Value::Coll(coll) = obj else {
+        return Err(EvalError::UnexpectedValue(format!(
+            "endsWith: expected Coll, found {obj:?}"
+        )));
+    };
+    let Some(Value::Coll(suffix)) = args.first() else {
+        return Err(EvalError::UnexpectedValue(format!(
+            "endsWith: expected Coll argument, found {:?}",
+            args.first(),
+        )));
+    };
+    if suffix.elem_tpe() != coll.elem_tpe() {
+        return Err(EvalError::UnexpectedValue(format!(
+            "endsWith: expected suffix to be of type {:?}, found {:?}",
+            coll.elem_tpe(),
+            suffix.elem_tpe()
+        )));
+    }
+    Ok(Value::from(coll.ends_with(suffix)))
+};
+
+pub(crate) static GET_EVAL_FN: EvalFn = |_mc, _env, _ctx, obj, args| {
+    let Value::Coll(coll) = obj else {
+        return Err(EvalError::UnexpectedValue(format!(
+            "get: expected Coll, found {obj:?}"
+        )));
+    };
+    let index = args
+        .first()
+        .cloned()
+        .ok_or_else(|| EvalError::UnexpectedValue("Get: index argument not found".into()))?
+        .try_extract_into::<i32>()?;
+    Ok(Value::Opt(Box::new(
+        index.try_into().ok().and_then(|index| coll.get_val(index)),
+    )))
+};
+
 #[allow(clippy::unwrap_used)]
 #[cfg(test)]
 #[cfg(feature = "arbitrary")]
@@ -790,4 +860,85 @@ mod tests {
         .into();
         assert!(try_eval_out_wo_ctx::<Vec<i64>>(&expr).is_err());
     }
+
+    #[test]
+    fn eval_reverse() {
+        let arr = vec![1i64, 2i64, 3i64];
+        let coll_const: Constant = arr.clone().into();
+        let expr: Expr = MethodCall::new(
+            coll_const.into(),
+            scoll::REVERSE_METHOD
+                .clone()
+                .with_concrete_types(&[(STypeVar::t(), SType::SLong)].into_iter().collect()),
+            vec![],
+        )
+        .unwrap()
+        .into();
+        assert_eq!(
+            eval_out_wo_ctx::<Vec<i64>>(&expr),
+            arr.into_iter().rev().collect::<Vec<_>>(),
+        );
+    }
+    #[test]
+    fn eval_starts_with() {
+        fn starts_with(input: Vec<i64>, prefix: Vec<i64>) -> bool {
+            let mc: Expr = MethodCall::new(
+                Constant::from(input).into(),
+                scoll::STARTS_WITH_METHOD
+                    .clone()
+                    .with_concrete_types(&[(STypeVar::t(), SType::SLong)].into_iter().collect()),
+                vec![Constant::from(prefix).into()],
+            )
+            .unwrap()
+            .into();
+            eval_out_wo_ctx(&mc)
+        }
+        assert!(starts_with(vec![1, 2, 3], vec![1, 2]));
+        assert!(starts_with(vec![1, 2, 3], vec![1, 2, 3]));
+        assert!(!starts_with(vec![1, 2, 3], vec![1, 2, 4]));
+        assert!(!starts_with(vec![1, 2, 3], vec![1, 2, 3, 4]));
+        assert!(starts_with(vec![], vec![]));
+        assert!(starts_with(vec![1, 2], vec![]));
+    }
+    #[test]
+    fn eval_ends_with() {
+        fn ends_with(input: Vec<i64>, suffix: Vec<i64>) -> bool {
+            let mc: Expr = MethodCall::new(
+                Constant::from(input).into(),
+                scoll::ENDS_WITH_METHOD
+                    .clone()
+                    .with_concrete_types(&[(STypeVar::t(), SType::SLong)].into_iter().collect()),
+                vec![Constant::from(suffix).into()],
+            )
+            .unwrap()
+            .into();
+            eval_out_wo_ctx(&mc)
+        }
+        assert!(!ends_with(vec![1, 2, 3], vec![1, 2]));
+        assert!(ends_with(vec![1, 2, 3], vec![2, 3]));
+        assert!(!ends_with(vec![1, 2, 3], vec![2, 3, 4]));
+        assert!(ends_with(vec![1, 2, 3], vec![1, 2, 3]));
+        assert!(ends_with(vec![], vec![]));
+        assert!(ends_with(vec![1, 2], vec![]));
+    }
+    #[test]
+    fn eval_get() {
+        fn get(input: Vec<i64>, index: i32) -> Option<i64> {
+            let mc: Expr = MethodCall::new(
+                Constant::from(input).into(),
+                scoll::GET_METHOD
+                    .clone()
+                    .with_concrete_types(&[(STypeVar::t(), SType::SLong)].into_iter().collect()),
+                vec![Constant::from(index).into()],
+            )
+            .unwrap()
+            .into();
+            eval_out_wo_ctx(&mc)
+        }
+        assert_eq!(get(vec![1, 2], 0), Some(1));
+        assert_eq!(get(vec![1, 2], 1), Some(2));
+        assert_eq!(get(vec![1, 2], -1), None);
+        assert_eq!(get(vec![1, 2], 2), None);
+        assert_eq!(get(vec![], 0), None);
+    }
 }
diff --git a/ergotree-ir/src/types/scoll.rs b/ergotree-ir/src/types/scoll.rs
@@ -32,6 +32,14 @@ pub const PATCH_METHOD_ID: MethodId = MethodId(19);
 pub const UPDATED_METHOD_ID: MethodId = MethodId(20);
 /// Coll.updateMany
 pub const UPDATE_MANY_METHOD_ID: MethodId = MethodId(21);
+/// Coll.reverse
+pub const REVERSE_METHOD_ID: MethodId = MethodId(30);
+/// Coll.startsWith
+pub const STARTS_WITH_METHOD_ID: MethodId = MethodId(31);
+/// Coll.endsWith
+pub const ENDS_WITH_METHOD_ID: MethodId = MethodId(32);
+/// Coll.get
+pub const GET_METHOD_ID: MethodId = MethodId(33);
 
 lazy_static! {
     /// Coll method descriptors
@@ -44,6 +52,10 @@ lazy_static! {
             &UPDATED_METHOD_DESC,
             &UPDATE_MANY_METHOD_DESC,
             &PATCH_METHOD_DESC,
+            &REVERSE_METHOD_DESC,
+            &STARTS_WITH_METHOD_DESC,
+            &ENDS_WITH_METHOD_DESC,
+            &GET_METHOD_DESC
         ]
     ;
 }
@@ -185,6 +197,75 @@ lazy_static! {
     pub static ref UPDATE_MANY_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, UPDATE_MANY_METHOD_DESC.clone());
 }
 
+lazy_static! {
+    static ref REVERSE_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: REVERSE_METHOD_ID,
+        name: "reverse",
+        tpe: SFunc::new(
+            vec![SType::SColl(SType::STypeVar(STypeVar::t()).into())],
+            SType::SColl(SType::STypeVar(STypeVar::t()).into())
+        ),
+        explicit_type_args: vec![],
+        min_version: ErgoTreeVersion::V3
+    };
+    /// Coll.reverse
+    pub static ref REVERSE_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, REVERSE_METHOD_DESC.clone());
+}
+
+lazy_static! {
+    static ref STARTS_WITH_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: STARTS_WITH_METHOD_ID,
+        name: "startsWith",
+        tpe: SFunc::new(
+            vec![
+                SType::SColl(SType::STypeVar(STypeVar::t()).into()),
+                SType::SColl(SType::STypeVar(STypeVar::t()).into())
+            ],
+            SType::SBoolean
+        ),
+        explicit_type_args: vec![],
+        min_version: ErgoTreeVersion::V3
+    };
+    /// Coll.startsWith
+    pub static ref STARTS_WITH_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, STARTS_WITH_METHOD_DESC.clone());
+}
+
+lazy_static! {
+    static ref ENDS_WITH_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: ENDS_WITH_METHOD_ID,
+        name: "endsWith",
+        tpe: SFunc::new(
+            vec![
+                SType::SColl(SType::STypeVar(STypeVar::t()).into()),
+                SType::SColl(SType::STypeVar(STypeVar::t()).into())
+            ],
+            SType::SBoolean
+        ),
+        explicit_type_args: vec![],
+        min_version: ErgoTreeVersion::V3
+    };
+    /// Coll.endsWith
+    pub static ref ENDS_WITH_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, ENDS_WITH_METHOD_DESC.clone());
+}
+
+lazy_static! {
+    static ref GET_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: GET_METHOD_ID,
+        name: "get",
+        tpe: SFunc::new(
+            vec![
+                SType::SColl(SType::STypeVar(STypeVar::t()).into()),
+                SType::SInt
+            ],
+            SType::SOption(SType::STypeVar(STypeVar::t()).into())
+        ),
+        explicit_type_args: vec![],
+        min_version: ErgoTreeVersion::V3
+    };
+    /// Coll.get
+    pub static ref GET_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, GET_METHOD_DESC.clone());
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -203,5 +284,14 @@ mod tests {
             SMethod::from_ids(TYPE_CODE, UPDATE_MANY_METHOD_ID).map(|e| e.name())
                 == Ok("updateMany")
         );
+        assert!(SMethod::from_ids(TYPE_CODE, REVERSE_METHOD_ID).map(|e| e.name()) == Ok("reverse"));
+        assert!(
+            SMethod::from_ids(TYPE_CODE, STARTS_WITH_METHOD_ID).map(|e| e.name())
+                == Ok("startsWith")
+        );
+        assert!(
+            SMethod::from_ids(TYPE_CODE, ENDS_WITH_METHOD_ID).map(|e| e.name()) == Ok("endsWith")
+        );
+        assert!(SMethod::from_ids(TYPE_CODE, GET_METHOD_ID).map(|e| e.name()) == Ok("get"));
     }
 }