feat: remove duplicated array reads at constant indices

compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs

@@ -284,7 +284,7 @@
     #[tracing::instrument(level = "trace", skip_all)]
     pub(crate) fn into_acir(self, brillig: &Brillig) -> Result<Artifacts, RuntimeError> {
         let mut acirs = Vec::new();
         // TODO: can we parallelise this?
         let mut shared_context = SharedContext::default();
         for function in self.functions.values() {
             let context = Context::new(&mut shared_context);
@@ -976,16 +976,23 @@
             }
         };
 
-        if self.handle_constant_index(instruction, dfg, index, array, store_value)? {
+        let array_id = dfg.resolve(array);
+        let array_typ = dfg.type_of_value(array_id);
+        // Compiler sanity checks
+        assert!(!array_typ.is_nested_slice(), "ICE: Nested slice type has reached ACIR generation");
+        let (Type::Array(_, _) | Type::Slice(_)) = &array_typ else {
+            unreachable!("ICE: expected array or slice type");
+        };
+
+        if self.handle_constant_index_wrapper(instruction, dfg, array, index, store_value)? {
             return Ok(());
         }
 
         // Get an offset such that the type of the array at the offset is the same as the type at the 'index'
         // If we find one, we will use it when computing the index under the enable_side_effect predicate
         // If not, array_get(..) will use a fallback costing one multiplication in the worst case.
         // cf. https://github.com/noir-lang/noir/pull/4971
-        let array_id = dfg.resolve(array);
-        let array_typ = dfg.type_of_value(array_id);
+
         // For simplicity we compute the offset only for simple arrays
         let is_simple_array = dfg.instruction_results(instruction).len() == 1
             && can_omit_element_sizes_array(&array_typ);
@@ -1018,83 +1025,106 @@
         Ok(())
     }
 
-    /// Handle constant index: if there is no predicate and we have the array values,
-    /// we can perform the operation directly on the array
-    fn handle_constant_index(
+    fn handle_constant_index_wrapper(
         &mut self,
         instruction: InstructionId,
         dfg: &DataFlowGraph,
+        array: ValueId,
         index: ValueId,
-        array_id: ValueId,
         store_value: Option<ValueId>,
     ) -> Result<bool, RuntimeError> {
-        let index_const = dfg.get_numeric_constant(index);
-        let value_type = dfg.type_of_value(array_id);
+        let array_id = dfg.resolve(array);
+        let array_typ = dfg.type_of_value(array_id);
         // Compiler sanity checks
-        assert!(
-            !value_type.is_nested_slice(),
-            "ICE: Nested slice type has reached ACIR generation"
-        );
-        let (Type::Array(_, _) | Type::Slice(_)) = &value_type else {
+        assert!(!array_typ.is_nested_slice(), "ICE: Nested slice type has reached ACIR generation");
+        let (Type::Array(_, _) | Type::Slice(_)) = &array_typ else {
             unreachable!("ICE: expected array or slice type");
         };
 
         match self.convert_value(array_id, dfg) {
             AcirValue::Var(acir_var, _) => {
-                return Err(RuntimeError::InternalError(InternalError::Unexpected {
+                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,
-                            });
-                        }
-                    };
+                // `AcirValue::Array` supports reading/writing to constant indices at compile-time in some cases.
+                if let Some(constant_index) = dfg.get_numeric_constant(index) {
+                    let store_value = store_value.map(|value| self.convert_value(value, dfg));
+                    self.handle_constant_index(instruction, dfg, array, constant_index, store_value)
+                } else {
+                    Ok(false)
+                }
+            }
+            AcirValue::DynamicArray(_) => Ok(false),
+        }
+    }
 
-                    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(),
-                            };
+    /// Handle constant index: if there is no predicate and we have the array values,
+    /// we can perform the operation directly on the array
+    fn handle_constant_index(
+        &mut self,
+        instruction: InstructionId,
+        dfg: &DataFlowGraph,
+        array: Vector<AcirValue>,
+        index: FieldElement,
+        store_value: Option<AcirValue>,
+    ) -> Result<bool, RuntimeError> {
+        let array_size: usize = array.len();
+        let index = match index.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,
+                });
+            }
+        };
 
-                            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);
-                    }
+        let side_effects_always_enabled =
+            self.acir_context.is_constant_one(&self.current_side_effects_enabled_var);
+        let index_out_of_bounds = index >= array_size;
+
+        // Note that the value of `side_effects_always_enabled` doesn't affect the value which we return here for valid
+        // indices, just whether we return an error for invalid indices at compile time or defer until execution.
+        match (side_effects_always_enabled, index_out_of_bounds) {
+            (true, false) => {
+                let value = match store_value {
+                    Some(store_value) => AcirValue::Array(array.update(index, store_value)),
+                    None => array[index].clone(),
+                };
+
+                self.define_result(dfg, instruction, value);
+                Ok(true)
+            }
+            (false, false) => {
+                if store_value.is_none() {
+                    // If there is a predicate and the index is not out of range, we can optimistically perform the
+                    // read at compile time as if the predicate is true.
+                    //
+                    // This is as if the predicate is false, any side-effects will be disabled so the value returned
+                    // will not affect the rest of execution.
+                    self.define_result(dfg, instruction, array[index].clone());
+                    Ok(true)
+                } else {
+                    // We do not do this for a array writes however.
+                    Ok(false)
                 }
             }
-            AcirValue::DynamicArray(_) => (),
-        };
 
-        Ok(false)
+            // Report the error if side effects are enabled.
+            (true, true) => {
+                let call_stack = self.acir_context.get_call_stack();
+                Err(RuntimeError::IndexOutOfBounds { index, array_size, call_stack })
+            }
+            // Index is out of bounds but predicate may result in this array operation being skipped
+            // so we don't return an error now.
+            (false, true) => Ok(false),
+        }
     }
 
     /// We need to properly setup the inputs for array operations in ACIR.

compiler/noirc_evaluator/src/ssa/ir/instruction.rs

@@ -379,9 +379,17 @@ impl Instruction {
             {
                 true
             }
-            Instruction::EnableSideEffects { .. }
-            | Instruction::ArrayGet { .. }
-            | Instruction::ArraySet { .. } => true,
+
+            // `ArrayGet`s which read from "known good" indices from an array don't need a predicate.
+            Instruction::ArrayGet { array, index } => {
+                #[allow(clippy::match_like_matches_macro)]
+                match (dfg.type_of_value(*array), dfg.get_numeric_constant(*index)) {
+                    (Type::Array(_, len), Some(index)) if index.to_u128() < (len as u128) => false,
+                    _ => true,
+                }
+            }
+
+            Instruction::EnableSideEffects { .. } | Instruction::ArraySet { .. } => true,
 
             Instruction::Call { func, .. } => match dfg[*func] {
                 Value::Function(_) => true,

compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs

@@ -288,7 +288,7 @@ mod test {
             value::{Value, ValueId},
         },
     };
-    use acvm::acir::AcirField;
+    use acvm::{acir::AcirField, FieldElement};
 
     #[test]
     fn simple_constant_fold() {
@@ -545,6 +545,73 @@ mod test {
         assert_eq!(instruction, &Instruction::Cast(v0, Type::unsigned(32)));
     }
 
+    #[test]
+    fn constant_index_array_access_deduplication() {
+        // fn main f0 {
+        //   b0(v0: [Field; 4], v1: u32, v2: bool, v3: bool):
+        //     enable_side_effects v2
+        //     v4 = array_get v0 u32 0
+        //     v5 = array_get v0 v1
+        //     enable_side_effects v3
+        //     v6 = array_get v0 u32 0
+        //     v7 = array_get v0 v1
+        //     constrain v4 v6
+        // }
+        //
+        // After constructing this IR, we run constant folding which should replace the second constant-index array get
+        // with a reference to the results to the first. This then allows us to optimize away
+        // the constrain instruction as both inputs are known to be equal.
+        //
+        let main_id = Id::test_new(0);
+
+        // Compiling main
+        let mut builder = FunctionBuilder::new("main".into(), main_id);
+
+        let v0 = builder.add_parameter(Type::Array(Rc::new(vec![Type::field()]), 4));
+        let v1 = builder.add_parameter(Type::unsigned(32));
+        let v2 = builder.add_parameter(Type::unsigned(1));
+        let v3 = builder.add_parameter(Type::unsigned(1));
+
+        let zero = builder.numeric_constant(FieldElement::zero(), Type::length_type());
+
+        builder.insert_enable_side_effects_if(v2);
+        let v4 = builder.insert_array_get(v0, zero, Type::field());
+        let _v5 = builder.insert_array_get(v0, v1, Type::field());
+
+        builder.insert_enable_side_effects_if(v3);
+        let v6 = builder.insert_array_get(v0, zero, Type::field());
+        let _v7 = builder.insert_array_get(v0, v1, Type::field());
+
+        builder.insert_constrain(v4, v6, None);
+
+        let ssa = builder.finish();
+
+        println!("{ssa}");
+
+        let main = ssa.main();
+        let instructions = main.dfg[main.entry_block()].instructions();
+        assert_eq!(instructions.len(), 7);
+
+        // Expected output:
+        //
+        // fn main f0 {
+        //   b0(v0: [Field; 4], v1: u32, v2: bool, v3: bool):
+        //     enable_side_effects v2
+        //     v10 = array_get v0 u32 0
+        //     v11 = array_get v0 v1
+        //     enable_side_effects v3
+        //     v12 = array_get v0 v1
+        // }
+        let ssa = ssa.fold_constants();
+
+        println!("{ssa}");
+
+        let main = ssa.main();
+        let instructions = main.dfg[main.entry_block()].instructions();
+
+        assert_eq!(instructions.len(), 5);
+    }
+
     #[test]
     fn constraint_decomposition() {
         // fn main f0 {