Support partial evaluation of IndexRange calls (#2727)

This updates the logic for accessing fields of UDTs during partial
evaluation and using them as loop bounds such that usage of
`Std.Arrays.IndexRange` on arrays with compile-time known length works
as expected during code generation.

Author: swernli

source/compiler/qsc_partial_eval/src/lib.rs

@@ -30,16 +30,16 @@ use qsc_eval::{
 use qsc_fir::{
     fir::{
         self, BinOp, Block, BlockId, CallableDecl, CallableImpl, ExecGraph, Expr, ExprId, ExprKind,
-        Global, Ident, LocalVarId, Mutability, PackageId, PackageStore, PackageStoreLookup, Pat,
-        PatId, PatKind, Res, SpecDecl, SpecImpl, Stmt, StmtId, StmtKind, StoreBlockId, StoreExprId,
-        StoreItemId, StorePatId, StoreStmtId, UnOp,
+        Field, Global, Ident, LocalVarId, Mutability, PackageId, PackageStore, PackageStoreLookup,
+        Pat, PatId, PatKind, PrimField, Res, SpecDecl, SpecImpl, Stmt, StmtId, StmtKind,
+        StoreBlockId, StoreExprId, StoreItemId, StorePatId, StoreStmtId, UnOp,
     },
     ty::{Prim, Ty},
 };
 use qsc_lowerer::map_fir_package_to_hir;
 use qsc_rca::{
     ComputeKind, ComputePropertiesLookup, ItemComputeProperties, PackageStoreComputeProperties,
-    QuantumProperties, RuntimeFeatureFlags,
+    QuantumProperties, RuntimeFeatureFlags, RuntimeKind, ValueKind,
     errors::{
         Error as CapabilityError, generate_errors_from_runtime_features,
         get_missing_runtime_features,
@@ -1180,10 +1180,7 @@ impl<'a> PartialEvaluator<'a> {
                 "using a dynamic value in a fail statement is invalid".to_string(),
                 expr_package_span,
             )),
-            ExprKind::Field(_, _) => Err(Error::Unexpected(
-                "accessing a field of a dynamic user-defined type is invalid".to_string(),
-                expr_package_span,
-            )),
+            ExprKind::Field(expr_id, field) => self.eval_expr_field(*expr_id, field.clone()),
             ExprKind::Hole => Err(Error::Unexpected(
                 "hole expressions are not expected during partial evaluation".to_string(),
                 expr_package_span,
@@ -1978,6 +1975,47 @@ impl<'a> PartialEvaluator<'a> {
         Ok(EvalControlFlow::Continue(value))
     }
 
+    fn eval_expr_field(
+        &mut self,
+        record_id: ExprId,
+        field: Field,
+    ) -> Result<EvalControlFlow, Error> {
+        let control_flow = self.try_eval_expr(record_id)?;
+        let EvalControlFlow::Continue(record) = control_flow else {
+            return Err(Error::Unexpected(
+                "embedded return in field access expression".to_string(),
+                self.get_expr_package_span(record_id),
+            ));
+        };
+
+        let field_value = match (record, field) {
+            (Value::Range(inner), Field::Prim(PrimField::Start)) => Value::Int(
+                inner
+                    .start
+                    .expect("range access should be validated by compiler"),
+            ),
+            (Value::Range(inner), Field::Prim(PrimField::Step)) => Value::Int(inner.step),
+            (Value::Range(inner), Field::Prim(PrimField::End)) => Value::Int(
+                inner
+                    .end
+                    .expect("range access should be validated by compiler"),
+            ),
+            (mut record, Field::Path(path)) => {
+                for index in path.indices {
+                    let Value::Tuple(items, _) = record else {
+                        panic!("invalid tuple access");
+                    };
+                    record = items[index].clone();
+                }
+                record
+            }
+            (ref value, ref field) => {
+                panic!("invalid field access. value: {value:?}, field: {field:?}")
+            }
+        };
+        Ok(EvalControlFlow::Continue(field_value))
+    }
+
     fn eval_expr_return(&mut self, expr_id: ExprId) -> Result<EvalControlFlow, Error> {
         let control_flow = self.try_eval_expr(expr_id)?;
         Ok(EvalControlFlow::Return(control_flow.into_value()))
@@ -2155,17 +2193,19 @@ impl<'a> PartialEvaluator<'a> {
         condition_expr_id: ExprId,
         body_block_id: BlockId,
     ) -> Result<EvalControlFlow, Error> {
-        // Verify assumptions.
+        // Verify assumptions: the condition expression must either classical (such that it can be fully evaluated) or
+        // quantum but statically known at runtime (such that it can be partially evaluated to a known value).
         assert!(
-            self.is_classical_expr(condition_expr_id),
+            matches!(
+                self.get_expr_compute_kind(condition_expr_id),
+                ComputeKind::Classical
+                    | ComputeKind::Quantum(QuantumProperties {
+                        runtime_features: _,
+                        value_kind: ValueKind::Element(RuntimeKind::Static),
+                    })
+            ),
             "loop conditions must be purely classical"
         );
-        let body_block = self.get_block(body_block_id);
-        assert_eq!(
-            body_block.ty,
-            Ty::UNIT,
-            "the type of a loop block is expected to be Unit"
-        );
 
         // Evaluate the block until the loop condition is false.
         let condition_expr_span = self.get_expr_package_span(condition_expr_id);

source/compiler/qsc_partial_eval/src/tests/arrays.rs

@@ -1038,3 +1038,62 @@ fn result_array_copy_and_update_with_out_of_bounds_range_raises_error() {
         ]],
     );
 }
+
+#[test]
+fn result_array_index_range_returns_length_as_end() {
+    let program = get_rir_program(indoc! {r#"
+        namespace Test {
+            @EntryPoint()
+            operation Main() : Int {
+                use qs = Qubit[2];
+                let results = MResetEachZ(qs);
+                Std.Arrays.IndexRange(results).End
+            }
+        }
+    "#});
+    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: __quantum__qis__mresetz__body
+                    call_type: Measurement
+                    input_type:
+                        [0]: Qubit
+                        [1]: Result
+                    output_type: <VOID>
+                    body: <NONE>
+                Callable 2: Callable:
+                    name: __quantum__rt__int_record_output
+                    call_type: OutputRecording
+                    input_type:
+                        [0]: Integer
+                        [1]: Pointer
+                    output_type: <VOID>
+                    body: <NONE>
+            blocks:
+                Block 0: Block:
+                    Variable(0, Integer) = Store Integer(0)
+                    Variable(0, Integer) = Store Integer(1)
+                    Variable(0, Integer) = Store Integer(2)
+                    Variable(1, Integer) = Store Integer(0)
+                    Call id(1), args( Qubit(0), Result(0), )
+                    Variable(1, Integer) = Store Integer(1)
+                    Call id(1), args( Qubit(1), Result(1), )
+                    Variable(1, Integer) = Store Integer(2)
+                    Variable(2, Integer) = Store Integer(0)
+                    Variable(2, Integer) = Store Integer(1)
+                    Variable(2, Integer) = Store Integer(2)
+                    Call id(2), args( Integer(1), Pointer, )
+                    Return
+            config: Config:
+                capabilities: TargetCapabilityFlags(Adaptive | IntegerComputations | FloatingPointComputations | BackwardsBranching | HigherLevelConstructs | QubitReset)
+            num_qubits: 2
+            num_results: 2"#]].assert_eq(&program.to_string());
+}

source/compiler/qsc_partial_eval/src/tests/loops.rs

@@ -399,7 +399,6 @@ fn mutable_double_updated_in_loop() {
 
     assert_blocks(
         &program,
-        //BlockId(0),
         &expect![[r#"
             Blocks:
             Block 0:Block:
@@ -452,3 +451,95 @@ fn mutable_double_updated_in_loop() {
                 Jump(9)"#]],
     );
 }
+
+#[test]
+fn result_array_index_range_in_for_loop() {
+    let program = get_rir_program(indoc! {r#"
+        namespace Test {
+            @EntryPoint()
+            operation Main() : Int {
+                use qs = Qubit[2];
+                let results = MResetEachZ(qs);
+                mutable count = 0;
+                for i in Std.Arrays.IndexRange(results) {
+                    if results[i] == One {
+                        set count += 1;
+                    }
+                }
+                count
+            }
+        }
+    "#});
+    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: __quantum__qis__mresetz__body
+                    call_type: Measurement
+                    input_type:
+                        [0]: Qubit
+                        [1]: Result
+                    output_type: <VOID>
+                    body: <NONE>
+                Callable 2: Callable:
+                    name: __quantum__qis__read_result__body
+                    call_type: Readout
+                    input_type:
+                        [0]: Result
+                    output_type: Boolean
+                    body: <NONE>
+                Callable 3: Callable:
+                    name: __quantum__rt__int_record_output
+                    call_type: OutputRecording
+                    input_type:
+                        [0]: Integer
+                        [1]: Pointer
+                    output_type: <VOID>
+                    body: <NONE>
+            blocks:
+                Block 0: Block:
+                    Variable(0, Integer) = Store Integer(0)
+                    Variable(0, Integer) = Store Integer(1)
+                    Variable(0, Integer) = Store Integer(2)
+                    Variable(1, Integer) = Store Integer(0)
+                    Call id(1), args( Qubit(0), Result(0), )
+                    Variable(1, Integer) = Store Integer(1)
+                    Call id(1), args( Qubit(1), Result(1), )
+                    Variable(1, Integer) = Store Integer(2)
+                    Variable(2, Integer) = Store Integer(0)
+                    Variable(3, Integer) = Store Integer(0)
+                    Variable(4, Boolean) = Call id(2), args( Result(0), )
+                    Variable(5, Boolean) = Store Variable(4, Boolean)
+                    Branch Variable(5, Boolean), 2, 1
+                Block 1: Block:
+                    Variable(3, Integer) = Store Integer(1)
+                    Variable(6, Boolean) = Call id(2), args( Result(1), )
+                    Variable(7, Boolean) = Store Variable(6, Boolean)
+                    Branch Variable(7, Boolean), 4, 3
+                Block 2: Block:
+                    Variable(2, Integer) = Store Integer(1)
+                    Jump(1)
+                Block 3: Block:
+                    Variable(3, Integer) = Store Integer(2)
+                    Variable(9, Integer) = Store Variable(2, Integer)
+                    Variable(10, Integer) = Store Integer(0)
+                    Variable(10, Integer) = Store Integer(1)
+                    Variable(10, Integer) = Store Integer(2)
+                    Call id(3), args( Variable(9, Integer), Pointer, )
+                    Return
+                Block 4: Block:
+                    Variable(8, Integer) = Add Variable(2, Integer), Integer(1)
+                    Variable(2, Integer) = Store Variable(8, Integer)
+                    Jump(3)
+            config: Config:
+                capabilities: TargetCapabilityFlags(Adaptive | IntegerComputations | FloatingPointComputations | BackwardsBranching | HigherLevelConstructs | QubitReset)
+            num_qubits: 2
+            num_results: 2"#]].assert_eq(&program.to_string());
+}

source/compiler/qsc_rca/src/tests/arrays.rs

@@ -542,3 +542,42 @@ fn check_rca_for_array_with_static_size_bound_through_dynamic_tuple() {
                 dynamic_param_applications: <empty>"#]],
     );
 }
+
+#[test]
+fn check_rca_for_index_range_on_array_with_classical_contents() {
+    let mut compilation_context = CompilationContext::default();
+    compilation_context.update(
+        r#"
+        let arr = [0, 1, 2, 3, 4];
+        Std.Arrays.IndexRange(arr)"#,
+    );
+    let package_store_compute_properties = compilation_context.get_compute_properties();
+    check_last_statement_compute_properties(
+        package_store_compute_properties,
+        &expect![[r#"
+            ApplicationsGeneratorSet:
+                inherent: Classical
+                dynamic_param_applications: <empty>"#]],
+    );
+}
+
+#[test]
+fn check_rca_for_index_range_on_array_with_dynamic_contents() {
+    let mut compilation_context = CompilationContext::default();
+    compilation_context.update(
+        r#"
+        use qs = Qubit[5];
+        let arr = MResetEachZ(qs);
+        Std.Arrays.IndexRange(arr)"#,
+    );
+    let package_store_compute_properties = compilation_context.get_compute_properties();
+    check_last_statement_compute_properties(
+        package_store_compute_properties,
+        &expect![[r#"
+            ApplicationsGeneratorSet:
+                inherent: Quantum: QuantumProperties:
+                    runtime_features: RuntimeFeatureFlags(0x0)
+                    value_kind: Element(Static)
+                dynamic_param_applications: <empty>"#]],
+    );
+}