feat: integrate compute_output_permutation into _optimize_unitary

Extract with `up_to_perm=True` so that basic_optimization runs on a
SWAP-free circuit, then prepend the output permutation as SWAP gate
objects (1 gate each instead of 3 CNOTs).

Author: dlyongemallo

test/test_zxpass.py

@@ -694,6 +694,75 @@ def test_compute_output_permutation_non_bijective() -> None:
         compute_output_permutation(g)
 
 
+def test_permutation_swaps_in_pipeline() -> None:
+    """Test that _optimize_unitary uses SWAP gates (not CNOT decompositions).
+
+    After extraction with up_to_perm=True, the output permutation should be
+    prepended as SWAP gate objects.  Each SWAP counts as one gate instead of
+    three CNOTs, improving the gate-count comparison.
+    """
+    from zxpass.zxpass import _optimize_unitary  # pylint: disable=import-outside-toplevel
+
+    # Build a circuit that produces a non-trivial output permutation.
+    c = zx.Circuit(4)
+    c.add_gate("CNOT", 0, 1)
+    c.add_gate("CNOT", 1, 2)
+    c.add_gate("HAD", 0)
+    c.add_gate("CNOT", 2, 3)
+    c.add_gate("HAD", 3)
+    c.add_gate("CNOT", 3, 0)
+    c.add_gate("CNOT", 0, 2)
+    c.add_gate("HAD", 1)
+    c.add_gate("CNOT", 2, 1)
+    c.add_gate("CNOT", 1, 3)
+    c.add_gate("HAD", 2)
+    c.add_gate("CNOT", 3, 2)
+
+    optimized = _optimize_unitary(c)
+
+    # If the fallback didn't trigger, verify no CNOT-decomposed SWAPs leaked.
+    if len(optimized.gates) < len(c.gates):
+        # SWAPs should be at the beginning (prepended for the permutation).
+        for i, g in enumerate(optimized.gates):
+            if g.qasm_name != "swap":
+                # All subsequent gates should be non-SWAP.
+                rest = optimized.gates[i:]
+                assert all(r.qasm_name != "swap" for r in rest), (
+                    "SWAP gates should only appear at the beginning"
+                )
+                break
+
+
+def test_permutation_to_swaps_correctness() -> None:
+    """Test that _permutation_to_swaps produces a correct SWAP sequence."""
+    from zxpass.zxpass import _permutation_to_swaps  # pylint: disable=import-outside-toplevel
+
+    # Identity permutation: no SWAPs needed.
+    assert not _permutation_to_swaps({0: 0, 1: 1, 2: 2})
+
+    # Simple transposition.
+    swaps = _permutation_to_swaps({0: 1, 1: 0})
+    assert len(swaps) == 1
+    assert set(swaps[0]) == {0, 1}
+
+    # 3-cycle: needs 2 transpositions.
+    perm = {0: 1, 1: 2, 2: 0}
+    swaps = _permutation_to_swaps(perm)
+    # Verify the SWAPs implement the permutation.
+    state = list(range(3))
+    for i, j in swaps:
+        state[i], state[j] = state[j], state[i]
+    assert state == [perm[k] for k in range(3)]
+
+    # Larger permutation with multiple cycles.
+    perm = {0: 2, 1: 0, 2: 1, 3: 4, 4: 3}
+    swaps = _permutation_to_swaps(perm)
+    state = list(range(5))
+    for i, j in swaps:
+        state[i], state[j] = state[j], state[i]
+    assert state == [perm[k] for k in range(5)]
+
+
 def test_post_extraction_cleanup() -> None:
     """Test that ``_optimize_unitary`` applies ``basic_optimization`` after extraction.
 
@@ -738,32 +807,35 @@ def test_post_extraction_cleanup() -> None:
 
 
 def test_post_extraction_cleanup_equivalence() -> None:
-    """Test that the post-extraction cleanup preserves circuit equivalence.
-
-    Runs multiple circuits through the full ZXPass pipeline and verifies
-    statevector equivalence after the basic_optimization post-pass.
-    """
-    circuits = []
-
-    # Bell state preparation with extra gates.
+    """Test that post-extraction cleanup and permutation integration preserve equivalence."""
     qc1 = QuantumCircuit(3)
     qc1.h(0)
     qc1.cx(0, 1)
     qc1.cx(1, 2)
     qc1.h(2)
     qc1.cx(2, 0)
-    circuits.append(qc1)
 
-    # Multi-CX circuit.
-    qc2 = QuantumCircuit(4)
-    for i in range(3):
+    qc2 = QuantumCircuit(5)
+    for i in range(5):
+        qc2.h(i)
+    for i in range(4):
         qc2.cx(i, i + 1)
-    qc2.h(0)
-    qc2.cx(3, 0)
-    qc2.h(1)
-    circuits.append(qc2)
-
-    for qc in circuits:
+    qc2.cx(4, 0)
+    for i in range(5):
+        qc2.rz(0.3 * (i + 1), i)
+
+    qc3 = QuantumCircuit(4)
+    qc3.h(0)
+    qc3.cx(0, 1)
+    qc3.h(1)
+    qc3.cx(1, 2)
+    qc3.h(2)
+    qc3.cx(2, 3)
+    qc3.cx(3, 0)
+    qc3.cx(2, 1)
+    qc3.cx(1, 0)
+
+    for qc in [qc1, qc2, qc3]:
         assert _run_zxpass(qc), f"Equivalence failed for circuit:\n{qc}"
 
 

zxpass/zxpass.py

@@ -177,21 +177,52 @@ def compute_output_permutation(g: Any) -> Dict[int, int]:
     return perm
 
 
+def _permutation_to_swaps(perm: Dict[int, int]) -> List[Tuple[int, int]]:
+    """Decompose a permutation into transpositions (SWAP pairs).
+
+    Uses cycle decomposition.  The returned list, applied left to right,
+    implements the forward permutation (wire *j* ends up holding the state
+    of input qubit ``perm[j]``).
+    """
+    n = len(perm)
+    current = [perm[i] for i in range(n)]
+    swaps: List[Tuple[int, int]] = []
+    for i in range(n):
+        while current[i] != i:
+            j = current[i]
+            swaps.append((i, j))
+            current[i], current[j] = current[j], current[i]
+    swaps.reverse()
+    return swaps
+
+
 def _optimize_unitary(c: zx.Circuit) -> zx.Circuit:
     """Optimise a purely unitary PyZX circuit using full_reduce and extraction.
 
-    After extraction, ``basic_optimization`` converts HAD-CZ-HAD sequences to
-    CNOTs and cancels redundant single-qubit gates.  If the result still has at
-    least as many gates as the original circuit, the original is returned
-    unchanged to avoid regressions on small circuits with compact multi-qubit
-    gates (e.g. Toffoli, Fredkin). The comparison counts PyZX gate objects
-    directly; since ``_recover_dag`` emits one Qiskit op per PyZX gate, this
-    matches the Qiskit-side ``size()`` that downstream passes see.
+    Extracts with ``up_to_perm=True`` so that ``basic_optimization`` runs on a
+    circuit free of SWAP-decomposition clutter.  The output permutation is then
+    prepended as SWAP gates (each counting as one gate rather than three CNOTs),
+    giving a fairer gate-count comparison against the original circuit.  If the
+    result still has at least as many gates as the original, the original is
+    returned unchanged to avoid regressions on small circuits with compact
+    multi-qubit gates (e.g. Toffoli, Fredkin). The comparison counts PyZX gate
+    objects directly; since ``_recover_dag`` emits one Qiskit op per PyZX gate,
+    this matches the Qiskit-side ``size()`` that downstream passes see.
     """
     g = c.to_graph()
     zx.simplify.full_reduce(g)
-    optimized = zx.extract.extract_circuit(g)
+    optimized = zx.extract.extract_circuit(g, up_to_perm=True)
+    perm = compute_output_permutation(g)
     optimized = basic_optimization(optimized.to_basic_gates(), do_swaps=False)
+    # Prepend SWAP gates for the output permutation.
+    swap_pairs = _permutation_to_swaps(perm)
+    if swap_pairs:
+        with_perm = zx.Circuit(c.qubits)
+        for i, j in swap_pairs:
+            with_perm.add_gate(SWAP(i, j))
+        for gate in optimized.gates:
+            with_perm.add_gate(gate)
+        optimized = with_perm
     # TODO: Consider a two-axis comparison keyed primarily on 2-qubit gate
     # count (``twoqubitcount()``), with total gate count as a tiebreaker. The
     # 2-qubit count is the dominant hardware cost and is naturally apples-to-