Fix qubit state across dynamic branches

source/compiler/qsc/src/codegen/tests.rs

@@ -7,6 +7,7 @@ mod adaptive_profile;
 mod adaptive_ri_profile;
 mod adaptive_rif_profile;
 mod base_profile;
+mod early_return_cleanup;
 
 use std::{io::Cursor, rc::Rc, sync::Arc};
 
@@ -91,6 +92,68 @@ fn compile_source_to_qir_result(
     )
 }
 
+#[test]
+fn adaptive_qubit_array_without_loop_compiles() {
+    let source =
+        r#"@EntryPoint(Adaptive) operation Main() : Bool { use qs = Qubit[4]; return false; }"#;
+    let _ = compile_source_to_qir(source, Profile::Adaptive.into());
+}
+
+#[test]
+fn adaptive_loop_with_scalar_qubit_and_early_return_compiles() {
+    let source = r#"@EntryPoint(Adaptive) operation Main() : Bool { for e in [0] { use q = Qubit(); return false; } return false; }"#;
+    let _ = compile_source_to_qir(source, Profile::Adaptive.into());
+}
+
+#[test]
+fn adaptive_loop_with_qubit_array_without_early_return_compiles() {
+    let source = r#"@EntryPoint(Adaptive) operation Main() : Bool { for e in [0] { use qs = Qubit[4]; } return false; }"#;
+    let _ = compile_source_to_qir(source, Profile::Adaptive.into());
+}
+
+#[test]
+fn adaptive_range_loop_with_qubit_array_and_early_return_compiles() {
+    let source = r#"@EntryPoint(Adaptive) operation Main() : Bool { for e in 0..0 { use qs = Qubit[4]; return false; } return false; }"#;
+    let _ = compile_source_to_qir(source, Profile::Adaptive.into());
+}
+
+#[test]
+fn adaptive_loop_with_local_qubit_array_and_early_return_compiles() {
+    let source = r#"
+        @EntryPoint(Adaptive)
+        operation Main() : Bool {
+            for e0 in [0] {
+                use qa0 = Qubit[4];
+                let ret0 = false;
+                return ret0;
+            }
+            return false;
+        }
+    "#;
+
+    let _ = compile_source_to_qir(source, Profile::Adaptive.into());
+}
+
+// Return unification introduces a loop around the early return. Indexing the
+// qubit array in that loop used to panic with "qubit should still be alive".
+#[test]
+fn adaptive_loop_with_indexed_qubit_array_and_early_return_compiles() {
+    let source = r#"
+        @EntryPoint(Adaptive)
+        operation Main() : Int {
+            use qs = Qubit[2];
+            mutable acc = 0;
+            for j in 0..1 {
+                if j == 9 { return 0; }
+                set acc += MResetZ(qs[j]) == One ? 1 | 0;
+            }
+            acc
+        }
+    "#;
+
+    let _ = compile_source_to_qir(source, Profile::Adaptive.into());
+}
+
 #[test]
 fn expanding_polymorphic_recursion_reports_source_errors_on_every_compile_profile() {
     let sources = [
@@ -4805,6 +4868,49 @@ fn early_return_in_dynamic_branch_reinvoke_original_compiles_parity_target() {
     );
 }
 
+#[test]
+fn callable_input_qubit_release_is_not_hidden_by_external_owner() {
+    let caps = TargetCapabilityFlags::from(Profile::AdaptiveRIF)
+        | TargetCapabilityFlags::DynamicQubitAllocation;
+    let source = indoc::indoc! {r#"
+        namespace Test {
+            operation Run(q : Qubit) : Unit {
+                QIR.Runtime.__quantum__rt__qubit_release(q);
+                let fresh = QIR.Runtime.__quantum__rt__qubit_allocate();
+                X(q);
+                QIR.Runtime.__quantum__rt__qubit_release(fresh);
+            }
+        }
+    "#};
+    let (store, pkg, items) = compile_and_locate_items(source, &[("Run", true)], caps);
+    let external = Rc::new(qsc_eval::val::Qubit(0));
+    let args = Value::Qubit((&external).into());
+    let (codegen_fir, backend) =
+        prepare_codegen_fir_from_callable_args(&store, hir_id_for(pkg, items["Run"]), &args, caps)
+            .unwrap_or_else(|errors| {
+                panic!("preparation failed: {}", format_interpret_errors(errors))
+            });
+    let CallableArgsBackend::ReinvokeOriginal { callable, args } = backend else {
+        panic!("qubit arguments must use reinvocation");
+    };
+    // Numeric-ID reuse and distinct allocation identities are asserted directly
+    // by imported_qubit_preserves_identity_across_release_and_reuse. This
+    // integration test checks the diagnostic through callable QIR generation.
+    let error = qsc_codegen::qir::fir_to_qir_from_callable(
+        &codegen_fir.fir_store,
+        caps,
+        &codegen_fir.compute_properties,
+        callable,
+        args,
+    )
+    .expect_err("a caller-owned Rc must not hide release");
+    assert!(matches!(
+        error,
+        qsc_partial_eval::Error::EvaluationFailed(message, _)
+            if message == "qubit used after release"
+    ));
+}
+
 #[test]
 fn synthetic_path_udt_wrapped_controlled_callable_preserves_functor() {
     let source = indoc::indoc! {r#"

source/compiler/qsc/src/codegen/tests/early_return_cleanup.rs

@@ -0,0 +1,103 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+//! Compilation-only checks for the straight-line return suffix rule.
+
+use super::compile_source_to_qir;
+use qsc_data_structures::target::Profile;
+
+fn compile_body(body: &str) -> String {
+    compile_source_to_qir(
+        &format!(
+            r#"
+            operation Before(q : Qubit) : Unit {{ body intrinsic; }}
+            operation Operand(q : Qubit) : Unit {{ body intrinsic; }}
+            operation After(q : Qubit) : Unit {{ body intrinsic; }}
+            operation Tail(q : Qubit) : Unit {{ body intrinsic; }}
+            operation Value(q : Qubit) : Bool {{ Operand(q); false }}
+            operation Run(stop : Bool, q : Qubit) : Bool {{ {body} }}
+            @EntryPoint(Adaptive) operation Main() : Bool {{
+                use q = Qubit();
+                Run(MResetZ(q) == One, q)
+            }}
+            "#
+        ),
+        Profile::Adaptive.into(),
+    )
+}
+
+#[test]
+fn straight_line_nested_return_keeps_operand_effect_and_drops_dead_effects() {
+    let qir = compile_body(
+        "for _ in [0] { Before(q); { { return Value(q); } After(q); } Tail(q); } false",
+    );
+    assert!(qir.contains("call void @Before("), "{qir}");
+    assert!(qir.contains("call void @Operand("), "{qir}");
+    assert!(!qir.contains("call void @After("), "{qir}");
+    assert!(!qir.contains("call void @Tail("), "{qir}");
+}
+
+#[test]
+fn conditional_return_preserves_fallthrough_effects() {
+    let qir =
+        compile_body("for _ in [0] { { if stop { return Value(q); } } After(q); } Tail(q); false");
+    assert!(qir.contains("call void @Operand("), "{qir}");
+    assert!(qir.contains("call void @After("), "{qir}");
+    assert!(qir.contains("call void @Tail("), "{qir}");
+}
+
+#[test]
+fn loop_return_preserves_zero_iteration_fallthrough_effects() {
+    let qir = compile_body("{ while stop { return Value(q); } } After(q); false");
+    assert!(qir.contains("call void @Operand("), "{qir}");
+    assert!(qir.contains("call void @After("), "{qir}");
+}
+
+#[test]
+fn nested_array_and_scalar_cleanup_keeps_return_operand_live() {
+    let qir = compile_body(
+        "use outer = Qubit(); for _ in [0] { use qs = Qubit[2]; { use inner = Qubit(); Before(outer); Before(qs[0]); return Value(inner); } After(q); } Tail(q); false",
+    );
+    assert!(qir.contains("call void @Before("), "{qir}");
+    assert!(qir.contains("call void @Operand("), "{qir}");
+    assert!(!qir.contains("call void @After("), "{qir}");
+    assert!(!qir.contains("call void @Tail("), "{qir}");
+}
+
+#[test]
+fn conditional_array_cleanup_preserves_both_paths() {
+    let qir = compile_body(
+        "for _ in [0] { use qs = Qubit[2]; if stop { return Value(qs[0]); } After(qs[1]); } Tail(q); false",
+    );
+    assert!(qir.contains("call void @Operand("), "{qir}");
+    assert!(qir.contains("call void @After("), "{qir}");
+    assert!(qir.contains("call void @Tail("), "{qir}");
+}
+
+#[test]
+fn conditional_operand_return_compiles_for_adaptive_profiles() {
+    // Return-bearing conditional operands compile across adaptive target profiles.
+    let source = r#"
+        function Run(stop : Bool) : Int {
+            let v = stop ? 1 + (return 1) | 2;
+            v
+        }
+        @EntryPoint(Adaptive_RIF) operation Main() : Int { Run(true) }
+    "#;
+    let _ = compile_source_to_qir(source, Profile::AdaptiveRIF.into());
+    let source = r#"
+        operation Run(stop : Bool, q : Qubit) : Int {
+            X(q);
+            let m = MResetZ(q) == One ? 1 | 0;
+            let v = stop ? 1 + (return 100 + m) | m + 1;
+            v * 10 + m
+        }
+        @EntryPoint(Adaptive_RIF) operation Main() : (Int, Int) {
+            use q = Qubit();
+            (Run(true, q), Run(false, q))
+        }
+    "#;
+    for profile in [Profile::AdaptiveRI, Profile::AdaptiveRIF, Profile::Adaptive] {
+        let _ = compile_source_to_qir(source, profile.into());
+    }
+}

source/compiler/qsc_fir_transforms/src/return_unify/lower.rs

@@ -65,6 +65,23 @@ fn contains_return_in_while_stmt(package: &Package, stmt_id: StmtId) -> bool {
     }
 }
 
+// Only follow straight-line blocks: a return inside a loop or one arm of
+// a conditional does not establish that the containing statement returns.
+fn stmt_unconditionally_returns(package: &Package, stmt_id: StmtId) -> bool {
+    let (StmtKind::Expr(expr_id) | StmtKind::Semi(expr_id)) = package.get_stmt(stmt_id).kind else {
+        return false;
+    };
+    match package.get_expr(expr_id).kind {
+        ExprKind::Return(_) => true,
+        ExprKind::Block(block_id) => package
+            .get_block(block_id)
+            .stmts
+            .iter()
+            .any(|&stmt_id| stmt_unconditionally_returns(package, stmt_id)),
+        _ => false,
+    }
+}
+
 fn sync_block_type_to_stmt_or_unit(package: &mut Package, block_id: BlockId) {
     let trailing_ty = match package.get_block(block_id).stmts.last() {
         Some(&stmt_id) => match package.get_stmt(stmt_id).kind {
@@ -247,6 +264,10 @@ enum FlagBlockOutput {
     Unit,
 }
 
+fn ends_unit_block(package: &Package, stmt_id: StmtId, output: FlagBlockOutput) -> bool {
+    matches!(output, FlagBlockOutput::Unit) && stmt_unconditionally_returns(package, stmt_id)
+}
+
 impl FlagBlockOutput {
     fn lazy(self) -> Self {
         match self {
@@ -338,11 +359,9 @@ fn transform_block_stmts_with_flags(
         // - does it hold a `while` loop that itself contains a `return`?
         // - does it contain a `return` anywhere?
         // - is it the block's final trailing expression (the block's value)?
-        let has_return_in_while = match &package.get_stmt(stmt_id).kind {
-            StmtKind::Expr(e) | StmtKind::Semi(e) => contains_return_in_while_expr(package, *e),
-            _ => false,
-        };
+        let has_return_in_while = contains_return_in_while_stmt(package, stmt_id);
         let has_return = contains_return_in_stmt(package, stmt_id);
+        let ends_unit_block = ends_unit_block(package, stmt_id, output);
         let is_final_trailing_expr = output.final_trailing_expr_strategy().is_some()
             && index == original_stmts.len() - 1
             && matches!(package.get_stmt(stmt_id).kind, StmtKind::Expr(_));
@@ -451,6 +470,11 @@ fn transform_block_stmts_with_flags(
             has_return,
             seen_return_bearing_stmt,
         );
+        // The return has already emitted its cleanup and flag/slot writes.
+        // Any remaining statements on this straight-line path are unreachable.
+        if ends_unit_block {
+            break;
+        }
     }
 
     new_stmts

source/compiler/qsc_fir_transforms/src/return_unify/normalize/tests/flag_strategy.rs

@@ -193,9 +193,6 @@ fn while_body_with_call_arg_return() {
                         __ret_val = 42;
                         __has_returned = true;
                     };
-                    if (not __has_returned) {
-                        i += 1;
-                    };
                 }
 
                 if __has_returned {

source/compiler/qsc_fir_transforms/src/return_unify/normalize/tests/hoist_expression.rs

@@ -556,29 +556,6 @@ fn hoist_return_in_range_endpoint() {
                         __ret_val = 5;
                         __has_returned = true;
                     };
-                    mutable _index_id_31 : Int = if (not __has_returned) {
-                        _range_id_28.Start
-                    } else {
-                        0
-                    };
-                    let _step_id_36 : Int = if (not __has_returned) {
-                        _range_id_28.Step
-                    } else {
-                        0
-                    };
-                    let _end_id_41 : Int = if (not __has_returned) {
-                        _range_id_28.End
-                    } else {
-                        0
-                    };
-                    if (not __has_returned) {
-                        while ((_step_id_36 > 0) and (_index_id_31 <= _end_id_41)) or ((_step_id_36 < 0) and (_index_id_31 >= _end_id_41)) {
-                            let i : Int = _index_id_31;
-                            sum += i;
-                            _index_id_31 += _step_id_36;
-                        }
-
-                    };
                 }
 
                 if __has_returned {
@@ -668,29 +645,6 @@ fn hoist_return_in_local_init_preserves_binding() {
                         __ret_val = 5;
                         __has_returned = true;
                     };
-                    mutable _index_id_95 : Int = if (not __has_returned) {
-                        _range_id_92.Start
-                    } else {
-                        0
-                    };
-                    let _step_id_100 : Int = if (not __has_returned) {
-                        _range_id_92.Step
-                    } else {
-                        0
-                    };
-                    let _end_id_105 : Int = if (not __has_returned) {
-                        _range_id_92.End
-                    } else {
-                        0
-                    };
-                    if (not __has_returned) {
-                        while ((_step_id_100 > 0) and (_index_id_95 <= _end_id_105)) or ((_step_id_100 < 0) and (_index_id_95 >= _end_id_105)) {
-                            let i : Int = _index_id_95;
-                            sum += i;
-                            _index_id_95 += _step_id_100;
-                        }
-
-                    };
                 }
 
                 if __has_returned {