From c3fad3c48be8d40be6d19e85985453eb9e14d7a3 Mon Sep 17 00:00:00 2001
From: seyon <seyon.sivarajah@cambridgequantum.com>
Date: Thu, 17 Feb 2022 14:27:24 +0000
Subject: [PATCH] [feature] add auto_squash_pass in vein of auto_rebase

---
 pytket/docs/changelog.rst           |  4 ++-
 pytket/pytket/passes/__init__.py    |  2 +-
 pytket/pytket/passes/auto_rebase.py | 12 +++++++
 pytket/tests/transform_test.py      | 53 +++++++++++++++++++++++++++--
 4 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/pytket/docs/changelog.rst b/pytket/docs/changelog.rst
index 6a6e205864..887a38b946 100644
--- a/pytket/docs/changelog.rst
+++ b/pytket/docs/changelog.rst
@@ -19,7 +19,9 @@ API changes:
 Minor new features:
 
 * Add ``delay_measures`` option to ``DefaultMappingPass``.
-* New ``pytket.passes.auto_rebase_pass`` which attempts to construct a rebase pass given a target gate set from known decompositions.
+* New ``pytket.passes.auto_rebase_pass`` and ``pytket.passes.auto_squash_pass``
+  which attempt to construct rebase and squash passess given a target gate set from known
+  decompositions.
 
 0.19.1 (February 2022)
 ----------------------
diff --git a/pytket/pytket/passes/__init__.py b/pytket/pytket/passes/__init__.py
index aa68496e71..b617a99486 100644
--- a/pytket/pytket/passes/__init__.py
+++ b/pytket/pytket/passes/__init__.py
@@ -17,4 +17,4 @@
 from pytket._tket.passes import *  # type: ignore
 
 from .script import compilation_pass_from_script, compilation_pass_grammar
-from .auto_rebase import auto_rebase_pass
+from .auto_rebase import auto_rebase_pass, auto_squash_pass
diff --git a/pytket/pytket/passes/auto_rebase.py b/pytket/pytket/passes/auto_rebase.py
index af3946a4b0..0dca9c5734 100644
--- a/pytket/pytket/passes/auto_rebase.py
+++ b/pytket/pytket/passes/auto_rebase.py
@@ -100,3 +100,15 @@ def auto_rebase_pass(gateset: Set[OpType]) -> RebaseCustom:
     return RebaseCustom(
         gateset, get_cx_decomposition(gateset), get_TK1_decomposition_function(gateset)
     )
+
+
+def auto_squash_pass(gateset: Set[OpType]) -> SquashCustom:
+    """Attempt to generate a squash pass automatically for the given target
+    single qubit gateset.
+
+    :param gateset: Available single qubit gateset
+    :type gateset: Set[OpType]
+    :return: Squash to target gateset
+    :rtype: SquashCustom
+    """
+    return SquashCustom(gateset, get_TK1_decomposition_function(gateset))
diff --git a/pytket/tests/transform_test.py b/pytket/tests/transform_test.py
index 32ae7c72f7..0dd44d20c4 100644
--- a/pytket/tests/transform_test.py
+++ b/pytket/tests/transform_test.py
@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 from pathlib import Path
+from typing import List
 from pytket.circuit import Circuit, OpType, PauliExpBox  # type: ignore
 from pytket._tket.circuit import _library  # type: ignore
 from pytket.pauli import Pauli  # type: ignore
-from pytket.passes import RemoveRedundancies, KAKDecomposition, ThreeQubitSquash, CommuteThroughMultis, PauliSquash, FullPeepholeOptimise, GlobalisePhasedX, RebaseCustom  # type: ignore
-from pytket.passes import auto_rebase_pass
+from pytket.passes import RemoveRedundancies, KAKDecomposition, SquashCustom, CommuteThroughMultis, PauliSquash, FullPeepholeOptimise, GlobalisePhasedX, RebaseCustom  # type: ignore
+from pytket.passes import auto_rebase_pass, auto_squash_pass
 from pytket.passes.auto_rebase import _CX_CIRCS, NoAutoRebase
 from pytket.predicates import CompilationUnit  # type: ignore
 from pytket.transform import Transform, CXConfigType, PauliSynthStrat  # type: ignore
@@ -780,6 +782,53 @@ def test_auto_rebase() -> None:
     assert "TK1" in str(cx_err.value)
 
 
+def test_auto_squash() -> None:
+    pass_params = [
+        ({OpType.Rz, OpType.Rx}, _library._TK1_to_RzRx),
+        (
+            {OpType.Rz, OpType.SX},
+            _library._TK1_to_RzSX,
+        ),
+        (
+            {OpType.T, OpType.Rz, OpType.H},
+            _library._TK1_to_RzH,
+        ),
+        (
+            {OpType.T, OpType.Rz, OpType.H},
+            _library._TK1_to_RzH,
+        ),
+        (
+            {OpType.PhasedX, OpType.Rz},
+            _library._TK1_to_PhasedXRz,
+        ),
+        (
+            {OpType.TK1, OpType.U3},
+            _library._TK1_to_TK1,
+        ),
+    ]
+
+    for gateset, TK1_func in pass_params:
+        circ = Circuit(1)
+        for gate in itertools.islice(itertools.cycle(gateset), 5):
+            # make a sequence of 5 gates from gateset to make sure squash does
+            # something
+            params: List[float] = []
+            while True:
+                try:
+                    circ.add_gate(gate, params, [0])
+                    break
+                except (RuntimeError, TypeError):
+                    params.append(0.1)
+        squash = auto_squash_pass(gateset)
+        assert squash.to_dict() == SquashCustom(gateset, TK1_func).to_dict()
+
+        assert squash.apply(circ)
+
+    with pytest.raises(NoAutoRebase) as tk_err:
+        _ = auto_squash_pass({OpType.H, OpType.T})
+    assert "TK1" in str(tk_err.value)
+
+
 if __name__ == "__main__":
     test_remove_redundancies()
     test_reduce_singles()