Fix RIR transform bugs (#1514)

This fixes multiple bugs in RIR transformations:

- Updates variable remapping logic during SSA pass to check whether the
proposed replacement itself has a pending remap, fixing scenario where
mutable variables are mapped to other mutable variables.
- Includes variables created by `Store` instructions in the count for
identifying highest variable id to use for inserted phi-nodes, fixing
issues where a phi variable may overlap with an existing remapped
mutable variable.
- Avoids wrapping subtraction in qubit remapping mass when there are no
qubits in the program.
- Fixes type check pass to validate that comparison instruction produces
a Boolean value rather than the same value as arguments.

This change adds some new test cases and expands test coverage
significatly by running all programs produced in partial eval tests
through RIR transforms in the test utility.

compiler/qsc/src/codegen/tests.rs

@@ -625,8 +625,8 @@ mod adaptive_ri_profile {
             block_2:
               br label %block_3
             block_3:
-              %var_2 = phi i64 [0, %block_1], [1, %block_2]
-              call void @__quantum__rt__integer_record_output(i64 %var_2, i8* null)
+              %var_3 = phi i64 [0, %block_1], [1, %block_2]
+              call void @__quantum__rt__integer_record_output(i64 %var_3, i8* null)
               ret void
             }
 

compiler/qsc_partial_eval/tests/test_utils.rs

@@ -9,7 +9,10 @@ use qsc_frontend::compile::{PackageStore as HirPackageStore, SourceMap};
 use qsc_lowerer::{map_hir_package_to_fir, Lowerer};
 use qsc_partial_eval::{partially_evaluate, Error, ProgramEntry};
 use qsc_rca::{Analyzer, PackageStoreComputeProperties};
-use qsc_rir::rir::{BlockId, CallableId, Program};
+use qsc_rir::{
+    passes::check_and_transform,
+    rir::{BlockId, CallableId, Program},
+};
 
 pub fn assert_block_instructions(program: &Program, block_id: BlockId, expected_insts: &Expect) {
     let block = program.get_block(block_id);
@@ -60,7 +63,11 @@ pub fn get_partial_evaluation_error_with_capabilities(
 pub fn get_rir_program(source: &str) -> Program {
     let maybe_program = compile_and_partially_evaluate(source, TargetCapabilityFlags::all());
     match maybe_program {
-        Ok(program) => program,
+        Ok(program) => {
+            // Verify the program can go through transformations.
+            check_and_transform(&mut program.clone());
+            program
+        }
         Err(error) => panic!("partial evaluation failed: {error:?}"),
     }
 }

compiler/qsc_rir/src/passes/reindex_qubits.rs

@@ -4,7 +4,7 @@
 #[cfg(test)]
 mod tests;
 
-use std::collections::hash_map::Entry;
+use std::{collections::hash_map::Entry, ops::Sub};
 
 use qsc_data_structures::index_map::IndexMap;
 use rustc_hash::FxHashMap;
@@ -47,7 +47,11 @@ pub fn reindex_qubits(program: &mut Program) {
         used_cx,
         cx_id,
         mresetz_id,
-        highest_used_id: program.num_qubits - 1,
+        // For this calculation, qubit IDs can never be lower than zero but a program may not use
+        // any qubits. Since `highest_used_id` is only needed for remapping pass and a program without any
+        // qubits won't do any remapping, it's safe to treat this as 1 and let `highest_used_id` default
+        // to zero.
+        highest_used_id: program.num_qubits.max(1).sub(1),
     };
 
     let pred_map = build_predecessors_map(program);

compiler/qsc_rir/src/passes/ssa_transform.rs

@@ -17,19 +17,19 @@ pub fn transform_to_ssa(program: &mut Program, preds: &IndexMap<BlockId, Vec<Blo
     // Ensure that the graph is acyclic before proceeding. Current approach does not support cycles.
     ensure_acyclic(preds);
 
-    // First, remove store instructions and propagate variables through individual blocks.
-    // This produces a per-block map of dynamic variables to their values.
-    // Orphan variables may be left behind where a variable is defined in one block and used in another, which
-    // will be resolved by inserting phi nodes.
-    let mut block_var_map = map_store_to_dominated_ssa(program, preds);
-
     // Get the next available variable ID for use in newly generated phi nodes.
     let mut next_var_id = get_variable_assignments(program)
         .iter()
         .last()
         .map(|(var_id, _)| var_id.successor())
         .unwrap_or_default();
 
+    // First, remove store instructions and propagate variables through individual blocks.
+    // This produces a per-block map of dynamic variables to their values.
+    // Orphan variables may be left behind where a variable is defined in one block and used in another, which
+    // will be resolved by inserting phi nodes.
+    let mut block_var_map = map_store_to_dominated_ssa(program, preds);
+
     // Insert phi nodes where necessary, mapping any remaining orphaned uses to the new variable
     // created by the phi node.
     // This can be done in one pass because the graph is assumed to be acyclic.
@@ -160,14 +160,14 @@ fn map_store_to_dominated_ssa(
 
 // Propagates stored variables through a block, tracking the latest stored value and replacing
 // usage of the variable with the stored value.
-fn map_variable_use_in_block(block: &mut Block, var_map: &mut impl VariableMapper) {
+fn map_variable_use_in_block(block: &mut Block, var_map: &mut FxHashMap<VariableId, Operand>) {
     let instrs = block.0.drain(..).collect::<Vec<_>>();
 
     for mut instr in instrs {
         match &mut instr {
             // Track the new value of the variable and omit the store instruction.
             Instruction::Store(operand, var) => {
-                var_map.insert(*var, *operand);
+                var_map.insert(var.variable_id, *operand);
                 continue;
             }
 
@@ -176,15 +176,15 @@ fn map_variable_use_in_block(block: &mut Block, var_map: &mut impl VariableMappe
                 *args = args
                     .iter()
                     .map(|arg| match arg {
-                        Operand::Variable(var) => var_map.to_operand(*var),
+                        Operand::Variable(var) => var.map_to_operand(var_map),
                         Operand::Literal(_) => *arg,
                     })
                     .collect();
             }
 
             // Replace the branch condition with the new value of the variable.
             Instruction::Branch(var, _, _) => {
-                *var = var_map.to_variable(*var);
+                *var = var.map_to_variable(var_map);
             }
 
             // Two variable instructions, replace left and right operands with new values.
@@ -219,37 +219,35 @@ fn map_variable_use_in_block(block: &mut Block, var_map: &mut impl VariableMappe
 }
 
 impl Operand {
-    fn mapped(&self, var_map: &impl VariableMapper) -> Operand {
+    fn mapped(&self, var_map: &FxHashMap<VariableId, Operand>) -> Operand {
         match self {
             Operand::Literal(_) => *self,
-            Operand::Variable(var) => var_map.to_operand(*var),
+            Operand::Variable(var) => var.map_to_operand(var_map),
         }
     }
 }
 
-trait VariableMapper {
-    fn insert(&mut self, var: Variable, operand: Operand);
-    fn to_operand(&self, var: Variable) -> Operand;
-    fn to_variable(&self, var: Variable) -> Variable;
-}
-
-impl VariableMapper for FxHashMap<VariableId, Operand> {
-    fn insert(&mut self, var: Variable, operand: Operand) {
-        self.insert(var.variable_id, operand);
-    }
-
-    fn to_operand(&self, var: Variable) -> Operand {
-        self.get(&var.variable_id)
-            .copied()
-            .unwrap_or(Operand::Variable(var))
+impl Variable {
+    fn map_to_operand(self, var_map: &FxHashMap<VariableId, Operand>) -> Operand {
+        let mut var = self;
+        while let Some(operand) = var_map.get(&var.variable_id) {
+            if let Operand::Variable(new_var) = operand {
+                var = *new_var;
+            } else {
+                return *operand;
+            }
+        }
+        Operand::Variable(var)
     }
 
-    fn to_variable(&self, var: Variable) -> Variable {
-        self.get(&var.variable_id)
-            .copied()
-            .map_or(var, |operand| match operand {
-                Operand::Literal(_) => panic!("literal not supported in this context"),
-                Operand::Variable(var) => var,
-            })
+    fn map_to_variable(self, var_map: &FxHashMap<VariableId, Operand>) -> Variable {
+        let mut var = self;
+        while let Some(operand) = var_map.get(&var.variable_id) {
+            let Operand::Variable(new_var) = operand else {
+                panic!("literal not supported in this context");
+            };
+            var = *new_var;
+        }
+        var
     }
 }

compiler/qsc_rir/src/passes/ssa_transform/tests.rs

@@ -2221,3 +2221,121 @@ fn ssa_transform_propagates_updates_from_multiple_predecessors_to_later_single_s
             num_qubits: 0
             num_results: 0"#]].assert_eq(&program.to_string());
 }
+
+#[test]
+fn ssa_transform_maps_store_instrs_that_use_values_from_other_store_instrs() {
+    let mut program = new_program();
+    program.callables.insert(
+        CallableId(1),
+        Callable {
+            name: "dynamic_bool".to_string(),
+            input_type: Vec::new(),
+            output_type: Some(Ty::Boolean),
+            body: None,
+            call_type: CallableType::Regular,
+        },
+    );
+
+    program.blocks.insert(
+        BlockId(0),
+        Block(vec![
+            Instruction::Call(
+                CallableId(1),
+                Vec::new(),
+                Some(Variable {
+                    variable_id: VariableId(0),
+                    ty: Ty::Boolean,
+                }),
+            ),
+            Instruction::Store(
+                Operand::Variable(Variable {
+                    variable_id: VariableId(0),
+                    ty: Ty::Boolean,
+                }),
+                Variable {
+                    variable_id: VariableId(1),
+                    ty: Ty::Boolean,
+                },
+            ),
+            Instruction::Store(
+                Operand::Variable(Variable {
+                    variable_id: VariableId(1),
+                    ty: Ty::Boolean,
+                }),
+                Variable {
+                    variable_id: VariableId(2),
+                    ty: Ty::Boolean,
+                },
+            ),
+            Instruction::LogicalNot(
+                Operand::Variable(Variable {
+                    variable_id: VariableId(2),
+                    ty: Ty::Boolean,
+                }),
+                Variable {
+                    variable_id: VariableId(3),
+                    ty: Ty::Boolean,
+                },
+            ),
+            Instruction::Return,
+        ]),
+    );
+
+    // Before
+    expect![[r#"
+        Program:
+            entry: 0
+            callables:
+                Callable 0: Callable:
+                    name: main
+                    call_type: Regular
+                    input_type: <VOID>
+                    output_type: <VOID>
+                    body: 0
+                Callable 1: Callable:
+                    name: dynamic_bool
+                    call_type: Regular
+                    input_type: <VOID>
+                    output_type: Boolean
+                    body: <NONE>
+            blocks:
+                Block 0: Block:
+                    Variable(0, Boolean) = Call id(1), args( )
+                    Variable(1, Boolean) = Store Variable(0, Boolean)
+                    Variable(2, Boolean) = Store Variable(1, Boolean)
+                    Variable(3, Boolean) = LogicalNot Variable(2, Boolean)
+                    Return
+            config: Config:
+                capabilities: Base
+            num_qubits: 0
+            num_results: 0"#]]
+    .assert_eq(&program.to_string());
+
+    // After
+    transform_program(&mut program);
+    expect![[r#"
+        Program:
+            entry: 0
+            callables:
+                Callable 0: Callable:
+                    name: main
+                    call_type: Regular
+                    input_type: <VOID>
+                    output_type: <VOID>
+                    body: 0
+                Callable 1: Callable:
+                    name: dynamic_bool
+                    call_type: Regular
+                    input_type: <VOID>
+                    output_type: Boolean
+                    body: <NONE>
+            blocks:
+                Block 0: Block:
+                    Variable(0, Boolean) = Call id(1), args( )
+                    Variable(3, Boolean) = LogicalNot Variable(0, Boolean)
+                    Return
+            config: Config:
+                capabilities: TargetCapabilityFlags(Adaptive | IntegerComputations | FloatingPointComputations | BackwardsBranching | HigherLevelConstructs | QubitReset)
+            num_qubits: 0
+            num_results: 0"#]].assert_eq(&program.to_string());
+}

compiler/qsc_rir/src/passes/type_check.rs

@@ -31,12 +31,16 @@ fn check_instr_types(program: &Program, instr: &Instruction) {
         | Instruction::LogicalOr(opr1, opr2, var)
         | Instruction::BitwiseAnd(opr1, opr2, var)
         | Instruction::BitwiseOr(opr1, opr2, var)
-        | Instruction::BitwiseXor(opr1, opr2, var)
-        | Instruction::Icmp(_, opr1, opr2, var) => {
+        | Instruction::BitwiseXor(opr1, opr2, var) => {
             assert_eq!(opr1.get_type(), opr2.get_type());
             assert_eq!(opr1.get_type(), var.ty);
         }
 
+        Instruction::Icmp(_, opr1, opr2, var) => {
+            assert_eq!(opr1.get_type(), opr2.get_type());
+            assert_eq!(Ty::Boolean, var.ty);
+        }
+
         Instruction::Store(opr, var)
         | Instruction::LogicalNot(opr, var)
         | Instruction::BitwiseNot(opr, var) => {

compiler/qsc_rir/src/rir.rs

@@ -6,7 +6,7 @@ use qsc_data_structures::{index_map::IndexMap, target::TargetCapabilityFlags};
 use std::fmt::{self, Display, Formatter, Write};
 
 /// The root of the RIR.
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct Program {
     pub entry: CallableId,
     pub callables: IndexMap<CallableId, Callable>,
@@ -63,7 +63,7 @@ impl Program {
     }
 }
 
-#[derive(Default)]
+#[derive(Default, Clone, Copy)]
 pub struct Config {
     pub capabilities: TargetCapabilityFlags,
 }
@@ -129,7 +129,7 @@ impl BlockId {
 }
 
 /// A block is a collection of instructions.
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct Block(pub Vec<Instruction>);
 
 /// A unique identifier for a callable in a RIR program.
@@ -156,7 +156,7 @@ impl CallableId {
 }
 
 /// A callable.
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Callable {
     /// The name of the callable.
     pub name: String,

compiler/qsc_rir/src/utils.rs

@@ -72,6 +72,8 @@ pub fn build_predecessors_map(program: &Program) -> IndexMap<BlockId, Vec<BlockI
 #[must_use]
 pub fn get_variable_assignments(program: &Program) -> IndexMap<VariableId, (BlockId, usize)> {
     let mut assignments = IndexMap::default();
+    let mut has_store = false;
+    let mut has_phi = false;
     for (block_id, block) in program.blocks.iter() {
         for (idx, instr) in block.0.iter().enumerate() {
             match instr {
@@ -97,18 +99,24 @@ pub fn get_variable_assignments(program: &Program) -> IndexMap<VariableId, (Bloc
                         "Duplicate assignment to {:?} in {block_id:?}, instruction {idx}",
                         var.variable_id
                     );
+                    has_phi |= matches!(instr, Instruction::Phi(_, _));
                     assignments.insert(var.variable_id, (block_id, idx));
                 }
+                Instruction::Store(_, var) => {
+                    has_store = true;
+                    assignments.insert(var.variable_id, (block_id, idx));
+                }
+
                 Instruction::Call(_, _, None)
                 | Instruction::Jump(..)
                 | Instruction::Branch(..)
                 | Instruction::Return => {}
-
-                Instruction::Store(..) => {
-                    panic!("Unexpected Store at {block_id:?}, instruction {idx}")
-                }
             }
         }
     }
+    assert!(
+        !(has_store && has_phi),
+        "Program has both store and phi instructions."
+    );
     assignments
 }