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

Issue:
[150 SigmaUSD] Implement `Coll.indices` method
- Add `SMethod` to `ergotree-ir::types::scoll`(method id is 29); 
- Add `EvalFn` in `ertotree-interpreter::eval::scoll` and include it in `ergotree-interpreter::eval::smethod_eval_fn`

See existing method implementations in `scoll` and `sbox` for examples;

Reference Scala version:
eval - https://github.com/ScorexFoundation/sigmastate-interpreter/blob/b60edae20f2af3352e4064d139d45a15fc507f7d/library-impl/src/main/scala/special/collection/CollsOverArrays.scala#L74



Discussion / hints:

Fix patch:
diff --git a/ergotree-interpreter/src/eval.rs b/ergotree-interpreter/src/eval.rs
index e6e285dfc..2ec02b703 100644
--- a/ergotree-interpreter/src/eval.rs
+++ b/ergotree-interpreter/src/eval.rs
@@ -204,6 +204,8 @@ fn smethod_eval_fn(method: &SMethod) -> Result<EvalFn, EvalError> {
             scoll::INDEX_OF_METHOD_ID => self::scoll::INDEX_OF_EVAL_FN,
             scoll::FLATMAP_METHOD_ID => self::scoll::FLATMAP_EVAL_FN,
             scoll::ZIP_METHOD_ID => self::scoll::ZIP_EVAL_FN,
+            scoll::INDICES_METHOD_ID => self::scoll::INDICES_EVAL_FN,
+            scoll::UPDATED_METHOD_ID => self::scoll::UPDATED_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
index 11cf5e429..7df7893ce 100644
--- a/ergotree-interpreter/src/eval/scoll.rs
+++ b/ergotree-interpreter/src/eval/scoll.rs
@@ -5,8 +5,10 @@ use ergotree_ir::mir::constant::TryExtractInto;
 use ergotree_ir::mir::value::CollKind;
 use ergotree_ir::mir::value::Value;
 use ergotree_ir::types::stuple::STuple;
+use ergotree_ir::types::stype::SType::SInt;
 
 use super::EvalFn;
+use std::convert::TryFrom;
 
 pub(crate) static INDEX_OF_EVAL_FN: EvalFn = |_env, _ctx, obj, args| {
     Ok(Value::Int({
@@ -122,6 +124,66 @@ pub(crate) static ZIP_EVAL_FN: EvalFn = |_env, _ctx, obj, args| {
     }
 };
 
+pub(crate) static INDICES_EVAL_FN: EvalFn = |_env, _ctx, obj, _args| {
+    let normalized_input_vals: Vec<Value> = match obj {
+        Value::Coll(coll) => Ok(coll.as_vec()),
+        _ => Err(EvalError::UnexpectedValue(format!(
+            "expected obj to be Value::Coll, got: {0:?}",
+            obj
+        ))),
+    }?;
+    let indices_i32 = normalized_input_vals
+        .into_iter()
+        .enumerate()
+        .map(|(i, _)| i32::try_from(i))
+        .collect::<Result<Vec<i32>, _>>();
+    let indices_val =
+        indices_i32.map(|vec_i32| vec_i32.into_iter().map(Value::Int).collect::<Vec<Value>>());
+    match indices_val {
+        Ok(vec_val) => match CollKind::from_vec(SInt, vec_val) {
+            Ok(coll) => Ok(Value::Coll(coll)),
+            Err(e) => Err(EvalError::TryExtractFrom(e)),
+        },
+        Err(e) => Err(EvalError::UnexpectedValue(format!(
+            "Coll length overflow: {0:?}",
+            e
+        ))),
+    }
+};
+
+pub(crate) static UPDATED_EVAL_FN: EvalFn = |_env, _ctx, obj, args| {
+    let (input_tpe, normalized_input_vals) = match obj {
+        Value::Coll(coll) => Ok((coll.elem_tpe().clone(), coll.as_vec())),
+        _ => Err(EvalError::UnexpectedValue(format!(
+            "expected obj to be Value::Coll, got: {0:?}",
+            obj
+        ))),
+    }?;
+
+    let target_index_val = args
+        .get(0)
+        .cloned()
+        .ok_or_else(|| EvalError::NotFound("updated: missing first arg (index)".to_string()))?;
+    let update_val = args
+        .get(1)
+        .cloned()
+        .ok_or_else(|| EvalError::NotFound("updated: missing second arg (update)".to_string()))?;
+
+    let target_index_usize = target_index_val.clone().try_extract_into::<i32>()? as usize;
+    let mut res = normalized_input_vals;
+
+    match res.get_mut(target_index_usize) {
+        Some(elem) => {
+            *elem = update_val;
+            Ok(Value::Coll(CollKind::from_vec(input_tpe, res)?))
+        }
+        None => Err(EvalError::UnexpectedValue(format!(
+            "updated: target index out of bounds, got: {:?}",
+            target_index_val
+        ))),
+    }
+};
+
 #[allow(clippy::unwrap_used)]
 #[cfg(test)]
 #[cfg(feature = "arbitrary")]
@@ -139,7 +201,7 @@ mod tests {
     use ergotree_ir::types::stype::SType;
     use ergotree_ir::types::stype_param::STypeVar;
 
-    use crate::eval::tests::eval_out_wo_ctx;
+    use crate::eval::tests::{eval_out_wo_ctx, try_eval_out_wo_ctx};
     use ergotree_ir::types::stype::SType::SBoolean;
 
     #[test]
@@ -352,4 +414,67 @@ mod tests {
         let res = eval_out_wo_ctx::<Vec<(i64, bool)>>(&expr);
         assert_eq!(res, vec![(1i64, true), (2, false)]);
     }
+
+    #[test]
+    fn eval_indices() {
+        let coll_const: Constant = vec![1i64, 2i64, 3i64].into();
+        let expr: Expr = MethodCall::new(
+            coll_const.into(),
+            scoll::INDICES_METHOD
+                .clone()
+                .with_concrete_types(&[(STypeVar::t(), SType::SLong)].iter().cloned().collect()),
+            vec![],
+        )
+        .unwrap()
+        .into();
+        let res = eval_out_wo_ctx::<Vec<i32>>(&expr);
+        assert_eq!(res, vec![0i32, 1i32, 2i32]);
+    }
+
+    #[test]
+    fn eval_indices_empty_coll() {
+        let coll_const: Constant = Vec::<i64>::new().into();
+        let expr: Expr = MethodCall::new(
+            coll_const.into(),
+            scoll::INDICES_METHOD
+                .clone()
+                .with_concrete_types(&[(STypeVar::t(), SType::SLong)].iter().cloned().collect()),
+            vec![],
+        )
+        .unwrap()
+        .into();
+        let res = eval_out_wo_ctx::<Vec<i32>>(&expr);
+        assert_eq!(res, vec![]);
+    }
+
+    #[test]
+    fn eval_update() {
+        let coll_const: Constant = vec![1i64, 2i64, 3i64].into();
+        let expr: Expr = MethodCall::new(
+            coll_const.into(),
+            scoll::UPDATED_METHOD
+                .clone()
+                .with_concrete_types(&[(STypeVar::t(), SType::SLong)].iter().cloned().collect()),
+            vec![1i32.into(), 5i64.into()],
+        )
+        .unwrap()
+        .into();
+        let res = eval_out_wo_ctx::<Vec<i64>>(&expr);
+        assert_eq!(res, vec![1i64, 5i64, 3i64]);
+    }
+
+    #[test]
+    fn eval_update_oob() {
+        let coll_const: Constant = vec![1i64, 2i64, 3i64].into();
+        let expr: Expr = MethodCall::new(
+            coll_const.into(),
+            scoll::UPDATED_METHOD
+                .clone()
+                .with_concrete_types(&[(STypeVar::t(), SType::SLong)].iter().cloned().collect()),
+            vec![5i32.into(), 5i64.into()],
+        )
+        .unwrap()
+        .into();
+        assert!(try_eval_out_wo_ctx::<Vec<i64>>(&expr).is_err());
+    }
 }
diff --git a/ergotree-ir/src/types/scoll.rs b/ergotree-ir/src/types/scoll.rs
index ed48ee684..8eeffa485 100644
--- a/ergotree-ir/src/types/scoll.rs
+++ b/ergotree-ir/src/types/scoll.rs
@@ -18,7 +18,11 @@ pub const INDEX_OF_METHOD_ID: MethodId = MethodId(26);
 /// Coll.flatmap
 pub const FLATMAP_METHOD_ID: MethodId = MethodId(15);
 /// Coll.zip
-pub const ZIP_METHOD_ID: MethodId = MethodId(14);
+pub const ZIP_METHOD_ID: MethodId = MethodId(29);
+/// Coll.indices
+pub const INDICES_METHOD_ID: MethodId = MethodId(14);
+/// Coll.updated
+pub const UPDATED_METHOD_ID: MethodId = MethodId(20);
 
 static S_COLL_TYPE_COMPANION_HEAD: STypeCompanionHead = STypeCompanionHead {
     type_id: TYPE_ID,
@@ -32,7 +36,9 @@ lazy_static! {
         vec![
             &INDEX_OF_METHOD_DESC,
             &FLATMAP_METHOD_DESC,
-            &ZIP_METHOD_DESC
+            &ZIP_METHOD_DESC,
+            &INDICES_METHOD_DESC,
+            &UPDATED_METHOD_DESC,
         ]
     );
 }
@@ -42,7 +48,10 @@ lazy_static! {
         method_id: INDEX_OF_METHOD_ID,
         name: "indexOf",
         tpe: SFunc {
-            t_dom: vec![SType::SColl(SType::STypeVar(STypeVar::t()).into()), STypeVar::t().into(), SType::SInt],
+            t_dom: vec![
+                SType::SColl(SType::STypeVar(STypeVar::t()).into()),
+                STypeVar::t().into(),
+                SType::SInt],
             t_range: SType::SInt.into(),
             tpe_params: vec![],
         },
@@ -88,6 +97,39 @@ lazy_static! {
     pub static ref ZIP_METHOD: SMethod = SMethod::new(&S_COLL_TYPE_COMPANION, ZIP_METHOD_DESC.clone());
 }
 
+lazy_static! {
+    static ref INDICES_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: INDICES_METHOD_ID,
+        name: "indices",
+        tpe: SFunc::new(
+            vec![
+                SType::SColl(SType::STypeVar(STypeVar::t()).into()),
+            ],
+            SType::SColl(SType::SInt.into())
+        )
+    };
+    /// Coll.indices
+    pub static ref INDICES_METHOD: SMethod = SMethod::new(&S_COLL_TYPE_COMPANION, INDICES_METHOD_DESC.clone());
+}
+
+lazy_static! {
+    static ref UPDATED_METHOD_DESC: SMethodDesc = SMethodDesc {
+        method_id: UPDATED_METHOD_ID,
+        name: "updated",
+        tpe: SFunc::new(
+            vec![
+                SType::SColl(SType::STypeVar(STypeVar::t()).into()),
+                SType::SInt,
+                SType::STypeVar(STypeVar::t())
+
+            ],
+            SType::SColl(SType::STypeVar(STypeVar::t()).into())
+        )
+    };
+    /// Coll.updated
+    pub static ref UPDATED_METHOD: SMethod = SMethod::new(&S_COLL_TYPE_COMPANION, UPDATED_METHOD_DESC.clone());
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -97,5 +139,7 @@ mod tests {
         assert!(SMethod::from_ids(TYPE_ID, INDEX_OF_METHOD_ID).map(|e| e.name()) == Ok("indexOf"));
         assert!(SMethod::from_ids(TYPE_ID, FLATMAP_METHOD_ID).map(|e| e.name()) == Ok("flatMap"));
         assert!(SMethod::from_ids(TYPE_ID, ZIP_METHOD_ID).map(|e| e.name()) == Ok("zip"));
+        assert!(SMethod::from_ids(TYPE_ID, INDICES_METHOD_ID).map(|e| e.name()) == Ok("indices"));
+        assert!(SMethod::from_ids(TYPE_ID, UPDATED_METHOD_ID).map(|e| e.name()) == Ok("updated"));
     }
 }


Test patch: