From 4ebc5819b0c2d49c8d2144008d3bf3f04f34edc6 Mon Sep 17 00:00:00 2001
From: WrathfulSpatula <stranoj@gmail.com>
Date: Sun, 10 Nov 2024 16:38:53 -0500
Subject: [PATCH] More QUnit classical shadows

---
 src/qunit.cpp | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/qunit.cpp b/src/qunit.cpp
index 001b34569..f6f713c79 100644
--- a/src/qunit.cpp
+++ b/src/qunit.cpp
@@ -2678,7 +2678,25 @@ void QUnit::ApplyEitherControlled(
         ebits[i] = &allBitsMapped[i];
     }
 
-    QInterfacePtr unit = EntangleInCurrentBasis(ebits.begin(), ebits.end());
+    QInterfacePtr unit;
+    try {
+        unit = EntangleInCurrentBasis(ebits.begin(), ebits.end());
+    } catch (const bad_alloc& e) {
+        if (targets.size() > 1U) {
+            throw e;
+        }
+        // We overallocated; use a really primitive classical shadow, at least.
+        for (bitLenInt i = 0U; i < controlVec.size(); ++i) {
+            if (Prob(controlVec[i]) < (ONE_R1_F / 2)) {
+                // Any control has less than 50% probability to be measured as |1>:
+                // Do nothing.
+                return;
+            }
+            // Act the classical shadow of the gate payload.
+            cfn(unit, {});
+        }
+        return;
+    }
 
     for (size_t i = 0U; i < controlVec.size(); ++i) {
         bitLenInt& c = controlVec[i];