From d379dea7876f43778738f011860aa782f2e49ef3 Mon Sep 17 00:00:00 2001
From: Jim Garrison <garrison@ibm.com>
Date: Thu, 25 Apr 2024 11:29:29 -0400
Subject: [PATCH] Early terminate from loops when generating cutting
 experiments (#565)

This is a follow-up to #556.
---
 .../cutting/cutting_experiments.py            |  8 ++++++-
 test/cutting/test_cutting_experiments.py      | 23 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/circuit_knitting/cutting/cutting_experiments.py b/circuit_knitting/cutting/cutting_experiments.py
index 4411d3687..e5634103a 100644
--- a/circuit_knitting/cutting/cutting_experiments.py
+++ b/circuit_knitting/cutting/cutting_experiments.py
@@ -409,7 +409,7 @@ def _remove_resets_in_zero_state(
         circuit = circuit.copy()
 
     # Keep up with which qubits have at least one non-reset instruction
-    active_qubits = set()
+    active_qubits: set[int] = set()
     remove_ids = []
     for i, inst in enumerate(circuit.data):
         qargs = [circuit.find_bit(q).index for q in inst.qubits]
@@ -419,6 +419,9 @@ def _remove_resets_in_zero_state(
         else:
             for q in qargs:
                 active_qubits.add(q)
+            # Early terminate once all qubits have become active
+            if len(active_qubits) == circuit.num_qubits:
+                break
 
     for i in sorted(remove_ids, reverse=True):
         del circuit.data[i]
@@ -446,6 +449,9 @@ def _remove_final_resets(
         else:
             for q in qargs:
                 qubit_ended.discard(q)
+            # Early terminate once all qubits have been touched
+            if not qubit_ended:
+                break
 
     for i in sorted(remove_ids, reverse=True):
         del circuit.data[i]
diff --git a/test/cutting/test_cutting_experiments.py b/test/cutting/test_cutting_experiments.py
index 4da5edca6..47c6e15c0 100644
--- a/test/cutting/test_cutting_experiments.py
+++ b/test/cutting/test_cutting_experiments.py
@@ -275,7 +275,7 @@ def test_optimize_single_reset_in_diff_qubits(self):
         _remove_final_resets(circuit)
         self.assertEqual(expected, circuit)
 
-    def test_optimize_single_reset(self):
+    def test_optimize_single_final_reset(self):
         """Remove a single final reset
         qr0:--[H]--|0>--   ==>    qr0:--[H]--
         """
@@ -291,6 +291,27 @@ def test_optimize_single_reset(self):
 
         self.assertEqual(expected, circuit)
 
+    def test_optimize_single_final_reset_2(self):
+        """Remove a single final reset on two qubits
+        qr0:--[H]--|0>--   ==>    qr0:--[H]-------
+            --[X]--[S]--              --[X]--[S]--
+        """
+        qr = QuantumRegister(2, "qr")
+        circuit = QuantumCircuit(qr)
+        circuit.h(0)
+        circuit.x(1)
+        circuit.reset(0)
+        circuit.s(1)
+
+        expected = QuantumCircuit(qr)
+        expected.h(0)
+        expected.x(1)
+        expected.s(1)
+
+        _remove_final_resets(circuit)
+
+        self.assertEqual(expected, circuit)
+
     def test_dont_optimize_non_final_reset(self):
         """Do not remove reset if not final instruction
         qr0:--|0>--[H]--   ==>    qr0:--|0>--[H]--