From 7fd37fd1c7bfe87b621c4c1a72876f969a2d18ae Mon Sep 17 00:00:00 2001
From: Neal Erickson <neal.erickson@quantinuum.com>
Date: Tue, 10 Oct 2023 10:37:53 -0600
Subject: [PATCH] Adding test for partial barriers

---
 tests/data/qasm/barrier_complex.qasm | 27 ++++++++++++
 tests/sample_data.py                 |  1 +
 tests/test_sharder.py                | 64 ++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+)
 create mode 100644 tests/data/qasm/barrier_complex.qasm

diff --git a/tests/data/qasm/barrier_complex.qasm b/tests/data/qasm/barrier_complex.qasm
new file mode 100644
index 0000000..bee4d0b
--- /dev/null
+++ b/tests/data/qasm/barrier_complex.qasm
@@ -0,0 +1,27 @@
+OPENQASM 2.0;
+include "hqslib1.inc";
+
+qreg q[4];
+creg c[4];
+
+h q[0];
+h q[1];
+h q[2];
+h q[3];
+c[3] = 1;
+
+barrier q[0], q[1], c[3];
+
+CX q[0], q[1];
+
+measure q[0]->c[0];
+
+h q[2];
+h q[3];
+x q[3];
+
+barrier q[2], q[3];
+
+CX q[2], q[3];
+
+measure q[2]->c[2];
diff --git a/tests/sample_data.py b/tests/sample_data.py
index 56300f1..983e8be 100644
--- a/tests/sample_data.py
+++ b/tests/sample_data.py
@@ -12,6 +12,7 @@ class QasmFiles(Enum):
     baby_with_rollup = 5
     simple_cond = 6
     cond_classical = 7
+    barrier_complex = 8
 
 
 def get_qasm_as_circuit(qasm_file: QasmFiles) -> Circuit:
diff --git a/tests/test_sharder.py b/tests/test_sharder.py
index caada75..de8bdf8 100644
--- a/tests/test_sharder.py
+++ b/tests/test_sharder.py
@@ -127,3 +127,67 @@ def test_simple_conditional(self) -> None:
         assert s2_sub_cmds[0].op.type == OpType.Conditional
         assert cast(Conditional, s2_sub_cmds[0].op).op.type == OpType.H
         assert s2_sub_cmds[0].qubits == [circuit.qubits[0]]
+
+    def test_complex_barriers(self) -> None:
+        circuit = get_qasm_as_circuit(QasmFiles.barrier_complex)
+        sharder = Sharder(circuit)
+        shards = sharder.shard()
+
+        assert len(shards) == 7
+
+        # shard 0: [], c[3] = 1
+        assert shards[0].primary_command.op.type == OpType.SetBits
+        assert len(shards[0].sub_commands.items()) == 0
+        assert shards[0].qubits_used == set()
+        assert shards[0].bits_written == {circuit.bits[3]}
+        assert shards[0].bits_read == {circuit.bits[3]}  # bits written are always read
+        assert shards[0].depends_upon == set()
+
+        # shard 1: [h q[0]; h q[1];] barrier q[0], q[1], c[3];
+        assert shards[1].primary_command.op.type == OpType.Barrier
+        assert len(shards[1].sub_commands.items()) == 2
+        # TODO: sub commands
+        assert shards[1].qubits_used == {circuit.qubits[0], circuit.qubits[1]}
+        assert shards[1].bits_written == {circuit.bits[3]}
+        assert shards[1].bits_read == {circuit.bits[3]}
+        assert shards[1].depends_upon == {shards[0].ID}
+
+        # shard 2: [] CX q[0], q[1];
+        assert shards[2].primary_command.op.type == OpType.CX
+        assert len(shards[2].sub_commands.items()) == 0
+        assert shards[2].qubits_used == {circuit.qubits[0], circuit.qubits[1]}
+        assert shards[2].bits_written == set()
+        assert shards[2].bits_read == set()
+        assert shards[2].depends_upon == {shards[1].ID}
+
+        # shard 3: measure q[0]->c[0];
+        assert shards[3].primary_command.op.type == OpType.Measure
+        assert len(shards[3].sub_commands.items()) == 0
+        assert shards[3].qubits_used == {circuit.qubits[0]}
+        assert shards[3].bits_written == {circuit.bits[0]}
+        assert shards[3].bits_read == {circuit.bits[0]}
+        assert shards[3].depends_upon == {shards[2].ID, shards[1].ID}
+
+        # shard 4: [] barrier q[2], q[3];
+        assert shards[4].primary_command.op.type == OpType.Barrier
+        assert len(shards[4].sub_commands.items()) == 2
+        assert shards[4].qubits_used == {circuit.qubits[2], circuit.qubits[3]}
+        assert shards[4].bits_written == set()
+        assert shards[4].bits_read == set()
+        assert shards[4].depends_upon == set()
+
+        # shard 5: [] CX q[2], q[3];
+        assert shards[5].primary_command.op.type == OpType.CX
+        assert len(shards[5].sub_commands.items()) == 0
+        assert shards[5].qubits_used == {circuit.qubits[2], circuit.qubits[3]}
+        assert shards[5].bits_written == set()
+        assert shards[5].bits_read == set()
+        assert shards[5].depends_upon == {shards[4].ID}
+
+        # shard 6: measure q[2]->c[2];
+        assert shards[6].primary_command.op.type == OpType.Measure
+        assert len(shards[6].sub_commands.items()) == 0
+        assert shards[6].qubits_used == {circuit.qubits[2]}
+        assert shards[6].bits_written == {circuit.bits[2]}
+        assert shards[6].bits_read == {circuit.bits[2]}
+        assert shards[6].depends_upon == {shards[5].ID, shards[4].ID}