Allow negative indexing (#2606)

This change relaxing the restriction on indexing into arrays to allow
negative indices. Previously the index limit was 0 to N-1, this updates
it such that the allowed range is -N to N-1, where negative indices
count from the end of the array.

Author: swernli

source/compiler/qsc_eval/src/lib.rs

@@ -1462,8 +1462,7 @@ impl State {
                 if index < 0 {
                     return Err(Error::InvalidNegativeInt(index, span));
                 }
-                let i = index.as_index(span)?;
-                self.update_array_index_single(env, globals, lhs, span, i, update)
+                self.update_array_index_single(env, globals, lhs, span, index, update)
             }
             range @ Value::Range(..) => {
                 self.update_array_index_range(env, globals, lhs, span, &range, update)
@@ -1598,16 +1597,14 @@ impl State {
         globals: &impl PackageStoreLookup,
         lhs: ExprId,
         span: PackageSpan,
-        index: usize,
+        index: i64,
         rhs: Value,
     ) -> Result<(), Error> {
         let lhs = globals.get_expr((self.package, lhs).into());
         match &lhs.kind {
             &ExprKind::Var(Res::Local(id), _) => match env.get_mut(id) {
                 Some(var) => {
-                    var.value.update_array(index, rhs).map_err(|idx| {
-                        Error::IndexOutOfRange(idx.try_into().expect("index should be valid"), span)
-                    })?;
+                    var.value.update_array(index, rhs, span)?;
                 }
                 None => return Err(Error::UnboundName(self.to_global_span(lhs.span))),
             },
@@ -1642,13 +1639,7 @@ impl State {
                         if idx < 0 {
                             return Err(Error::InvalidNegativeInt(idx, range_span));
                         }
-                        let i = idx.as_index(range_span)?;
-                        var.value.update_array(i, rhs.clone()).map_err(|idx| {
-                            Error::IndexOutOfRange(
-                                idx.try_into().expect("index should be valid"),
-                                range_span,
-                            )
-                        })?;
+                        var.value.update_array(idx, rhs.clone(), range_span)?;
                     }
                 }
                 None => return Err(Error::UnboundName(self.to_global_span(lhs.span))),

source/compiler/qsc_eval/src/tests.rs

@@ -1734,19 +1734,24 @@ fn array_slice_out_of_range_expr() {
 
 #[test]
 fn array_index_negative_expr() {
+    check_expr("", "[1, 2, 3][-2]", &expect!["2"]);
+}
+
+#[test]
+fn array_index_out_of_range_expr() {
     check_expr(
         "",
-        "[1, 2, 3][-2]",
+        "[1, 2, 3][4]",
         &expect![[r#"
-            InvalidIndex(
-                -2,
+            IndexOutOfRange(
+                4,
                 PackageSpan {
                     package: PackageId(
                         2,
                     ),
                     span: Span {
                         lo: 10,
-                        hi: 12,
+                        hi: 11,
                     },
                 },
             )
@@ -1755,13 +1760,13 @@ fn array_index_negative_expr() {
 }
 
 #[test]
-fn array_index_out_of_range_expr() {
+fn array_index_out_of_range_with_length_expr() {
     check_expr(
         "",
-        "[1, 2, 3][4]",
+        "[1, 2, 3][3]",
         &expect![[r#"
             IndexOutOfRange(
-                4,
+                3,
                 PackageSpan {
                     package: PackageId(
                         2,
@@ -1776,6 +1781,28 @@ fn array_index_out_of_range_expr() {
     );
 }
 
+#[test]
+fn array_index_out_of_range_with_negative_length_minus_one_expr() {
+    check_expr(
+        "",
+        "[1, 2, 3][-4]",
+        &expect![[r#"
+            IndexOutOfRange(
+                -4,
+                PackageSpan {
+                    package: PackageId(
+                        2,
+                    ),
+                    span: Span {
+                        lo: 10,
+                        hi: 12,
+                    },
+                },
+            )
+        "#]],
+    );
+}
+
 #[test]
 fn literal_big_int_expr() {
     check_expr(
@@ -2273,24 +2300,11 @@ fn update_array_with_range_out_of_range_err() {
 }
 
 #[test]
-fn update_array_with_range_negative_index_err() {
+fn update_array_with_range_negative_index() {
     check_expr(
         "",
         "[0, 1, 2, 3] w/ -1..0 <- [10, 11, 12, 13]",
-        &expect![[r#"
-            InvalidIndex(
-                -1,
-                PackageSpan {
-                    package: PackageId(
-                        2,
-                    ),
-                    span: Span {
-                        lo: 16,
-                        hi: 21,
-                    },
-                },
-            )
-        "#]],
+        &expect!["[11, 1, 2, 10]"],
     );
 }
 

source/compiler/qsc_eval/src/val.rs

@@ -275,17 +275,23 @@ impl Value {
     /// Updates a value in an array in-place.
     /// # Panics
     /// This will panic if the [Value] is not a [`Value::Array`].
-    pub fn update_array(&mut self, index: usize, value: Self) -> core::result::Result<(), usize> {
+    pub fn update_array(
+        &mut self,
+        index: i64,
+        value: Self,
+        span: PackageSpan,
+    ) -> core::result::Result<(), Error> {
         let Value::Array(arr) = self else {
             panic!("value should be Array, got {}", self.type_name());
         };
         let arr = Rc::get_mut(arr).expect("array should be uniquely referenced");
-        match arr.get_mut(index) {
+        let i = index_allowing_negative(arr.len(), index, span)?;
+        match arr.get_mut(i) {
             Some(v) => {
                 *v = value;
                 Ok(())
             }
-            None => Err(index),
+            None => Err(Error::IndexOutOfRange(index, span)),
         }
     }
 
@@ -479,13 +485,29 @@ pub fn index_array(
     index: i64,
     span: PackageSpan,
 ) -> std::result::Result<Value, Error> {
-    let i = index.as_index(span)?;
+    let i = index_allowing_negative(arr.len(), index, span)?;
     match arr.get(i) {
         Some(v) => Ok(v.clone()),
         None => Err(Error::IndexOutOfRange(index, span)),
     }
 }
 
+/// Converts an index to a usize, allowing for negative indices that count from the end of the array.
+/// Valid range is `-len` to `len - 1`, where `len` is the length of the array.
+pub fn index_allowing_negative(
+    len: usize,
+    index: i64,
+    span: PackageSpan,
+) -> std::result::Result<usize, Error> {
+    let i = if index >= 0 {
+        index.as_index(span)?
+    } else {
+        len.checked_sub((-index).as_index(span)?)
+            .ok_or(Error::IndexOutOfRange(index, span))?
+    };
+    Ok(i)
+}
+
 pub fn make_range(
     arr: &[Value],
     start: Option<i64>,
@@ -534,7 +556,7 @@ pub fn update_index_single(
     if index < 0 {
         return Err(Error::InvalidNegativeInt(index, span));
     }
-    let i = index.as_index(span)?;
+    let i = index_allowing_negative(values.len(), index, span)?;
     let mut values = values.to_vec();
     match values.get_mut(i) {
         Some(value) => {
@@ -557,7 +579,7 @@ pub fn update_index_range(
     let mut values = values.to_vec();
     let update = update.unwrap_array();
     for (idx, update) in range.into_iter().zip(update.iter()) {
-        let i = idx.as_index(span)?;
+        let i = index_allowing_negative(values.len(), idx, span)?;
         match values.get_mut(i) {
             Some(value) => {
                 *value = update.clone();

source/compiler/qsc_partial_eval/src/tests/arrays.rs

@@ -269,8 +269,8 @@ fn result_array_value_at_index() {
 }
 
 #[test]
-fn result_array_value_at_negative_index_raises_error() {
-    let error = get_partial_evaluation_error(indoc! {r#"
+fn result_array_value_at_negative_index_works() {
+    let program = get_rir_program(indoc! {r#"
         namespace Test {
             @EntryPoint()
             operation Main() : Result {
@@ -280,11 +280,15 @@ fn result_array_value_at_negative_index_raises_error() {
             }
         }
     "#});
-    assert_error(
-        &error,
-        &expect![[
-            r#"EvaluationFailed("value cannot be used as an index: -1", PackageSpan { package: PackageId(2), span: Span { lo: 177, hi: 179 } })"#
-        ]],
+    assert_block_instructions(
+        &program,
+        BlockId(0),
+        &expect![[r#"
+            Block:
+                Call id(1), args( Qubit(0), Result(0), )
+                Call id(1), args( Qubit(1), Result(1), )
+                Call id(2), args( Result(1), Pointer, )
+                Return"#]],
     );
 }