From 388b29deba615ea14c39054d19b4b05eaf3ed0ea Mon Sep 17 00:00:00 2001
From: Daniel Strano <stranoj@gmail.com>
Date: Thu, 23 Feb 2023 07:48:42 -0500
Subject: [PATCH] Optimize ProbMask()

---
 src/pinvoke_api.cpp           | 3 ---
 src/qinterface/qinterface.cpp | 4 ++++
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/pinvoke_api.cpp b/src/pinvoke_api.cpp
index 25e925826..e83cc5c6f 100644
--- a/src/pinvoke_api.cpp
+++ b/src/pinvoke_api.cpp
@@ -2050,9 +2050,6 @@ MICROSOFT_QUANTUM_DECL double PermutationProb(
 
     try {
         QInterfacePtr simulator = simulators[sid];
-        if ((simulator->GetMaxQPower() - 1U) == mask) {
-            return (double)simulator->ProbAll(perm);
-        }
         return (double)simulator->ProbMask(mask, perm);
     } catch (...) {
         simulatorErrors[sid] = 1;
diff --git a/src/qinterface/qinterface.cpp b/src/qinterface/qinterface.cpp
index 8a9d4d9f9..c341c97a9 100644
--- a/src/qinterface/qinterface.cpp
+++ b/src/qinterface/qinterface.cpp
@@ -270,6 +270,10 @@ real1_f QInterface::ProbReg(bitLenInt start, bitLenInt length, bitCapInt permuta
 /// Returns probability of permutation of the mask
 real1_f QInterface::ProbMask(bitCapInt mask, bitCapInt permutation)
 {
+    if ((maxQPower - 1U) == mask) {
+        return ProbAll(permutation);
+    }
+
     real1 prob = ZERO_R1;
     for (bitCapInt lcv = 0U; lcv < maxQPower; ++lcv) {
         if ((lcv & mask) == permutation) {