From 25767ecb7b9370c7067934aa5760585b82fc8ece Mon Sep 17 00:00:00 2001
From: Innokenty <innokenty@singularitynet.io>
Date: Mon, 28 Oct 2024 13:28:33 +0300
Subject: [PATCH] index-atom final version, documentation and tests added

---
 lib/src/metta/runner/arithmetics.rs       |  2 +-
 lib/src/metta/runner/stdlib.metta         |  7 +++++
 lib/src/metta/runner/stdlib.rs            | 37 +++++++++++++++++++++++
 lib/src/metta/runner/stdlib_minimal.metta |  7 +++++
 lib/src/metta/runner/stdlib_minimal.rs    |  8 +++++
 5 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/lib/src/metta/runner/arithmetics.rs b/lib/src/metta/runner/arithmetics.rs
index f76e45de3..bc8374b44 100644
--- a/lib/src/metta/runner/arithmetics.rs
+++ b/lib/src/metta/runner/arithmetics.rs
@@ -343,7 +343,7 @@ impl serial::Serializer for NumberSerializer {
     }
 }
 
-struct AsPrimitive<'a> {
+pub struct AsPrimitive<'a> {
     atom: &'a super::Atom
 }
 
diff --git a/lib/src/metta/runner/stdlib.metta b/lib/src/metta/runner/stdlib.metta
index 34b90e9b4..2139af922 100644
--- a/lib/src/metta/runner/stdlib.metta
+++ b/lib/src/metta/runner/stdlib.metta
@@ -415,6 +415,13 @@
     (@param "Tail of an expression")))
   (@return "New expression consists of two input arguments"))
 
+(@doc index-atom
+  (@desc "Returns atom from an expression (first argument) using index (second argument) or error if index is out of bounds")
+  (@params (
+    (@param "Expression")
+    (@param "Index")))
+  (@return "Atom from an expression in the place defined by index. Error if index is out of bounds of an expression"))
+
 (@doc println!
   (@desc "Prints a line of text to the console")
   (@params (
diff --git a/lib/src/metta/runner/stdlib.rs b/lib/src/metta/runner/stdlib.rs
index 47e26de37..0aee17415 100644
--- a/lib/src/metta/runner/stdlib.rs
+++ b/lib/src/metta/runner/stdlib.rs
@@ -1160,6 +1160,33 @@ impl CustomExecute for IntersectionAtomOp {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct IndexAtomOp {}
+
+grounded_op!(IndexAtomOp, "index-atom");
+
+impl Grounded for IndexAtomOp {
+    fn type_(&self) -> Atom {
+        Atom::expr([ARROW_SYMBOL, ATOM_TYPE_EXPRESSION, ATOM_TYPE_NUMBER, ATOM_TYPE_ATOM])
+    }
+
+    fn as_execute(&self) -> Option<&dyn CustomExecute> {
+        Some(self)
+    }
+}
+
+impl CustomExecute for IndexAtomOp {
+    fn execute(&self, args: &[Atom]) -> Result<Vec<Atom>, ExecError> {
+        let arg_error = || ExecError::from("index-atom expects two arguments: expression and atom");
+        let children = TryInto::<&ExpressionAtom>::try_into(args.get(0).ok_or_else(arg_error)?)?.children();
+        let index = AsPrimitive::from_atom(args.get(1).ok_or_else(arg_error)?).as_number().ok_or_else(arg_error)?;
+        match children.get(Into::<i64>::into(index) as usize) {
+            Some(atom) => Ok(vec![atom.clone()]),
+            None => Err(ExecError::from("Index is out of bounds")),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct SubtractionAtomOp {}
 
@@ -1737,6 +1764,8 @@ mod non_minimal_only_stdlib {
         tref.register_token(regex(r"cdr-atom"), move |_| { cdr_atom_op.clone() });
         let cons_atom_op = Atom::gnd(ConsAtomOp{});
         tref.register_token(regex(r"cons-atom"), move |_| { cons_atom_op.clone() });
+        let index_atom_op = Atom::gnd(IndexAtomOp{});
+        tref.register_token(regex(r"index-atom"), move |_| { index_atom_op.clone() });
         let println_op = Atom::gnd(PrintlnOp{});
         tref.register_token(regex(r"println!"), move |_| { println_op.clone() });
         let format_args_op = Atom::gnd(FormatArgsOp{});
@@ -2024,6 +2053,14 @@ mod tests {
         assert_eq!(res, vec![expr!(("A" "F") ("B" "C") "D")]);
     }
 
+    #[test]
+    fn index_atom_op() {
+        let res = IndexAtomOp{}.execute(&mut vec![expr!({Number::Integer(5)} {Number::Integer(4)} {Number::Integer(3)} {Number::Integer(2)} {Number::Integer(1)}), expr!({Number::Integer(2)})]).expect("No result returned");
+        assert_eq!(res, vec![expr!({Number::Integer(3)})]);
+        let res = IndexAtomOp{}.execute(&mut vec![expr!({Number::Integer(5)} {Number::Integer(4)} {Number::Integer(3)} {Number::Integer(2)} {Number::Integer(1)}), expr!({Number::Integer(5)})]);
+        assert_eq!(res, Err(ExecError::from("Index is out of bounds")));
+    }
+
     #[test]
     fn bind_new_space_op() {
         let tokenizer = Shared::new(Tokenizer::new());
diff --git a/lib/src/metta/runner/stdlib_minimal.metta b/lib/src/metta/runner/stdlib_minimal.metta
index 5f053be70..e4d43d89a 100644
--- a/lib/src/metta/runner/stdlib_minimal.metta
+++ b/lib/src/metta/runner/stdlib_minimal.metta
@@ -72,6 +72,13 @@
   (@return "Deconsed expression"))
 (: decons-atom (-> Expression Expression))
 
+(@doc index-atom
+  (@desc "Returns atom from an expression (first argument) using index (second argument) or error if index is out of bounds")
+  (@params (
+    (@param "Expression")
+    (@param "Index")))
+  (@return "Atom from an expression in the place defined by index. Error if index is out of bounds"))
+
 (@doc collapse-bind
   (@desc "Evaluates minimal MeTTa operation (first argument) and returns an expression which contains all alternative evaluations in a form (Atom Bindings). Bindings are represented in a form of a grounded atom.")
   (@params (
diff --git a/lib/src/metta/runner/stdlib_minimal.rs b/lib/src/metta/runner/stdlib_minimal.rs
index 927a33bb3..ff81ba839 100644
--- a/lib/src/metta/runner/stdlib_minimal.rs
+++ b/lib/src/metta/runner/stdlib_minimal.rs
@@ -432,6 +432,8 @@ pub fn register_common_tokens(tref: &mut Tokenizer, _tokenizer: Shared<Tokenizer
     tref.register_token(regex(r"nop"), move |_| { nop_op.clone() });
     let match_op = Atom::gnd(stdlib::MatchOp{});
     tref.register_token(regex(r"match"), move |_| { match_op.clone() });
+    let index_atom_op = Atom::gnd(stdlib::IndexAtomOp{});
+    tref.register_token(regex(r"index-atom"), move |_| { index_atom_op.clone() });
     let mod_space_op = Atom::gnd(stdlib::ModSpaceOp::new(metta.clone()));
     tref.register_token(regex(r"mod-space!"), move |_| { mod_space_op.clone() });
     let print_mods_op = Atom::gnd(stdlib::PrintModsOp::new(metta.clone()));
@@ -615,6 +617,12 @@ mod tests {
         assert_eq!(run_program(&format!("!(cdr-atom $a)")), Ok(vec![vec![expr!("Error" ("cdr-atom" a) {Str::from_str("cdr-atom expects a non-empty expression as an argument")})]]));
     }
 
+    #[test]
+    fn metta_index_atom() {
+        assert_eq!(run_program(&format!("!(index-atom (5 4 3 2 1) 2)")), Ok(vec![vec![expr!({Number::Integer(3)})]]));
+        assert_eq!(run_program(&format!("!(index-atom (A B C D E) 5)")), Ok(vec![vec![expr!("Error" ({ stdlib::IndexAtomOp{} } ("A" "B" "C" "D" "E") {Number::Integer(5)}) "Index is out of bounds")]]));
+    }
+
     #[test]
     fn metta_switch() {
         let result = run_program("!(eval (switch (A $b) ( (($a B) ($b $a)) ((B C) (C B)) )))");