Add indexing operator for expressions

The <nth> operator extracts the nth integer from a list.

2 files changed, 24 insertions, 0 deletions

diff --git a/src/types/expression_stack.rs b/src/types/expression_stack.rs
index 3f80dc8..f197ea7 100644
--- a/src/types/expression_stack.rs
+++ b/src/types/expression_stack.rs
@@ -68,6 +68,7 @@ impl ExpressionStack {
             Operator::BitXor           => { pop!(b); pop!(a); push!(op_bit_xor(a, b)) },
             Operator::BitNot           => {          pop!(a); push!(op_bit_not(a)) },
             Operator::Length           => {          pop!(a); push!(op_length(a)) },
+            Operator::Index            => { pop!(b); pop!(a); push!(op_index(a, b)) },
             Operator::Sum              => {          pop!(a); push!(op_sum(a)) },
             Operator::Absolute         => {          pop!(a); push!(op_absolute(a)) },
             Operator::Debug            => {          pop!(a); op_debug(&a, &source); push!(Ok(a)) },
@@ -99,6 +100,15 @@ fn to_isize(value: IntermediateValue) -> Result<isize, ExpressionError> {
     Err(ExpressionError::InvalidArgumentType("an integer", received))
 }
 
+fn to_list(value: IntermediateValue) -> Result<Vec<isize>, ExpressionError> {
+    let received = match value {
+        IntermediateValue::List(list) => return Ok(list.into_iter().map(|t| t.value).collect()),
+        IntermediateValue::Integer(_) => "an integer",
+        IntermediateValue::Block(_) => "a block",
+    };
+    Err(ExpressionError::InvalidArgumentType("a list", received))
+}
+
 fn from_isize(value: isize) -> IntermediateValue {
     IntermediateValue::Integer(Tracked::from(value, null_span()))
 }
@@ -175,6 +185,13 @@ fn op_length(l: IntermediateValue) -> Result<IntermediateValue, ExpressionError>
     };
     Ok(from_isize(length))
 }
+fn op_index(l: IntermediateValue, r: IntermediateValue) -> Result<IntermediateValue, ExpressionError> {
+    let l = to_list(l)?; let r = to_isize(r)?;
+    match r >= 0 && r < l.len() as isize {
+        true => Ok(from_isize(l[r as usize])),
+        false => Err(ExpressionError::IndexError(l.len(), r)),
+    }
+}
 fn op_sum(l: IntermediateValue) -> Result<IntermediateValue, ExpressionError> {
     let sum = match l {
         IntermediateValue::Integer(integer) => ones(*integer) as isize,
@@ -224,6 +241,8 @@ pub enum ExpressionError {
     InvalidReturnType(&'static str),
     // (expected, received)
     InvalidArgumentType(&'static str, &'static str),
+    // (length, index)
+    IndexError(usize, isize),
 }
 
 
@@ -240,6 +259,8 @@ pub fn report_expression_error(error: &Tracked<ExpressionError>, source_code: &s
             &format!("Expression must return an integer value, not {received} value"),
         ExpressionError::InvalidArgumentType(expected, received) =>
             &format!("Operator expected {expected} value, not {received} value"),
+        ExpressionError::IndexError(length, index) =>
+            &format!("Could not access element {index} from a list of length {length}"),
     };
 
     report_source_issue(LogLevel::Error, &context, message);
diff --git a/src/types/operator.rs b/src/types/operator.rs
index 8a76e52..5bf82b2 100644
--- a/src/types/operator.rs
+++ b/src/types/operator.rs
@@ -19,6 +19,7 @@ pub enum Operator {
     BitXor,
     BitNot,
     Length,
+    Index,
     Sum,
     Absolute,
     Debug,
@@ -60,6 +61,7 @@ impl Operator {
             "<xor>" => Some(Operator::BitXor),
             "<not>" => Some(Operator::BitNot),
             "<len>" => Some(Operator::Length),
+            "<nth>" => Some(Operator::Index),
             "<sum>" => Some(Operator::Sum),
             "<abs>" => Some(Operator::Absolute),
             "<dbg>" => Some(Operator::Debug),
@@ -90,6 +92,7 @@ impl std::fmt::Display for Operator {
             Operator::BitXor           => "<xor>",
             Operator::BitNot           => "<not>",
             Operator::Length           => "<len>",
+            Operator::Index            => "<nth>",
             Operator::Sum              => "<sum>",
             Operator::Absolute         => "<abs>",
             Operator::Debug            => "<dbg>",