Fixes #367 (#541)

* Fixes #367.

* Apply suggestions from code review

Co-authored-by: Mariia Mykhailova <mamykhai@microsoft.com>

* Addressing Mariia's comments.

Co-authored-by: Mariia Mykhailova <mamykhai@microsoft.com>

Standard/src/Diagnostics/Allows.qs

@@ -7,6 +7,11 @@ namespace Microsoft.Quantum.Diagnostics {
     /// Between a call to this operation and its adjoint, asserts that
     /// a given operation is called at most a certain number of times.
     ///
+    /// Operation calls are considered, if they contain the the specified
+    /// variant.  For example, if `op` is `X`, `Adjoint X` or `Controlled X`
+    /// are also counted, but if `op` is `Controlled X`, only `Controlled X`
+    /// or `Controlled Adjoint X` are counted.
+    ///
     /// # Input
     /// ## nTimes
     /// The maximum number of times that `op` may be called.
@@ -19,14 +24,27 @@ namespace Microsoft.Quantum.Diagnostics {
     /// The following snippet will fail when executed on machines which
     /// support this diagnostic:
     /// ```qsharp
-    /// using (register = Qubit[4]) {
-    ///     within {
-    ///         AllowAtMostNCallsCA(3, H, "Too many calls to H.");
-    ///     } apply {
-    ///         // Fails since this calls H four times, rather than the
-    ///         // allowed maximum of three.
-    ///         ApplyToEach(H, register);
-    ///     }
+    /// within {
+    ///     AllowAtMostNCallsCA(3, H, "Too many calls to H.");
+    /// } apply {
+    ///     use register = Qubit[4];
+    ///     // Fails since this calls H four times, rather than the
+    ///     // allowed maximum of three.
+    ///     ApplyToEach(H, register);
+    /// }
+    /// ```
+    ///
+    /// Another example illustrates how restricted calls are handled.
+    /// ```qsharp
+    /// within {
+    ///     // Both tests will pass in this case
+    ///     AllowAtMostNCallsCA(1, Controlled H, "Too many calls to Controlled H.");
+    ///     AllowAtMostNCallsCA(2, H, "Too many calls to H or Controlled H.");
+    /// } apply {
+    ///     use (a, b) = (Qubit(), Qubit());
+    ///     H(a);
+    ///     Controlled H([a], b);
+    ///     ResetAll([a, b]);
     /// }
     /// ```
     ///

Standard/src/Diagnostics/Emulation/AllowOperationCalls.cs

@@ -55,6 +55,15 @@ namespace Microsoft.Quantum.Diagnostics
                 bool IsSelf(ICallable callable) =>
                     callable.FullName == "Microsoft.Quantum.Diagnostics.AllowAtMostNCallsCA";
 
+                // Partial ordering on two variants; returns true, if `lhs` is included in `rhs`.
+                // For example, Body <= Body, and Body <= Adjoint, but Controlled is not less than or equal Adjoint.
+                bool LessThanOrEqual(OperationFunctor lhs, OperationFunctor rhs) => lhs switch {
+                        OperationFunctor.ControlledAdjoint => rhs == OperationFunctor.ControlledAdjoint,
+                        OperationFunctor.Controlled => rhs == OperationFunctor.ControlledAdjoint || rhs == OperationFunctor.Controlled,
+                        OperationFunctor.Adjoint => rhs == OperationFunctor.ControlledAdjoint || rhs == OperationFunctor.Adjoint,
+                        _ => true // OperationFunctor.Body
+                };
+
                 // Record whether or not the condition checked by this allow
                 // has failed, so that we can property unwind in the endOperation
                 // handler below.
@@ -65,7 +74,11 @@ namespace Microsoft.Quantum.Diagnostics
                     {
                         if (IsSelf(callable)) return;
                         callStack = callStack.Push(callable.FullName);
-                        if (callable.FullName == op.FullName)
+                        // `callable` is callable we just entered on the call stack, `op` is the callable
+                        // that we are monitoring with AllowAtMostNCallsCA.  We only increment the counter,
+                        // if both callables have the same fully qualified name and if the variant of `op`
+                        // is less restrictive than the one of `callable`.
+                        if (callable.FullName == op.FullName && LessThanOrEqual(op.Variant, callable.Variant))
                         {
                             callSites = callSites.Add(callStack);
                             if (callSites.Count > nTimes)

Standard/tests/Diagnostics/AllowTests.qs

@@ -64,6 +64,41 @@ namespace Microsoft.Quantum.Tests {
         }
     }
 
+    @Diag.Test("QuantumSimulator")
+    operation CheckAllowNCallsDistinguishControlled() : Unit {
+        within {
+            Diag.AllowAtMostNCallsCA(1, Controlled X, "Too many calls to Controlled X.");
+        } apply {
+            use (q1, q2) = (Qubit(), Qubit());
+            // Should use one Controlled X, exactly as many times as allowed.
+            // X will not be accounted for.
+            X(q2);
+            Controlled X([q1], q2);
+            X(q2);
+        }
+    }
+
+    @Diag.Test("QuantumSimulator")
+    operation CheckAllowNCallsDistinguishAdjoint() : Unit {
+        within {
+            Diag.AllowAtMostNCallsCA(4, S, "Too many calls to S.");
+            Diag.AllowAtMostNCallsCA(2, Adjoint S, "Too many calls to Adjoint S.");
+            Diag.AllowAtMostNCallsCA(2, Controlled S, "Too many calls to Controlled S.");
+            Diag.AllowAtMostNCallsCA(1, Adjoint Controlled S, "Too many calls to Adjoint Controlled S.");
+        } apply {
+            use (q1, q2) = (Qubit(), Qubit());
+            // Should use two Adjoint S (one via Adjoint Controlled S), but exactly
+            // one Adjoint Controlled S.
+            S(q1);
+            Controlled S([q1], q2);
+            Adjoint Controlled S([q2], q1);
+            Adjoint S(q1);
+
+            Reset(q1);
+            Reset(q2);
+        }
+    }
+
     @Diag.Test("ToffoliSimulator")
     operation CheckAllowNQubitsWithNestedCalls() : Unit {
         // Here, the total number of allocated qubits exceeds our policy,