feat: handle constant index operations on simple slices (#3464)

Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>
noir-lang · Nov 14, 2023 · 22e1e89 · 22e1e89
1 parent 259f04f
commit 22e1e89
Showing 1 changed file with 62 additions and 61 deletions.
diff --git a/compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs b/compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs
@@ -658,75 +658,76 @@ impl Context {
         store_value: Option<ValueId>,
     ) -> Result<bool, RuntimeError> {
         let index_const = dfg.get_numeric_constant(index);
-        match dfg.type_of_value(array) {
-            Type::Array(_, _) => {
-                match self.convert_value(array, dfg) {
-                    AcirValue::Var(acir_var, _) => {
-                        return Err(RuntimeError::InternalError(InternalError::UnExpected {
-                            expected: "an array value".to_string(),
-                            found: format!("{acir_var:?}"),
-                            call_stack: self.acir_context.get_call_stack(),
-                        }))
-                    }
-                    AcirValue::Array(array) => {
-                        if let Some(index_const) = index_const {
-                            let array_size = array.len();
-                            let index = match index_const.try_to_u64() {
-                                Some(index_const) => index_const as usize,
-                                None => {
-                                    let call_stack = self.acir_context.get_call_stack();
-                                    return Err(RuntimeError::TypeConversion {
-                                        from: "array index".to_string(),
-                                        into: "u64".to_string(),
-                                        call_stack,
-                                    });
+        let value_type = dfg.type_of_value(array);
+        let (Type::Array(element_types, _) | Type::Slice(element_types)) = &value_type else {
+            unreachable!("ICE: expected array or slice type");
+
+        };
+
+        // TODO(#3188): Need to be able to handle constant index for slices to seriously reduce
+        // constraint sizes of nested slices
+        // This can only be done if we accurately flatten nested slices as otherwise we will reach
+        // index out of bounds errors. If the slice is already flat then we can treat them similarly to arrays.
+        if matches!(value_type, Type::Slice(_))
+            && element_types.iter().any(|element| element.contains_slice_element())
+        {
+            return Ok(false);
+        }
+
+        match self.convert_value(array, dfg) {
+            AcirValue::Var(acir_var, _) => {
+                return Err(RuntimeError::InternalError(InternalError::UnExpected {
+                    expected: "an array value".to_string(),
+                    found: format!("{acir_var:?}"),
+                    call_stack: self.acir_context.get_call_stack(),
+                }))
+            }
+            AcirValue::Array(array) => {
+                if let Some(index_const) = index_const {
+                    let array_size = array.len();
+                    let index = match index_const.try_to_u64() {
+                        Some(index_const) => index_const as usize,
+                        None => {
+                            let call_stack = self.acir_context.get_call_stack();
+                            return Err(RuntimeError::TypeConversion {
+                                from: "array index".to_string(),
+                                into: "u64".to_string(),
+                                call_stack,
+                            });
+                        }
+                    };
+                    if self.acir_context.is_constant_one(&self.current_side_effects_enabled_var) {
+                        // Report the error if side effects are enabled.
+                        if index >= array_size {
+                            let call_stack = self.acir_context.get_call_stack();
+                            return Err(RuntimeError::IndexOutOfBounds {
+                                index,
+                                array_size,
+                                call_stack,
+                            });
+                        } else {
+                            let value = match store_value {
+                                Some(store_value) => {
+                                    let store_value = self.convert_value(store_value, dfg);
+                                    AcirValue::Array(array.update(index, store_value))
                                 }
+                                None => array[index].clone(),
                             };
-                            if self
-                                .acir_context
-                                .is_constant_one(&self.current_side_effects_enabled_var)
-                            {
-                                // Report the error if side effects are enabled.
-                                if index >= array_size {
-                                    let call_stack = self.acir_context.get_call_stack();
-                                    return Err(RuntimeError::IndexOutOfBounds {
-                                        index,
-                                        array_size,
-                                        call_stack,
-                                    });
-                                } else {
-                                    let value = match store_value {
-                                        Some(store_value) => {
-                                            let store_value = self.convert_value(store_value, dfg);
-                                            AcirValue::Array(array.update(index, store_value))
-                                        }
-                                        None => array[index].clone(),
-                                    };
 
-                                    self.define_result(dfg, instruction, value);
-                                    return Ok(true);
-                                }
-                            }
-                            // If there is a predicate and the index is not out of range, we can directly perform the read
-                            else if index < array_size && store_value.is_none() {
-                                self.define_result(dfg, instruction, array[index].clone());
-                                return Ok(true);
-                            }
+                            self.define_result(dfg, instruction, value);
+                            return Ok(true);
                         }
                     }
-                    AcirValue::DynamicArray(_) => (),
+                    // If there is a predicate and the index is not out of range, we can directly perform the read
+                    else if index < array_size && store_value.is_none() {
+                        self.define_result(dfg, instruction, array[index].clone());
+                        return Ok(true);
+                    }
                 }
             }
-            Type::Slice(_) => {
-                // TODO(#3188): Need to be able to handle constant index for slices to seriously reduce
-                // constraint sizes of nested slices
-                // This can only be done if we accurately flatten nested slices as other we will reach
-                // index out of bounds errors.
+            AcirValue::DynamicArray(_) => (),
+        };
 
-                // Do nothing we only want dynamic checks for slices
-            }
-            _ => unreachable!("ICE: expected array or slice type"),
-        }
         Ok(false)
     }