Merge branch 'microsoft:main' into Kata-issue-591-add-examples-edit-run

compiler/qsc_eval/src/intrinsic/tests.rs

@@ -1502,3 +1502,74 @@ fn two_qubit_rotation_neg_inf_error() {
         &expect!["invalid rotation angle: -inf"],
     );
 }
+
+#[test]
+fn stop_counting_operation_before_start_fails() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            Std.Diagnostics.StopCountingOperation(I);
+        }"},
+        &expect!["callable not counted"],
+    );
+}
+
+#[test]
+fn stop_counting_function_before_start_fails() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            function Foo() : Unit {}
+            Std.Diagnostics.StopCountingFunction(Foo);
+        }"},
+        &expect!["callable not counted"],
+    );
+}
+
+#[test]
+fn start_counting_operation_called_twice_before_stop_fails() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            Std.Diagnostics.StartCountingOperation(I);
+            Std.Diagnostics.StartCountingOperation(I);
+        }"},
+        &expect!["callable already counted"],
+    );
+}
+
+#[test]
+fn start_counting_function_called_twice_before_stop_fails() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            function Foo() : Unit {}
+            Std.Diagnostics.StartCountingFunction(Foo);
+            Std.Diagnostics.StartCountingFunction(Foo);
+        }"},
+        &expect!["callable already counted"],
+    );
+}
+
+#[test]
+fn stop_counting_qubits_before_start_fails() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            Std.Diagnostics.StopCountingQubits();
+        }"},
+        &expect!["qubits not counted"],
+    );
+}
+
+#[test]
+fn start_counting_qubits_called_twice_before_stop_fails() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            Std.Diagnostics.StartCountingQubits();
+            Std.Diagnostics.StartCountingQubits();
+        }"},
+        &expect!["qubits already counted"],
+    );
+}

compiler/qsc_eval/src/lib.rs

@@ -43,7 +43,7 @@ use qsc_fir::fir::{
 use qsc_fir::ty::Ty;
 use qsc_lowerer::map_fir_package_to_hir;
 use rand::{rngs::StdRng, SeedableRng};
-use rustc_hash::FxHashSet;
+use rustc_hash::{FxHashMap, FxHashSet};
 use std::ops;
 use std::{
     cell::RefCell,
@@ -61,6 +61,18 @@ pub enum Error {
     #[diagnostic(code("Qsc.Eval.ArrayTooLarge"))]
     ArrayTooLarge(#[label("this array has too many items")] PackageSpan),
 
+    #[error("callable already counted")]
+    #[diagnostic(help(
+        "counting for a given callable must be stopped before it can be started again"
+    ))]
+    #[diagnostic(code("Qsc.Eval.CallableAlreadyCounted"))]
+    CallableAlreadyCounted(#[label] PackageSpan),
+
+    #[error("callable not counted")]
+    #[diagnostic(help("counting for a given callable must be started before it can be stopped"))]
+    #[diagnostic(code("Qsc.Eval.CallableNotCounted"))]
+    CallableNotCounted(#[label] PackageSpan),
+
     #[error("invalid array length: {0}")]
     #[diagnostic(code("Qsc.Eval.InvalidArrayLength"))]
     InvalidArrayLength(i64, #[label("cannot be used as a length")] PackageSpan),
@@ -105,6 +117,16 @@ pub enum Error {
     #[diagnostic(code("Qsc.Eval.QubitUniqueness"))]
     QubitUniqueness(#[label] PackageSpan),
 
+    #[error("qubits already counted")]
+    #[diagnostic(help("counting for qubits must be stopped before it can be started again"))]
+    #[diagnostic(code("Qsc.Eval.QubitsAlreadyCounted"))]
+    QubitsAlreadyCounted(#[label] PackageSpan),
+
+    #[error("qubits not counted")]
+    #[diagnostic(help("counting for qubits must be started before it can be stopped"))]
+    #[diagnostic(code("Qsc.Eval.QubitsNotCounted"))]
+    QubitsNotCounted(#[label] PackageSpan),
+
     #[error("qubits are not separable")]
     #[diagnostic(help("subset of qubits provided as arguments must not be entangled with any qubits outside of the subset"))]
     #[diagnostic(code("Qsc.Eval.QubitsNotSeparable"))]
@@ -150,6 +172,8 @@ impl Error {
     pub fn span(&self) -> &PackageSpan {
         match self {
             Error::ArrayTooLarge(span)
+            | Error::CallableAlreadyCounted(span)
+            | Error::CallableNotCounted(span)
             | Error::DivZero(span)
             | Error::EmptyRange(span)
             | Error::IndexOutOfRange(_, span)
@@ -160,6 +184,8 @@ impl Error {
             | Error::InvalidNegativeInt(_, span)
             | Error::OutputFail(span)
             | Error::QubitUniqueness(span)
+            | Error::QubitsAlreadyCounted(span)
+            | Error::QubitsNotCounted(span)
             | Error::QubitsNotSeparable(span)
             | Error::RangeStepZero(span)
             | Error::ReleasedQubitNotZero(_, span)
@@ -435,6 +461,8 @@ struct Scope {
     frame_id: usize,
 }
 
+type CallableCountKey = (StoreItemId, bool, bool);
+
 pub struct State {
     exec_graph_stack: Vec<ExecGraph>,
     idx: u32,
@@ -446,6 +474,8 @@ pub struct State {
     call_stack: CallStack,
     current_span: Span,
     rng: RefCell<StdRng>,
+    call_counts: FxHashMap<CallableCountKey, i64>,
+    qubit_counter: Option<QubitCounter>,
 }
 
 impl State {
@@ -466,6 +496,8 @@ impl State {
             call_stack: CallStack::default(),
             current_span: Span::default(),
             rng,
+            call_counts: FxHashMap::default(),
+            qubit_counter: None,
         }
     }
 
@@ -962,9 +994,20 @@ impl State {
 
         let spec = spec_from_functor_app(functor);
         match &callee.implementation {
+            CallableImpl::Intrinsic if is_counting_call(&callee.name.name) => {
+                self.push_frame(Vec::new().into(), callee_id, functor);
+
+                let val = self.counting_call(&callee.name.name, arg, arg_span)?;
+
+                self.set_val_register(val);
+                self.leave_frame();
+                Ok(())
+            }
             CallableImpl::Intrinsic => {
                 self.push_frame(Vec::new().into(), callee_id, functor);
 
+                self.increment_call_count(callee_id, functor);
+
                 let name = &callee.name.name;
                 let val = intrinsic::call(
                     name,
@@ -981,6 +1024,14 @@ impl State {
                         callee_span,
                     ));
                 }
+
+                // If qubit counting is enabled, update the qubit counter when the intrinsic for allocation is called.
+                if let (Some(counter), "__quantum__rt__qubit_allocate", Value::Qubit(q)) =
+                    (&mut self.qubit_counter, name.as_ref(), &val)
+                {
+                    counter.allocated(q.0);
+                }
+
                 self.set_val_register(val);
                 self.leave_frame();
                 Ok(())
@@ -995,6 +1046,7 @@ impl State {
                 .expect("missing specialization should be a compilation error");
                 self.push_frame(spec_decl.exec_graph.clone(), callee_id, functor);
                 self.push_scope(env);
+                self.increment_call_count(callee_id, functor);
 
                 self.bind_args_for_spec(
                     env,
@@ -1436,6 +1488,58 @@ impl State {
             span,
         }
     }
+
+    fn counting_call(&mut self, name: &str, arg: Value, span: PackageSpan) -> Result<Value, Error> {
+        let counting_key = |arg: Value| match arg {
+            Value::Closure(closure) => make_counting_key(closure.id, closure.functor),
+            Value::Global(id, functor) => make_counting_key(id, functor),
+            _ => panic!("value should be callable"),
+        };
+        match name {
+            "StartCountingOperation" | "StartCountingFunction" => {
+                if self.call_counts.insert(counting_key(arg), 0).is_some() {
+                    Err(Error::CallableAlreadyCounted(span))
+                } else {
+                    Ok(Value::unit())
+                }
+            }
+            "StopCountingOperation" | "StopCountingFunction" => {
+                if let Some(count) = self.call_counts.remove(&counting_key(arg)) {
+                    Ok(Value::Int(count))
+                } else {
+                    Err(Error::CallableNotCounted(span))
+                }
+            }
+            "StartCountingQubits" => {
+                if self
+                    .qubit_counter
+                    .replace(QubitCounter::default())
+                    .is_some()
+                {
+                    Err(Error::QubitsAlreadyCounted(span))
+                } else {
+                    Ok(Value::unit())
+                }
+            }
+            "StopCountingQubits" => {
+                if let Some(qubit_counter) = self.qubit_counter.take() {
+                    Ok(Value::Int(qubit_counter.into_count()))
+                } else {
+                    Err(Error::QubitsNotCounted(span))
+                }
+            }
+            _ => panic!("unknown counting call"),
+        }
+    }
+
+    fn increment_call_count(&mut self, callee_id: StoreItemId, functor: FunctorApp) {
+        if let Some(count) = self
+            .call_counts
+            .get_mut(&make_counting_key(callee_id, functor))
+        {
+            *count += 1;
+        }
+    }
 }
 
 pub fn are_ctls_unique(ctls: &[Value], tup: &Value) -> bool {
@@ -1925,3 +2029,37 @@ fn is_updatable_in_place(env: &Env, expr: &Expr) -> (bool, bool) {
         _ => (false, false),
     }
 }
+
+fn is_counting_call(name: &str) -> bool {
+    matches!(
+        name,
+        "StartCountingOperation"
+            | "StopCountingOperation"
+            | "StartCountingFunction"
+            | "StopCountingFunction"
+            | "StartCountingQubits"
+            | "StopCountingQubits"
+    )
+}
+
+fn make_counting_key(id: StoreItemId, functor: FunctorApp) -> CallableCountKey {
+    (id, functor.adjoint, functor.controlled > 0)
+}
+
+#[derive(Default)]
+struct QubitCounter {
+    seen: FxHashSet<usize>,
+    count: i64,
+}
+
+impl QubitCounter {
+    fn allocated(&mut self, qubit: usize) {
+        if self.seen.insert(qubit) {
+            self.count += 1;
+        }
+    }
+
+    fn into_count(self) -> i64 {
+        self.count
+    }
+}

compiler/qsc_eval/src/tests.rs

@@ -3724,7 +3724,7 @@ fn controlled_operation_with_duplicate_controls_fails() {
                                 1,
                             ),
                             item: LocalItemId(
-                                124,
+                                130,
                             ),
                         },
                         caller: PackageId(
@@ -3774,7 +3774,7 @@ fn controlled_operation_with_target_in_controls_fails() {
                                 1,
                             ),
                             item: LocalItemId(
-                                124,
+                                130,
                             ),
                         },
                         caller: PackageId(

compiler/qsc_fir/src/fir.rs

@@ -324,7 +324,7 @@ pub enum Global<'a> {
 }
 
 /// A unique identifier for an item within a package store.
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq, Hash, Eq)]
 pub struct StoreItemId {
     /// The package ID.
     pub package: PackageId,