Fix panic in QIR generation when conditional branches use early return

compiler/qsc_partial_eval/src/evaluation_context.rs

@@ -262,14 +262,6 @@ impl Arg {
     }
 }
 
-/// Represents the possible control flow options that can result from a branch.
-pub enum BranchControlFlow {
-    /// The block ID corresponding to a branch.
-    Block(BlockId),
-    /// The return value resulting from a branch.
-    Return(Value),
-}
-
 /// Represents the possible control flow options that an evaluation can have.
 pub enum EvalControlFlow {
     Continue(Value),

compiler/qsc_partial_eval/src/lib.rs

@@ -12,9 +12,7 @@ mod evaluation_context;
 mod management;
 
 use core::panic;
-use evaluation_context::{
-    Arg, BlockNode, BranchControlFlow, EvalControlFlow, EvaluationContext, Scope,
-};
+use evaluation_context::{Arg, BlockNode, EvalControlFlow, EvaluationContext, Scope};
 use management::{QuantumIntrinsicsChecker, ResourceManager};
 use miette::Diagnostic;
 use qsc_data_structures::{functors::FunctorApp, span::Span, target::TargetCapabilityFlags};
@@ -1772,31 +1770,24 @@ impl<'a> PartialEvaluator<'a> {
         // Evaluate the body expression.
         // First, we cache the current static variable mappings so that we can restore them later.
         let cached_mappings = self.clone_current_static_var_map();
-        let if_true_branch_control_flow =
+        let if_true_block_id =
             self.eval_expr_if_branch(body_expr_id, continuation_block_node_id, maybe_if_expr_var)?;
-        let if_true_block_id = match if_true_branch_control_flow {
-            BranchControlFlow::Block(block_id) => block_id,
-            BranchControlFlow::Return(value) => return Ok(EvalControlFlow::Return(value)),
-        };
 
         // Evaluate the otherwise expression (if any), and determine the block to branch to if the condition is false.
         let if_false_block_id = if let Some(otherwise_expr_id) = otherwise_expr_id {
             // Cache the mappings after the true block so we can compare afterwards.
             let post_if_true_mappings = self.clone_current_static_var_map();
             // Restore the cached mappings from before evaluating the true block.
             self.overwrite_current_static_var_map(cached_mappings);
-            let if_false_branch_control_flow = self.eval_expr_if_branch(
+            let if_false_block_id = self.eval_expr_if_branch(
                 otherwise_expr_id,
                 continuation_block_node_id,
                 maybe_if_expr_var,
             )?;
             // Only keep the static mappings that are the same in both blocks; when they are different,
             // the variable is no longer static across the if expression.
             self.keep_matching_static_var_mappings(&post_if_true_mappings);
-            match if_false_branch_control_flow {
-                BranchControlFlow::Block(block_id) => block_id,
-                BranchControlFlow::Return(value) => return Ok(EvalControlFlow::Return(value)),
-            }
+            if_false_block_id
         } else {
             // Only keep the static mappings that are the same after the true block as before; when they are different,
             // the variable is no longer static across the if expression.
@@ -1837,7 +1828,7 @@ impl<'a> PartialEvaluator<'a> {
         branch_body_expr_id: ExprId,
         continuation_block_id: rir::BlockId,
         if_expr_var: Option<rir::Variable>,
-    ) -> Result<BranchControlFlow, Error> {
+    ) -> Result<rir::BlockId, Error> {
         // Create the block node that corresponds to the branch body and push it as the active one.
         let block_node_id = self.create_program_block();
         let block_node = BlockNode {
@@ -1849,7 +1840,8 @@ impl<'a> PartialEvaluator<'a> {
         // Evaluate the branch body expression.
         let body_control = self.try_eval_expr(branch_body_expr_id)?;
         if body_control.is_return() {
-            return Ok(BranchControlFlow::Return(body_control.into_value()));
+            let body_span = self.get_expr_package_span(branch_body_expr_id);
+            return Err(Error::Unimplemented("early return".to_string(), body_span));
         }
 
         // If there is a variable to save the value of the if expression to, add a store instruction.
@@ -1863,7 +1855,7 @@ impl<'a> PartialEvaluator<'a> {
         let jump_ins = Instruction::Jump(continuation_block_id);
         self.get_current_rir_block_mut().0.push(jump_ins);
         let _ = self.eval_context.pop_block_node();
-        Ok(BranchControlFlow::Block(block_node_id))
+        Ok(block_node_id)
     }
 
     fn eval_expr_if_with_classical_condition(

compiler/qsc_partial_eval/src/tests/branching.rs

@@ -2311,3 +2311,95 @@ fn if_expression_with_dynamic_operand_from_hybrid_doubles_array() {
                 Jump(6)"#]],
     );
 }
+
+#[test]
+fn if_expression_with_implicit_return_in_callable_supported() {
+    let program = get_rir_program(indoc! {r#"
+        function Choose(r : Result) : Int {
+            if r == One {
+                1
+            } else {
+                0
+            }
+        }
+        @EntryPoint()
+        operation Main() : Int {
+            use q = Qubit();
+            Choose(MResetZ(q))
+        }
+        "#
+    });
+
+    assert_blocks(
+        &program,
+        &expect![[r#"
+            Blocks:
+            Block 0:Block:
+                Call id(1), args( Qubit(0), Result(0), )
+                Variable(0, Boolean) = Call id(2), args( Result(0), )
+                Variable(1, Boolean) = Store Variable(0, Boolean)
+                Branch Variable(1, Boolean), 2, 3
+            Block 1:Block:
+                Variable(3, Integer) = Store Variable(2, Integer)
+                Call id(3), args( Variable(3, Integer), Pointer, )
+                Return
+            Block 2:Block:
+                Variable(2, Integer) = Store Integer(1)
+                Jump(1)
+            Block 3:Block:
+                Variable(2, Integer) = Store Integer(0)
+                Jump(1)"#]],
+    );
+}
+
+#[test]
+fn if_expression_with_explicit_return_in_callable_fails() {
+    let error = get_partial_evaluation_error(indoc! {r#"
+        function Choose(r : Result) : Int {
+            if r == One {
+                return 1;
+            } else {
+                return 0;
+            }
+        }
+        @EntryPoint()
+        operation Main() : Int {
+            use q = Qubit();
+            Choose(MResetZ(q))
+        }
+        "#
+    });
+
+    assert_error(
+        &error,
+        &expect![[
+            r#"Unimplemented("early return", PackageSpan { package: PackageId(2), span: Span { lo: 53, hi: 78 } })"#
+        ]],
+    );
+}
+
+#[test]
+fn if_expression_with_explicit_return_in_one_branch_and_fallthrough_else_in_callable_fails() {
+    let error = get_partial_evaluation_error(indoc! {r#"
+        function Choose(r : Result) : Int {
+            if r == One {
+                return 1;
+            }
+
+            return 3;
+        }
+        @EntryPoint()
+        operation Main() : Int {
+            use q = Qubit();
+            Choose(MResetZ(q))
+        }
+        "#
+    });
+
+    assert_error(
+        &error,
+        &expect![[
+            r#"Unimplemented("early return", PackageSpan { package: PackageId(2), span: Span { lo: 53, hi: 78 } })"#
+        ]],
+    );
+}

compiler/qsc_partial_eval/src/tests/returns.rs

@@ -365,7 +365,7 @@ fn non_classical_entry_point_with_classical_early_return_within_non_classical_br
     assert_error(
         &error,
         &expect![[
-            r#"Unimplemented("early return", PackageSpan { package: PackageId(2), span: Span { lo: 163, hi: 213 } })"#
+            r#"Unimplemented("early return", PackageSpan { package: PackageId(2), span: Span { lo: 176, hi: 213 } })"#
         ]],
     );
 }
@@ -389,7 +389,7 @@ fn non_classical_entry_point_with_non_classical_early_return_within_non_classica
     assert_error(
         &error,
         &expect![[
-            r#"Unimplemented("early return", PackageSpan { package: PackageId(2), span: Span { lo: 185, hi: 278 } })"#
+            r#"Unimplemented("early return", PackageSpan { package: PackageId(2), span: Span { lo: 199, hi: 278 } })"#
         ]],
     );
 }

compiler/qsc_passes/src/capabilitiesck/tests_adaptive.rs

@@ -579,11 +579,18 @@ fn use_of_dynamic_rhs_exp_binop_yields_errors() {
 }
 
 #[test]
-fn return_within_dynamic_scope_yields_no_errors() {
+fn return_within_dynamic_scope_yields_errors() {
     check_profile(
         RETURN_WITHIN_DYNAMIC_SCOPE,
         &expect![[r#"
-            []
+            [
+                ReturnWithinDynamicScope(
+                    Span {
+                        lo: 128,
+                        hi: 136,
+                    },
+                ),
+            ]
         "#]],
     );
 }

compiler/qsc_passes/src/capabilitiesck/tests_adaptive_plus_integers.rs

@@ -455,11 +455,18 @@ fn use_of_dynamic_rhs_exp_binop_yields_errors() {
 }
 
 #[test]
-fn return_within_dynamic_scope_yields_no_errors() {
+fn return_within_dynamic_scope_yields_errors() {
     check_profile(
         RETURN_WITHIN_DYNAMIC_SCOPE,
         &expect![[r#"
-            []
+            [
+                ReturnWithinDynamicScope(
+                    Span {
+                        lo: 128,
+                        hi: 136,
+                    },
+                ),
+            ]
         "#]],
     );
 }

compiler/qsc_passes/src/capabilitiesck/tests_adaptive_plus_integers_and_floats.rs

@@ -440,11 +440,18 @@ fn use_of_dynamic_rhs_exp_binop_yields_errors() {
 }
 
 #[test]
-fn return_within_dynamic_scope_yields_no_errors() {
+fn return_within_dynamic_scope_yields_errors() {
     check_profile(
         RETURN_WITHIN_DYNAMIC_SCOPE,
         &expect![[r#"
-            []
+            [
+                ReturnWithinDynamicScope(
+                    Span {
+                        lo: 128,
+                        hi: 136,
+                    },
+                ),
+            ]
         "#]],
     );
 }

compiler/qsc_rca/src/lib.rs

@@ -890,7 +890,7 @@ impl RuntimeFeatureFlags {
             capabilities |= TargetCapabilityFlags::HigherLevelConstructs;
         }
         if self.contains(RuntimeFeatureFlags::ReturnWithinDynamicScope) {
-            capabilities |= TargetCapabilityFlags::Adaptive;
+            capabilities |= TargetCapabilityFlags::HigherLevelConstructs;
         }
         if self.contains(RuntimeFeatureFlags::LoopWithDynamicCondition) {
             capabilities |= TargetCapabilityFlags::BackwardsBranching;