From ae3d9c98d9f5cce1e509dc285bf77866dcb60b7c Mon Sep 17 00:00:00 2001
From: Thien Nguyen <nguyentm@ornl.gov>
Date: Tue, 17 Aug 2021 10:30:06 -0400
Subject: [PATCH] Related to https://github.com/eclipse/xacc/issues/478

The procedure was not properly taking into account non-IBM gateset, i.e., 2-q gates may not be CNOT.

Signed-off-by: Thien Nguyen <nguyentm@ornl.gov>
---
 .../optimizers/simple/CircuitOptimizer.cpp    | 15 ++++++++++++++
 .../simple/tests/CircuitOptimizerTester.cpp   | 20 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/quantum/plugins/optimizers/simple/CircuitOptimizer.cpp b/quantum/plugins/optimizers/simple/CircuitOptimizer.cpp
index f009b996f..110c3e25d 100644
--- a/quantum/plugins/optimizers/simple/CircuitOptimizer.cpp
+++ b/quantum/plugins/optimizers/simple/CircuitOptimizer.cpp
@@ -678,6 +678,21 @@ bool CircuitOptimizer::tryRotationMergingUsingPhasePolynomials(std::shared_ptr<C
             }
           }
           else if (instruction->bits().size() == 2) {
+            // If this gate not involved the qubits that we are checking
+            if (!container::contains(qubits, instruction->bits()[0]) &&
+                !container::contains(qubits, instruction->bits()[1])) {
+              // Skip
+              continue;
+            }
+
+            // This 2-q gate **involves** at least one qubit:
+            if (instruction->name() != "CNOT") {
+              // Not a CNOT: we need to terminate, hence prune the subcircuit.
+              // Move the outer loop index to after this gate.
+              i = idx + 1;
+              break;
+            }
+
             assert(instruction->name() == "CNOT");
             // If the control is *outside* the boundary, we need to terminate, hence prune the subcircuit.
             const auto controlIdx = instruction->bits()[0];
diff --git a/quantum/plugins/optimizers/simple/tests/CircuitOptimizerTester.cpp b/quantum/plugins/optimizers/simple/tests/CircuitOptimizerTester.cpp
index 8d828585e..196d03867 100644
--- a/quantum/plugins/optimizers/simple/tests/CircuitOptimizerTester.cpp
+++ b/quantum/plugins/optimizers/simple/tests/CircuitOptimizerTester.cpp
@@ -818,6 +818,26 @@ measure q[3] -> c[3];
   EXPECT_EQ(truthTableBefore, truthTableAfter);
 }
 
+TEST(CircuitOptimizerTester, checkCZ) {
+  xacc::set_verbose(true);
+  auto compiler = xacc::getService<xacc::Compiler>("xasm");
+  auto program = compiler
+                     ->compile(R"(__qpu__ void testCz(qbit q) {
+        CNOT(q[0], q[1]);
+        CZ(q[0], q[1]);
+        Measure(q[0]);
+        Measure(q[1]);
+    })")
+                     ->getComposites()[0];
+
+  const auto before_xasm_str = program->toString();
+  auto optimizer = xacc::getService<IRTransformation>("circuit-optimizer");
+  optimizer->apply(program, nullptr);
+  std::cout << "FINAL CIRCUIT:\n" << program->toString() << "\n";
+  const auto after_xasm_str = program->toString();
+  EXPECT_EQ(before_xasm_str, after_xasm_str);
+}
+
 int main(int argc, char **argv) {
   xacc::Initialize(argc, argv);
   ::testing::InitGoogleTest(&argc, argv);