From cc285673678de371f7c281685bdce320a1ad516b Mon Sep 17 00:00:00 2001
From: "Michael A. Perlin" <mika.perlin@gmail.com>
Date: Thu, 2 Jan 2025 15:23:00 -0500
Subject: [PATCH] Permute all variables of the Steane code (#114)

* permute all variables of the Steane code

* more descriptive docstring

* compute the swap of CRegs

* minor cleanup

* fix bug

* formatting fix
---
 .../src/pecos/qeclib/steane/steane_class.py   | 23 +++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/python/quantum-pecos/src/pecos/qeclib/steane/steane_class.py b/python/quantum-pecos/src/pecos/qeclib/steane/steane_class.py
index 16da9920..826b36b2 100644
--- a/python/quantum-pecos/src/pecos/qeclib/steane/steane_class.py
+++ b/python/quantum-pecos/src/pecos/qeclib/steane/steane_class.py
@@ -285,7 +285,7 @@ def nonft_t_tel(self, aux: Steane):
             aux.cx(self),
             self.mz(self.t_meas),
             If(self.t_meas == 1).Then(aux.x(), aux.sz()),
-            Permute(self.d, aux.d),
+            self.permute(aux),
         )
 
     def t_tel(
@@ -308,7 +308,7 @@ def t_tel(
             aux.cx(self),
             self.mz(self.t_meas),
             If(self.t_meas == 1).Then(aux.x(), aux.sz()),  # SZ/S correction.
-            Permute(self.d, aux.d),
+            self.permute(aux),
         )
 
     def nonft_tdg_tel(self, aux: Steane):
@@ -326,7 +326,7 @@ def nonft_tdg_tel(self, aux: Steane):
             aux.cx(self),
             self.mz(self.tdg_meas),
             If(self.tdg_meas == 1).Then(aux.x(), aux.szdg()),
-            Permute(self.d, aux.d),
+            self.permute(aux),
         )
 
     def tdg_tel(
@@ -349,7 +349,7 @@ def tdg_tel(
             aux.cx(self),
             self.mz(self.tdg_meas),
             If(self.t_meas == 1).Then(aux.x(), aux.szdg()),  # SZdg/Sdg correction.
-            Permute(self.d, aux.d),
+            self.permute(aux),
         )
 
     def t_cor(
@@ -495,3 +495,18 @@ def qec(self, flag: Bit | None = None):
         if flag is not None:
             block.extend(If(self.flags != 0).Then(flag.set(1)))
         return block
+
+    def permute(self, other: Steane):
+        """Permute this code block (including both quantum and classical registers) with another."""
+        block = Block(
+            Permute(self.d, other.d),
+            Permute(self.a, other.a),
+        )
+        for var_a, var_b in zip(self.vars, other.vars):
+            if isinstance(var_a, CReg):
+                block.extend(
+                    var_a.set(var_a ^ var_b),
+                    var_b.set(var_b ^ var_a),
+                    var_a.set(var_a ^ var_b),
+                )
+        return block