From 765ad493dbc3ccf055a34611db7ca8b943ee347a Mon Sep 17 00:00:00 2001
From: "Stefan J. Wernli" <swernli@microsoft.com>
Date: Wed, 15 May 2024 09:41:24 -0700
Subject: [PATCH] Fix panic when creating unneeded Phi node (#1527)

This fixes an issue where a mutable variable A used only in one scope
and not across any others is set to a different mutable variable B that
spans multiple scopes, including scopes that are not dominated by
variable B. Variables like B do not need phi nodes as they are assumed
to be unused in later blocks, and any actual usage in later blocks will
be caught by the SSA check pass by using the dominator tree.
---
 compiler/qsc_partial_eval/tests/bindings.rs  | 86 ++++++++++++++++++++
 compiler/qsc_rir/src/passes/ssa_transform.rs | 20 +++--
 2 files changed, 99 insertions(+), 7 deletions(-)

diff --git a/compiler/qsc_partial_eval/tests/bindings.rs b/compiler/qsc_partial_eval/tests/bindings.rs
index 568dd0452b..c55f568255 100644
--- a/compiler/qsc_partial_eval/tests/bindings.rs
+++ b/compiler/qsc_partial_eval/tests/bindings.rs
@@ -397,3 +397,89 @@ fn mutable_int_binding_does_generate_store_instruction() {
                 Jump(1)"#]],
     );
 }
+
+#[test]
+fn mutable_variable_in_outer_scope_set_to_mutable_from_inner_scope() {
+    let program = get_rir_program(indoc! {
+        r#"
+        namespace Test {
+            @EntryPoint()
+            operation Main() : Int {
+                use q = Qubit();
+                mutable i = 0;
+                if MResetZ(q) == One {
+                    mutable j = 1;
+                    set i = j;
+                }
+                else {
+                    set i = 2;
+                }
+                return i;
+            }
+        }
+        "#,
+    });
+
+    let measurement_callable_id = CallableId(1);
+    assert_callable(
+        &program,
+        measurement_callable_id,
+        &expect![[r#"
+            Callable:
+                name: __quantum__qis__mresetz__body
+                call_type: Measurement
+                input_type:
+                    [0]: Qubit
+                    [1]: Result
+                output_type: <VOID>
+                body: <NONE>"#]],
+    );
+    let read_result_callable_id = CallableId(2);
+    assert_callable(
+        &program,
+        read_result_callable_id,
+        &expect![[r#"
+            Callable:
+                name: __quantum__qis__read_result__body
+                call_type: Readout
+                input_type:
+                    [0]: Result
+                output_type: Boolean
+                body: <NONE>"#]],
+    );
+    let output_recording_callable_id = CallableId(3);
+    assert_callable(
+        &program,
+        output_recording_callable_id,
+        &expect![[r#"
+            Callable:
+                name: __quantum__rt__int_record_output
+                call_type: OutputRecording
+                input_type:
+                    [0]: Integer
+                    [1]: Pointer
+                output_type: <VOID>
+                body: <NONE>"#]],
+    );
+    assert_blocks(
+        &program,
+        &expect![[r#"
+        Blocks:
+        Block 0:Block:
+            Variable(0, Integer) = Store Integer(0)
+            Call id(1), args( Qubit(0), Result(0), )
+            Variable(1, Boolean) = Call id(2), args( Result(0), )
+            Variable(2, Boolean) = Store Variable(1, Boolean)
+            Branch Variable(2, Boolean), 2, 3
+        Block 1:Block:
+            Call id(3), args( Variable(0, Integer), Pointer, )
+            Return
+        Block 2:Block:
+            Variable(3, Integer) = Store Integer(1)
+            Variable(0, Integer) = Store Integer(1)
+            Jump(1)
+        Block 3:Block:
+            Variable(0, Integer) = Store Integer(2)
+            Jump(1)"#]],
+    );
+}
diff --git a/compiler/qsc_rir/src/passes/ssa_transform.rs b/compiler/qsc_rir/src/passes/ssa_transform.rs
index 38be3936dc..040adad607 100644
--- a/compiler/qsc_rir/src/passes/ssa_transform.rs
+++ b/compiler/qsc_rir/src/passes/ssa_transform.rs
@@ -59,10 +59,10 @@ pub fn transform_to_ssa(program: &mut Program, preds: &IndexMap<BlockId, Vec<Blo
         } else {
             // Check each variable in the first predecessor's variable map, and if any other
             // predecessor has a different value for the variable, a phi node is needed.
-            for (var_id, operand) in block_var_map
+            let first_pred_map = block_var_map
                 .get(*first_pred)
-                .expect("block should have variable map")
-            {
+                .expect("block should have variable map");
+            'var_loop: for (var_id, operand) in first_pred_map {
                 let mut phi_nodes = FxHashMap::default();
 
                 if rest_preds.iter().any(|pred| {
@@ -74,14 +74,20 @@ pub fn transform_to_ssa(program: &mut Program, preds: &IndexMap<BlockId, Vec<Blo
                 }) {
                     // Some predecessors have different values for this variable, so a phi node is needed.
                     // Start with the first predecessor's value and block id, then add the values from the other predecessors.
-                    let mut phi_args = vec![(*operand, *first_pred)];
+                    let mut phi_args = vec![(operand.mapped(first_pred_map), *first_pred)];
                     for pred in rest_preds {
                         let pred_var_map = block_var_map
                             .get(*pred)
                             .expect("block should have variable map");
-                        let mut pred_operand = *pred_var_map
-                            .get(var_id)
-                            .expect("variable should be in predecessor's variable map");
+                        let mut pred_operand = match pred_var_map.get(var_id) {
+                            Some(operand) => *operand,
+                            None => {
+                                // If the variable is not defined in this predecessor, it does not dominate this block.
+                                // Assume it is not used and skip creating a phi node for this variable. If the variable is used,
+                                // the ssa check will detect it and panic later.
+                                continue 'var_loop;
+                            }
+                        };
                         pred_operand = pred_operand.mapped(pred_var_map);
                         phi_args.push((pred_operand, *pred));
                     }