Support call tracking within Q# code

compiler/qsc_data_structures/src/functors.rs

@@ -8,7 +8,7 @@ use std::{
 };
 
 /// A functor application.
-#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash)]
 pub struct FunctorApp {
     /// An invocation is either adjoint or not, with each successive use of `Adjoint` functor switching
     /// between the two, so a bool is sufficient to track.

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,
@@ -449,6 +449,7 @@ pub struct State {
     call_stack: CallStack,
     current_span: Span,
     rng: RefCell<StdRng>,
+    call_counts: FxHashMap<(StoreItemId, FunctorApp), i64>,
 }
 
 impl State {
@@ -473,6 +474,7 @@ impl State {
             call_stack: CallStack::default(),
             current_span: Span::default(),
             rng,
+            call_counts: FxHashMap::default(),
         }
     }
 
@@ -974,9 +976,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);
+
+                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,
@@ -1007,6 +1020,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,
@@ -1448,6 +1462,31 @@ impl State {
             span,
         }
     }
+
+    fn counting_call(&mut self, name: &str, arg: Value) -> Value {
+        let callable = if let Value::Closure(closure) = arg {
+            (closure.id, closure.functor)
+        } else {
+            arg.unwrap_global()
+        };
+        match name {
+            "StartCountingOperation" | "StartCountingFunction" => {
+                self.call_counts.insert(callable, 0);
+                Value::unit()
+            }
+            "StopCountingOperation" | "StopCountingFunction" => {
+                let count = self.call_counts.remove(&callable).unwrap_or(-1);
+                Value::Int(count)
+            }
+            _ => panic!("unknown counting call"),
+        }
+    }
+
+    fn increment_call_count(&mut self, callee_id: StoreItemId, functor: FunctorApp) {
+        if let Some(count) = self.call_counts.get_mut(&(callee_id, functor)) {
+            *count += 1;
+        }
+    }
 }
 
 pub fn are_ctls_unique(ctls: &[Value], tup: &Value) -> bool {
@@ -1937,3 +1976,13 @@ 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"
+    )
+}

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,

library/qs_source/src/std/diagnostics.qs

@@ -149,5 +149,25 @@ namespace Microsoft.Quantum.Diagnostics {
         areEqual
     }
 
-    export DumpMachine, DumpRegister, CheckZero, CheckAllZero, Fact, CheckOperationsAreEqual;
+    @Config(Unrestricted)
+    operation StartCountingOperation<'In, 'Out>(callable : 'In => 'Out) : Unit {
+        body intrinsic;
+    }
+
+    @Config(Unrestricted)
+    operation StopCountingOperation<'In, 'Out>(callable : 'In => 'Out) : Int {
+        body intrinsic;
+    }
+
+    @Config(Unrestricted)
+    operation StartCountingFunction<'In, 'Out>(callable : 'In -> 'Out) : Unit {
+        body intrinsic;
+    }
+
+    @Config(Unrestricted)
+    operation StopCountingFunction<'In, 'Out>(callable : 'In -> 'Out) : Int {
+        body intrinsic;
+    }
+
+    export DumpMachine, DumpRegister, CheckZero, CheckAllZero, Fact, CheckOperationsAreEqual, StartCountingOperation, StopCountingOperation, StartCountingFunction, StopCountingFunction;
 }

library/src/tests/diagnostics.rs

@@ -40,3 +40,119 @@ fn check_operations_are_equal() {
         ),
     );
 }
+
+#[test]
+fn check_start_stop_counting_operation_called_3_times() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StartCountingOperation;
+            import Microsoft.Quantum.Diagnostics.StopCountingOperation;
+
+            operation op1() : Unit {}
+            operation op2() : Unit { op1(); }
+            StartCountingOperation(op1);
+            StartCountingOperation(op2);
+            op1(); op1(); op2();
+            (StopCountingOperation(op1), StopCountingOperation(op2))
+        }",
+        &Value::Tuple([Value::Int(3), Value::Int(1)].into()),
+    );
+}
+
+#[test]
+fn check_start_stop_counting_operation_called_0_times() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StartCountingOperation;
+            import Microsoft.Quantum.Diagnostics.StopCountingOperation;
+
+            operation op1() : Unit {}
+            operation op2() : Unit { op1(); }
+            StartCountingOperation(op1);
+            StartCountingOperation(op2);
+            (StopCountingOperation(op1), StopCountingOperation(op2))
+        }",
+        &Value::Tuple([Value::Int(0), Value::Int(0)].into()),
+    );
+}
+
+#[test]
+fn check_stop_counting_operation_without_start() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StopCountingOperation;
+
+            operation op1() : Unit {}
+            StopCountingOperation(op1)
+        }",
+        &Value::Int(-1),
+    );
+}
+
+#[test]
+fn check_counting_operation_differentiates_between_body_adj_ctl() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StartCountingOperation;
+            import Microsoft.Quantum.Diagnostics.StopCountingOperation;
+
+            operation op1() : Unit is Adj + Ctl {}
+            StartCountingOperation(op1);
+            StartCountingOperation(Adjoint op1);
+            StartCountingOperation(Controlled op1);
+            op1();
+            Adjoint op1(); Adjoint op1();
+            Controlled op1([], ()); Controlled op1([], ()); Controlled op1([], ());
+            (StopCountingOperation(op1), StopCountingOperation(Adjoint op1), StopCountingOperation(Controlled op1))
+        }",
+        &Value::Tuple([Value::Int(1), Value::Int(2), Value::Int(3)].into()),
+    );
+}
+
+#[test]
+fn check_start_stop_counting_function_called_3_times() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StartCountingFunction;
+            import Microsoft.Quantum.Diagnostics.StopCountingFunction;
+
+            function f1() : Unit {}
+            function f2() : Unit { f1(); }
+            StartCountingFunction(f1);
+            StartCountingFunction(f2);
+            f1(); f1(); f2();
+            (StopCountingFunction(f1), StopCountingFunction(f2))
+        }",
+        &Value::Tuple([Value::Int(3), Value::Int(1)].into()),
+    );
+}
+
+#[test]
+fn check_start_stop_counting_function_called_0_times() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StartCountingFunction;
+            import Microsoft.Quantum.Diagnostics.StopCountingFunction;
+
+            function f1() : Unit {}
+            function f2() : Unit { f1(); }
+            StartCountingFunction(f1);
+            StartCountingFunction(f2);
+            (StopCountingFunction(f1), StopCountingFunction(f2))
+        }",
+        &Value::Tuple([Value::Int(0), Value::Int(0)].into()),
+    );
+}
+
+#[test]
+fn check_stop_counting_function_without_start() {
+    test_expression(
+        "{
+            import Microsoft.Quantum.Diagnostics.StopCountingFunction;
+
+            function f1() : Unit {}
+            StopCountingFunction(f1)
+        }",
+        &Value::Int(-1),
+    );
+}