diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs index 676bb48c4d9..6aa9acaca22 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs @@ -818,25 +818,96 @@ fn try_optimize_array_get_from_previous_set( SimplifyResult::None } +/// 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. +/// +/// Simple case: +/// v3 = array_get v1, index v2 +/// v5 = array_set v1, index v2, value v3 +/// +/// If we could not immediately simplify the array set from its value, we can try to follow +/// the array set backwards in the case we have constant indices: +/// +/// v3 = array_get v1, index 1 +/// v5 = array_set v1, index 2, value [Field 100, Field 101, Field 102] +/// v7 = array_set mut v5, index 1, value v3 +/// +/// We want to optimize `v7` to `v5`. We see that `v3` comes from an array get to `v1`. We follow `v5` backwards and see an array set +/// to `v1` and see that the previous array set occurs to a different constant index. +/// +/// 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. +/// - Otherwise, we check the array value of the `array_set`. We will refer to this array as array'. +/// In the case above, array' is `v1` from `v5 = array set ...` +/// - If the original `array_set` value comes from an `array_get`, check the array in that `array_get` against array'. +/// - If the two values are equal we can simplify. +/// - Continuing the example above, as we have `v3 = array_get v1, index 1`, `v1` is +/// what we want to check against array'. We now know we can simplify `v7` to `v5` as it is unchanged. +/// - If they are not equal, recur marking the current `array_set` array as the new array id to use in the checks fn try_optimize_array_set_from_previous_get( dfg: &DataFlowGraph, - array_id: ValueId, + mut array_id: ValueId, target_index: ValueId, target_value: ValueId, ) -> SimplifyResult { - match &dfg[target_value] { + let array_from_get = match &dfg[target_value] { Value::Instruction { instruction, .. } => match &dfg[*instruction] { Instruction::ArrayGet { array, index } => { if *array == array_id && *index == target_index { - SimplifyResult::SimplifiedTo(array_id) + // If array and index match from the value, we can immediately simplify + return SimplifyResult::SimplifiedTo(array_id); + } else if *index == target_index { + *array } else { - SimplifyResult::None + return SimplifyResult::None; } } - _ => SimplifyResult::None, + _ => return SimplifyResult::None, }, - _ => SimplifyResult::None, + _ => return SimplifyResult::None, + }; + + // At this point we have determined that the value we are writing in the `array_set` instruction + // comes from an `array_get` from the same index at which we want to write it at. + // It's possible that we're acting on the same array where other indices have been mutated in between + // the `array_get` and `array_set` (resulting in the `array_id` not matching). + // + // We then inspect the set of `array_set`s which which led to the current array the `array_set` is acting on. + // If we can work back to the array on which the `array_get` was reading from without having another `array_set` + // act on the same index then we can be sure that the new `array_set` can be removed without affecting the final result. + let Some(target_index) = dfg.get_numeric_constant(target_index) else { + return SimplifyResult::None; + }; + + let original_array_id = array_id; + // 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, .. } => { + let Some(index) = dfg.get_numeric_constant(*index) else { + return SimplifyResult::None; + }; + + if index == target_index { + return SimplifyResult::None; + } + + if *array == array_from_get { + return SimplifyResult::SimplifiedTo(original_array_id); + } + + array_id = *array; // recur + } + _ => return SimplifyResult::None, + }, + _ => return SimplifyResult::None, + } } + + SimplifyResult::None } pub(crate) type ErrorType = HirType;