Override ctrl system for Toffoli (#1552)

Author: anurudhp

qualtran/bloqs/basic_gates/toffoli.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 import itertools
 from functools import cached_property
-from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
+from typing import cast, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 from attrs import frozen
@@ -25,6 +25,7 @@
     BloqDocSpec,
     CompositeBloq,
     Connection,
+    CtrlSpec,
     DecomposeTypeError,
     QBit,
     Register,
@@ -35,6 +36,7 @@
     import cirq
     import quimb.tensor as qtn
 
+    from qualtran import AddControlledT, BloqBuilder, SoquetT
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.drawing import WireSymbol
     from qualtran.simulation.classical_sim import ClassicalValT
@@ -127,6 +129,29 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
             return ModPlus()
         raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.basic_gates import CNOT
+        from qualtran.bloqs.mcmt import ControlledViaAnd
+
+        if ctrl_spec != CtrlSpec():
+            return super().get_ctrl_system(ctrl_spec)
+
+        cc_cnot = ControlledViaAnd(CNOT(), CtrlSpec(cvs=[1, 1]))
+
+        def add_controlled(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            (new_ctrl,) = ctrl_soqs
+            ctrl0, ctrl1 = cast(NDArray, in_soqs.pop('ctrl'))
+
+            (new_ctrl, ctrl0), ctrl1, target = bb.add(
+                cc_cnot, ctrl2=np.array([new_ctrl, ctrl0]), ctrl=ctrl1, target=in_soqs.pop('target')
+            )
+
+            return [new_ctrl], [ctrl0, ctrl1, target]
+
+        return cc_cnot, add_controlled
+
 
 @bloq_example
 def _toffoli() -> Toffoli:

qualtran/bloqs/basic_gates/toffoli_test.py

@@ -16,10 +16,12 @@
 import cirq
 import numpy as np
 
-from qualtran import BloqBuilder
+from qualtran import BloqBuilder, CtrlSpec
 from qualtran.bloqs.basic_gates import Toffoli, ZeroState
 from qualtran.bloqs.basic_gates.toffoli import _toffoli
+from qualtran.bloqs.mcmt import And
 from qualtran.drawing.musical_score import Circle, ModPlus
+from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
 from qualtran.testing import assert_wire_symbols_match_expected
 
 
@@ -85,3 +87,18 @@ def test_toffoli_tensors():
     unitary = tof.tensor_contract()
     cirq_unitary = cirq.unitary(cirq.TOFFOLI)
     np.testing.assert_allclose(cirq_unitary, unitary)
+
+
+def test_ctrl_toffoli_cost():
+    ctrl_tof = Toffoli().controlled()
+
+    _, sigma = ctrl_tof.call_graph()
+    assert sigma == {Toffoli(): 1, And(): 1, And().adjoint(): 1}
+
+    gc = get_cost_value(ctrl_tof, QECGatesCost())
+    assert gc == GateCounts(and_bloq=1, toffoli=1, clifford=1, measurement=1)
+
+    cc_tof = Toffoli().controlled(CtrlSpec(cvs=[1, 1]))
+
+    _, sigma = cc_tof.call_graph()
+    assert sigma == {Toffoli(): 1, And(): 2, And().adjoint(): 2}