From 5fcff6a8129964fabeb2a799b6ffb8ee932422c0 Mon Sep 17 00:00:00 2001
From: WrathfulSpatula <stranoj@gmail.com>
Date: Thu, 14 Nov 2024 09:26:02 -0500
Subject: [PATCH] Ability to disable QUnit fidelity exceptions

---
 include/qunit.hpp | 15 ++++++++++++---
 src/qunit.cpp     | 43 ++++++++++++++-----------------------------
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/include/qunit.hpp b/include/qunit.hpp
index b9c2d9c2d..6d3ae3680 100644
--- a/include/qunit.hpp
+++ b/include/qunit.hpp
@@ -80,6 +80,17 @@ class QUnit : public QParity, public QInterface {
         return abs(at) / PI_R1;
     }
 
+    void CheckFidelity()
+    {
+#if ENABLE_ENV_VARS
+        if ((!(bool)getenv("QRACK_DISABLE_QUNIT_FIDELITY_GUARD")) && (logFidelity <= FIDELITY_MIN)) {
+#else
+        if (logFidelity <= FIDELITY_MIN) {
+#endif
+            throw std::runtime_error("QUnit fidelity is effectively 0!");
+        }
+    }
+
 public:
     QUnit(std::vector<QInterfaceEngine> eng, bitLenInt qBitCount, const bitCapInt& initState = ZERO_BCI,
         qrack_rand_gen_ptr rgp = nullptr, const complex& phaseFac = CMPLX_DEFAULT_ARG, bool doNorm = false,
@@ -727,9 +738,7 @@ class QUnit : public QParity, public QInterface {
             SeparateBit(true, qubit);
         }
 
-        if (logFidelity <= FIDELITY_MIN) {
-            throw std::runtime_error("QUnit fidelity is effectively 0!");
-        }
+        CheckFidelity();
     }
 
     void TransformX2x2(const complex* mtrxIn, complex* mtrxOut)
diff --git a/src/qunit.cpp b/src/qunit.cpp
index f1bdbf6a9..336e66778 100644
--- a/src/qunit.cpp
+++ b/src/qunit.cpp
@@ -182,9 +182,7 @@ void QUnit::SetQuantumState(const complex* inputState)
             shard.amp1 = shard.amp0 / abs(shard.amp0);
             shard.amp0 = ZERO_R1;
         }
-        if (logFidelity <= FIDELITY_MIN) {
-            throw std::runtime_error("QUnit fidelity is effectively 0!");
-        }
+        CheckFidelity();
         return;
     }
 
@@ -408,6 +406,10 @@ bool QUnit::Detach(bitLenInt start, bitLenInt length, QUnitPtr dest, bool isTry,
 QInterfacePtr QUnit::EntangleInCurrentBasis(
     std::vector<bitLenInt*>::iterator first, std::vector<bitLenInt*>::iterator last)
 {
+    if (std::distance(last, first) == 0U) {
+        throw std::invalid_argument("QUnit::EntangleInCurrentBasis() vector argument cannot be empty!");
+    }
+
     const QUnitPtr backupCopy = std::dynamic_pointer_cast<QUnit>(Copy());
 
     for (auto bit = first; bit < last; ++bit) {
@@ -743,10 +745,7 @@ bool QUnit::TrySeparate(bitLenInt qubit)
     ShardAI(qubit, azimuth, inclination);
 
     logFidelity += (double)log(clampProb(1.0 - oneMinR / 2));
-
-    if (logFidelity <= FIDELITY_MIN) {
-        throw std::runtime_error("QUnit fidelity is effectively 0!");
-    }
+    CheckFidelity();
 
     return true;
 }
@@ -977,9 +976,7 @@ real1_f QUnit::ProbBase(bitLenInt qubit)
             amps[0U] = ZERO_CMPLX;
         }
 
-        if (logFidelity <= FIDELITY_MIN) {
-            throw std::runtime_error("QUnit fidelity is effectively 0!");
-        }
+        CheckFidelity();
 
         shard.amp0 = amps[0U];
         shard.amp1 = amps[1U];
@@ -1010,9 +1007,7 @@ real1_f QUnit::ProbBase(bitLenInt qubit)
         SeparateBit(true, qubit);
     }
 
-    if (logFidelity <= FIDELITY_MIN) {
-        throw std::runtime_error("QUnit fidelity is effectively 0!");
-    }
+    CheckFidelity();
 
     return clampProb(norm(shard.amp1));
 }
@@ -1821,10 +1816,6 @@ void QUnit::EitherISwap(bitLenInt qubit1, bitLenInt qubit2, bool isInverse)
                 logFidelity += log(ONE_R1_F - pLo);
             }
 
-            if (logFidelity <= FIDELITY_MIN) {
-                throw std::runtime_error("QUnit fidelity is effectively 0!");
-            }
-
             if (isInverse) {
                 Swap(qubit1, qubit2);
             } else {
@@ -1834,6 +1825,7 @@ void QUnit::EitherISwap(bitLenInt qubit1, bitLenInt qubit2, bool isInverse)
 
             return;
         }
+
         if (isInverse) {
             unit->IISwap(shard1.mapped, shard2.mapped);
         } else {
@@ -1846,6 +1838,7 @@ void QUnit::EitherISwap(bitLenInt qubit1, bitLenInt qubit2, bool isInverse)
             TrySeparate(qubit1);
             TrySeparate(qubit2);
         }
+
         return;
     }
 
@@ -2783,9 +2776,7 @@ void QUnit::ApplyEitherControlled(
         } else {
             logFidelity += log(ONE_R1_F - p);
         }
-        if (logFidelity <= FIDELITY_MIN) {
-            throw std::runtime_error("QUnit fidelity is effectively 0!");
-        }
+        CheckFidelity();
 
         return;
     }
@@ -4020,9 +4011,7 @@ void QUnit::ApplyBuffer(PhaseShardPtr phaseShard, bitLenInt control, bitLenInt t
         bool pState = abs(pHi - HALF_R1) >= abs(pLo - HALF_R1);
 
         logFidelity += angleFrac(polarBottom) * log(pState ? pHi : (ONE_R1_F - pLo));
-        if (logFidelity <= FIDELITY_MIN) {
-            throw std::runtime_error("QUnit fidelity is effectively 0!");
-        }
+        CheckFidelity();
 
         if (pState) {
             if (!ptHi) {
@@ -4041,9 +4030,7 @@ void QUnit::ApplyBuffer(PhaseShardPtr phaseShard, bitLenInt control, bitLenInt t
         pState = abs(pHi - HALF_R1) >= abs(pLo - HALF_R1);
 
         logFidelity += angleFrac(polarTop) * log(pState ? pHi : (ONE_R1_F - pLo));
-        if (logFidelity <= FIDELITY_MIN) {
-            throw std::runtime_error("QUnit fidelity is effectively 0!");
-        }
+        CheckFidelity();
 
         if (!pState) {
             if (ptHi) {
@@ -4066,9 +4053,7 @@ void QUnit::ApplyBuffer(PhaseShardPtr phaseShard, bitLenInt control, bitLenInt t
             pState = abs(pHi - HALF_R1) >= abs(pLo - HALF_R1);
 
             logFidelity += log(pState ? pHi : (ONE_R1_F - pLo));
-            if (logFidelity <= FIDELITY_MIN) {
-                throw std::runtime_error("QUnit fidelity is effectively 0!");
-            }
+            CheckFidelity();
 
             if (pState) {
                 if (!ptHi) {