Reverting the Allocate() change. (#857)

microsoft · Oct 19, 2021 · e8c2ed0 · e8c2ed0
1 parent 76dddfd
commit e8c2ed0
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 9 deletions.
diff --git a/src/Simulation/Common/QubitManager.cs b/src/Simulation/Common/QubitManager.cs
@@ -321,13 +321,13 @@ public void Disable(IQArray<Qubit> qubitsToDisable)
         /// Allocates a qubit.
         /// Returns null if the qubit cannot be allocated.
         /// </summary>
-        protected virtual Qubit Allocate(bool usedOnlyForBorrowing)
+        protected virtual Qubit? Allocate(bool usedOnlyForBorrowing)
         {
             if (free == None)
             {
                 if (!MayExtendCapacity)
                 {
-                    throw new NotEnoughQubits(1, this.FreeQubitsCount);
+                    return null;
                 }
 
                 long oldNumQubits = NumQubits;
@@ -398,7 +398,12 @@ protected virtual Qubit Allocate(bool usedOnlyForBorrowing)
         /// </summary>
         public Qubit Allocate()
         {
-            return Allocate(usedOnlyForBorrowing: false);
+            Qubit? qb = Allocate(usedOnlyForBorrowing: false);
+            if (qb == null)
+            {
+                throw new NotEnoughQubits(1, this.FreeQubitsCount);
+            }
+            return qb;
         }
 
         /// <summary>
@@ -424,7 +429,15 @@ public IQArray<Qubit> Allocate(long numToAllocate)
             }
             for (int i = 0; i < numToAllocate; i++)
             {
-                Qubit allocated = Allocate(usedOnlyForBorrowing: false);
+                Qubit? allocated = Allocate(usedOnlyForBorrowing: false);
+                if (allocated == null)
+                {
+                    for (int k = 0; k < i; k++)
+                    {
+                        Release(result[k], wasUsedOnlyForBorrowing: false);
+                    }
+                    throw new NotEnoughQubits(numToAllocate, this.FreeQubitsCount);
+                }
                 result.Modify(i, allocated);
             }
 
@@ -581,7 +594,15 @@ internal IQArray<Qubit> Borrow(long numToBorrow, HashSet<Qubit> qubitsInUse)
             { // Not enough qubits to borrow. Allocate what was not borrowed.
                 for (long i = numBorrowed; i < numToBorrow; i++)
                 {
-                    Qubit allocated = Allocate(usedOnlyForBorrowing: true);
+                    Qubit? allocated = Allocate(usedOnlyForBorrowing: true);
+                    if (allocated == null)
+                    {
+                        for (long k = numBorrowed; k < i; k++)
+                        {
+                            Release(borrowed[(int)k], wasUsedOnlyForBorrowing: true);
+                        }
+                        throw new NotEnoughQubits(numToBorrow, numBorrowed + this.FreeQubitsCount);
+                    }
                     borrowed.Modify(i, allocated);
                 }
             }

diff --git a/src/Simulation/Simulators/CommonNativeSimulator/QubitManager.cs b/src/Simulation/Simulators/CommonNativeSimulator/QubitManager.cs
@@ -42,11 +42,14 @@ public override Qubit CreateQubitObject(long id)
                 return new QSimQubit((int)id, Simulator);
             }
 
-            protected override Qubit Allocate(bool usedOnlyForBorrowing)
+            protected override Qubit? Allocate(bool usedOnlyForBorrowing)
             {
-                Qubit qubit = base.Allocate(usedOnlyForBorrowing);
-                Debug.Assert(Simulator != null);
-                Simulator.AllocateOne((uint)qubit.Id); 
+                Qubit? qubit = base.Allocate(usedOnlyForBorrowing);
+                if (qubit != null) 
+                {
+                    Debug.Assert(Simulator != null);
+                    Simulator.AllocateOne((uint)qubit.Id); 
+                }
                 return qubit;
             }