feat(optimization): Follow past array_sets when optimizing `array_g…

…et`s (#5772)

# Description

## Problem\*

Resolves #5767

## Summary\*

Instead of checking if an array_get is from a constant array, we can
follow array_sets to find a previous constant array. If the index was
ever unknown we fail to optimize. If the index is known but changed by
an array set we fail as well instead of optimizing to the set value.
This is in case the side effect variables are different for the get &
set.

## Additional Context

I set an arbitrary maximum of 5 array_sets to check before giving up and
failing to find a constant array.

## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: Tom French <[email protected]>

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

@@ -608,16 +608,11 @@ impl Instruction {
                 }
             }
             Instruction::ArrayGet { array, index } => {
-                let array = dfg.get_array_constant(*array);
-                let index = dfg.get_numeric_constant(*index);
-                if let (Some((array, _)), Some(index)) = (array, index) {
-                    let index =
-                        index.try_to_u32().expect("Expected array index to fit in u32") as usize;
-                    if index < array.len() {
-                        return SimplifiedTo(array[index]);
-                    }
+                if let Some(index) = dfg.get_numeric_constant(*index) {
+                    try_optimize_array_get_from_previous_set(dfg, *array, index)
+                } else {
+                    None
                 }
-                None
             }
             Instruction::ArraySet { array, index, value, .. } => {
                 let array = dfg.get_array_constant(*array);
@@ -744,6 +739,65 @@ impl Instruction {
     }
 }
 
+/// Given a chain of operations like:
+/// v1 = array_set [10, 11, 12], index 1, value: 5
+/// v2 = array_set v1, index 2, value: 6
+/// v3 = array_set v2, index 2, value: 7
+/// v4 = array_get v3, index 1
+///
+/// We want to optimize `v4` to `10`. To do this we need to follow the array value
+/// through several array sets. For each array set:
+/// - If the index is non-constant we fail the optimization since any index may be changed
+/// - If the index is constant and is our target index, we conservatively fail the optimization
+///   in case the array_set is disabled from a previous `enable_side_effects_if` and the array get
+///   was not.
+/// - Otherwise, we check the array value of the array set.
+///   - If the array value is constant, we use that array.
+///   - If the array value is from a previous array-set, we recur.
+fn try_optimize_array_get_from_previous_set(
+    dfg: &DataFlowGraph,
+    mut array_id: Id<Value>,
+    target_index: FieldElement,
+) -> SimplifyResult {
+    let mut elements = None;
+
+    // Arbitrary number of maximum tries just to prevent this optimization from taking too long.
+    let max_tries = 5;
+    for _ in 0..max_tries {
+        match &dfg[array_id] {
+            Value::Instruction { instruction, .. } => {
+                match &dfg[*instruction] {
+                    Instruction::ArraySet { array, index, value, .. } => {
+                        if let Some(constant) = dfg.get_numeric_constant(*index) {
+                            if constant == target_index {
+                                return SimplifyResult::SimplifiedTo(*value);
+                            }
+
+                            array_id = *array; // recur
+                        } else {
+                            return SimplifyResult::None;
+                        }
+                    }
+                    _ => return SimplifyResult::None,
+                }
+            }
+            Value::Array { array, typ: _ } => {
+                elements = Some(array.clone());
+                break;
+            }
+            _ => return SimplifyResult::None,
+        }
+    }
+
+    if let (Some(array), Some(index)) = (elements, target_index.try_to_u64()) {
+        let index = index as usize;
+        if index < array.len() {
+            return SimplifyResult::SimplifiedTo(array[index]);
+        }
+    }
+    SimplifyResult::None
+}
+
 pub(crate) type ErrorType = HirType;
 
 pub(crate) fn error_selector_from_type(typ: &ErrorType) -> ErrorSelector {