feat(SSA): simplify array_get on slice_insert

compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs

@@ -2,7 +2,7 @@ use crate::ssa::ir::{
     basic_block::BasicBlockId,
     dfg::simplify::value_merger::ValueMerger,
     instruction::{
-        Binary, BinaryOp, ConstrainError, Instruction,
+        Binary, BinaryOp, ConstrainError, Instruction, Intrinsic,
         binary::{truncate, truncate_field},
     },
     types::{NumericType, Type},
@@ -111,20 +111,31 @@ pub(crate) fn simplify(
         }
         Instruction::ConstrainNotEqual(..) => None,
         Instruction::ArrayGet { array, index } => {
-            if let Some(index) = dfg.get_numeric_constant(*index) {
-                return try_optimize_array_get_from_previous_set(dfg, *array, index);
+            let array = *array;
+            let index = *index;
+            let constant_index = dfg.get_numeric_constant(index);
+
+            // If it's an `array_get` on a value that results from a call to `slice_insert` with a
+            // constant index we can do `array_get` on the original slice or use the inserted value
+            let result = try_optimize_array_get_on_slice_insert(dfg, array, index, constant_index);
+            if !matches!(result, None) {
+                return result;
             }
 
-            let array_or_slice_type = dfg.type_of_value(*array);
+            if let Some(constant_index) = constant_index {
+                return try_optimize_array_get_from_previous_set(dfg, array, index, constant_index);
+            }
+
+            // If the array is of length 1 then we know the only value which can be potentially read out of it.
+            // We can then simply assert that the index is equal to zero and return the array's contained value.
+            let array_or_slice_type = dfg.type_of_value(array);
             if matches!(array_or_slice_type, Type::Array(_, 1))
                 && array_or_slice_type.element_size() == 1
             {
-                // If the array is of length 1 then we know the only value which can be potentially read out of it.
-                // We can then simply assert that the index is equal to zero and return the array's contained value.
-                optimize_length_one_array_read(dfg, block, call_stack, *array, *index)
-            } else {
-                None
+                return optimize_length_one_array_read(dfg, block, call_stack, array, index);
             }
+
+            None
         }
         Instruction::ArraySet { array: array_id, index: index_id, value, .. } => {
             let array = dfg.get_array_constant(*array_id);
@@ -356,7 +367,7 @@ fn optimize_length_one_array_read(
     );
     dfg.insert_instruction_and_results(index_constraint, block, None, call_stack);
 
-    let result = try_optimize_array_get_from_previous_set(dfg, array, FieldElement::zero());
+    let result = try_optimize_array_get_from_previous_set(dfg, array, index, FieldElement::zero());
     if let SimplifyResult::None = result {
         SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet { array, index: zero })
     } else {
@@ -379,13 +390,15 @@ fn optimize_length_one_array_read(
 /// - 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.
+///   - If the array value is a parameter, we can directly get from that parameter
 ///
 /// That is, we have multiple `array_set` instructions setting various constant indexes
 /// of the same array, returning a modified version. We want to go backwards until we
 /// find the last `array_set` for the index we are interested in, and return the value set.
 fn try_optimize_array_get_from_previous_set(
     dfg: &DataFlowGraph,
     mut array_id: ValueId,
+    index: ValueId,
     target_index: FieldElement,
 ) -> SimplifyResult {
     let mut elements = None;
@@ -413,6 +426,16 @@ fn try_optimize_array_get_from_previous_set(
                 _ => return SimplifyResult::None,
             }
         } else {
+            // If we reach a parameter, directly get from that parameter, but only if the index is a constant
+            if dfg.get_numeric_constant(index).is_some() {
+                if let Value::Param { .. } = dfg[array_id] {
+                    return SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet {
+                        array: array_id,
+                        index,
+                    });
+                }
+            }
+
             return SimplifyResult::None;
         }
     }
@@ -426,6 +449,76 @@ fn try_optimize_array_get_from_previous_set(
     SimplifyResult::None
 }
 
+fn try_optimize_array_get_on_slice_insert(
+    dfg: &mut DataFlowGraph,
+    array: ValueId,
+    index: ValueId,
+    constant_index: Option<FieldElement>,
+) -> SimplifyResult {
+    let Some(constant_index) = constant_index else {
+        return SimplifyResult::None;
+    };
+
+    let Some(slice_insert) = dfg.get_intrinsic(Intrinsic::SliceInsert) else {
+        return SimplifyResult::None;
+    };
+
+    let Value::Instruction { instruction, .. } = dfg[array] else {
+        return SimplifyResult::None;
+    };
+
+    let Instruction::Call { func, arguments } = &dfg[instruction] else {
+        return SimplifyResult::None;
+    };
+
+    if func != slice_insert {
+        return SimplifyResult::None;
+    }
+
+    // slice_insert(length, slice, index, values...)
+
+    // For simplicity, we only optimize when only a single value is being inserted
+    if arguments.len() != 4 {
+        return SimplifyResult::None;
+    }
+
+    let mut slice = arguments[1];
+    let insert_index = arguments[2];
+
+    let Some(insert_index_constant) = dfg.get_numeric_constant(insert_index) else {
+        return SimplifyResult::None;
+    };
+
+    // If the slice came from an `as_slice` call, we can get from the original array
+    if let Some(as_slice) = dfg.get_intrinsic(Intrinsic::AsSlice) {
+        if let Value::Instruction { instruction, .. } = dfg[slice] {
+            if let Instruction::Call { func, arguments } = &dfg[instruction] {
+                if func == as_slice {
+                    slice = arguments[0];
+                }
+            }
+        }
+    }
+
+    match constant_index.cmp(&insert_index_constant) {
+        std::cmp::Ordering::Equal => {
+            // If we get from the index the value was inserted in, we can simplify it to that value
+            SimplifyResult::SimplifiedTo(arguments[3])
+        }
+        std::cmp::Ordering::Less => {
+            // If the index is before the slice insert index, we can directly read from the original slice
+            SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet { array: slice, index })
+        }
+        std::cmp::Ordering::Greater => {
+            // If the index is after the slice insert index, the index is one less in the original slice
+            // as in the new slice all values after the insert index have been shifted forward
+            let index =
+                dfg.make_constant(constant_index - FieldElement::one(), NumericType::length_type());
+            SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet { array: slice, index })
+        }
+    }
+}
+
 /// If we have an array set whose value is from an array get on the same array at the same index,
 /// we can simplify that array set to the array we were looking to perform an array set upon.
 ///
@@ -707,4 +800,82 @@ mod tests {
 
         assert_normalized_ssa_equals(ssa, src);
     }
+
+    #[test]
+    fn simplifies_array_get_on_slice_insert() {
+        let src = "
+        acir(inline) predicate_pure fn main f0 {
+        b0(v0: [Field; 3]):
+            v2, v3 = call as_slice(v0) -> (u32, [Field])
+            v8, v9 = call slice_insert(u32 3, v3, u32 1, Field 10) -> (u32, [Field])
+            v10 = array_get v9, index u32 0 -> Field
+            v11 = array_get v9, index u32 1 -> Field
+            v12 = array_get v9, index u32 2 -> Field
+            v13 = array_get v9, index u32 3 -> Field
+            return v10, v11, v12, v13
+        }
+        ";
+        let ssa = Ssa::from_str_simplifying(src).unwrap();
+
+        // We can see that the array_gets now read from v0 instead of v9,
+        // and for the index 1 we use the inserted value.
+        assert_ssa_snapshot!(ssa, @r"
+        acir(inline) predicate_pure fn main f0 {
+          b0(v0: [Field; 3]):
+            v2, v3 = call as_slice(v0) -> (u32, [Field])
+            v8, v9 = call slice_insert(u32 3, v3, u32 1, Field 10) -> (u32, [Field])
+            v11 = array_get v0, index u32 0 -> Field
+            v12 = array_get v0, index u32 1 -> Field
+            v14 = array_get v0, index u32 2 -> Field
+            return v11, Field 10, v12, v14
+        }
+        ");
+    }
+
+    #[test]
+    fn simplifies_array_get_from_previous_array_set_with_make_array() {
+        let src = "
+        acir(inline) predicate_pure fn main f0 {
+          b0():
+            v0 = make_array [Field 2, Field 3] : [Field; 2]
+            v1 = array_set mut v0, index u32 0, value Field 4
+            v2 = array_get v1, index u32 0 -> Field
+            v3 = array_get v1, index u32 1 -> Field
+            return v2, v3
+        }
+        ";
+        let ssa = Ssa::from_str_simplifying(src).unwrap();
+
+        assert_ssa_snapshot!(ssa, @r"
+        acir(inline) predicate_pure fn main f0 {
+          b0():
+            v2 = make_array [Field 2, Field 3] : [Field; 2]
+            v4 = make_array [Field 4, Field 3] : [Field; 2]
+            return Field 4, Field 3
+        }
+        ");
+    }
+
+    #[test]
+    fn simplifies_array_get_from_previous_array_set_with_array_param() {
+        let src = "
+        acir(inline) predicate_pure fn main f0 {
+          b0(v0: [Field; 2]):
+            v1 = array_set mut v0, index u32 0, value Field 4
+            v2 = array_get v1, index u32 0 -> Field
+            v3 = array_get v1, index u32 1 -> Field
+            return v2, v3
+        }
+        ";
+        let ssa = Ssa::from_str_simplifying(src).unwrap();
+
+        assert_ssa_snapshot!(ssa, @r"
+        acir(inline) predicate_pure fn main f0 {
+          b0(v0: [Field; 2]):
+            v3 = array_set mut v0, index u32 0, value Field 4
+            v5 = array_get v0, index u32 1 -> Field
+            return Field 4, v5
+        }
+        ");
+    }
 }

compiler/noirc_evaluator/src/ssa/mod.rs

@@ -163,8 +163,11 @@
             "Dead Instruction Elimination",
         ),
         SsaPass::new(Ssa::simplify_cfg, "Simplifying"),
-        SsaPass::new(Ssa::as_slice_optimization, "`as_slice` optimization")
-            .and_then(Ssa::remove_unreachable_functions),
+        SsaPass::new(
+            Ssa::slice_instrinsics_length_optimization,
+            "slice intrinsics length optimization",
+        )
+        .and_then(Ssa::remove_unreachable_functions),
         SsaPass::new_try(
             Ssa::evaluate_static_assert_and_assert_constant,
             "`static_assert` and `assert_constant`",
@@ -229,6 +232,12 @@
         // We cannot run mem2reg after DIE, because it removes Store instructions.
         // We have to run it before, to give it a chance to turn Store+Load into known values.
         SsaPass::new(Ssa::mem2reg, "Mem2Reg"),
+        // After several Mem2Reg passes, some slice intrinsics will be able to be further optimized
+        SsaPass::new(
+            Ssa::slice_instrinsics_length_optimization,
+            "slice intrinsics length optimization",
+        )
+        .and_then(Ssa::remove_unreachable_functions),
         SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"),
         SsaPass::new(Ssa::brillig_entry_point_analysis, "Brillig Entry Point Analysis")
             // Remove any potentially unnecessary duplication from the Brillig entry point analysis.