Improve interface for bloqs with specialized single-qubit-controlled versions

qualtran/_infra/controlled.py

@@ -23,6 +23,7 @@
     Sequence,
     Tuple,
     TYPE_CHECKING,
+    TypeAlias,
     Union,
 )
 
@@ -45,6 +46,9 @@
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
 
+ControlBit: TypeAlias = int
+"""A control bit, either 0 or 1."""
+
 
 def _cvs_convert(
     cvs: Union[
@@ -250,6 +254,18 @@ def from_cirq_cv(
             bloq_cvs.append(curr_cvs)
         return CtrlSpec(tuple(qdtypes), tuple(bloq_cvs))
 
+    def get_single_ctrl_bit(self) -> ControlBit:
+        """If controlled by a single qubit, return the control bit, otherwise raise"""
+        if self.num_qubits != 1:
+            raise ValueError(f"expected a single qubit control, got {self.num_qubits}")
+
+        (qdtype,) = self.qdtypes
+        (cv,) = self.cvs
+        (idx,) = Register('', qdtype, cv.shape).all_idxs()
+        (control_bit,) = qdtype.to_bits(cv[idx])
+
+        return int(control_bit)
+
 
 class AddControlledT(Protocol):
     """The signature for the `add_controlled` callback part of `ctrl_system`.

qualtran/_infra/controlled_test.py

@@ -100,6 +100,26 @@ def test_ctrl_spec_to_cirq_cv_roundtrip():
         assert CtrlSpec.from_cirq_cv(cirq_cv, qdtypes=ctrl_spec.qdtypes, shapes=ctrl_spec.shapes)
 
 
+@pytest.mark.parametrize(
+    "ctrl_spec", [CtrlSpec(), CtrlSpec(cvs=[1]), CtrlSpec(cvs=np.atleast_2d([1]))]
+)
+def test_ctrl_spec_single_bit_one(ctrl_spec: CtrlSpec):
+    assert ctrl_spec.get_single_ctrl_bit() == 1
+
+
+@pytest.mark.parametrize(
+    "ctrl_spec", [CtrlSpec(cvs=0), CtrlSpec(cvs=[0]), CtrlSpec(cvs=np.atleast_2d([0]))]
+)
+def test_ctrl_spec_single_bit_zero(ctrl_spec: CtrlSpec):
+    assert ctrl_spec.get_single_ctrl_bit() == 0
+
+
+@pytest.mark.parametrize("ctrl_spec", [CtrlSpec(cvs=[1, 1]), CtrlSpec(qdtypes=QUInt(2), cvs=0)])
+def test_ctrl_spec_single_bit_raises(ctrl_spec: CtrlSpec):
+    with pytest.raises(ValueError):
+        ctrl_spec.get_single_ctrl_bit()
+
+
 def _test_cirq_equivalence(bloq: Bloq, gate: 'cirq.Gate'):
     import cirq
 

qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py

@@ -20,17 +20,27 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from qualtran import bloq_example, BloqDocSpec, BQUInt, QAny, QBit, Register, Signature
+from qualtran import (
+    AddControlledT,
+    Bloq,
+    bloq_example,
+    BloqDocSpec,
+    BQUInt,
+    CtrlSpec,
+    QAny,
+    QBit,
+    Register,
+    Signature,
+)
 from qualtran._infra.gate_with_registers import total_bits
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
 from qualtran.bloqs.basic_gates import CSwap
 from qualtran.bloqs.multiplexers.apply_gate_to_lth_target import ApplyGateToLthQubit
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
 from qualtran.bloqs.multiplexers.selected_majorana_fermion import SelectedMajoranaFermion
 
 
 @attrs.frozen
-class SelectHubbard(SelectOracle, SpecializedSingleQubitControlledExtension):  # type: ignore[misc]
+class SelectHubbard(SelectOracle):
     r"""The SELECT operation optimized for the 2D Hubbard model.
 
     In contrast to SELECT for an arbitrary chemistry Hamiltonian, we:
@@ -180,6 +190,21 @@ def __str__(self):
             return f'C{s}'
         return s
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(
+            ctrl_spec=ctrl_spec,
+            current_ctrl_bit=self.control_val,
+            bloq_without_ctrl=attrs.evolve(self, control_val=None),
+            get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (
+                attrs.evolve(self, control_val=cv),
+                'control',
+            ),
+        )
+
 
 @bloq_example
 def _sel_hubb() -> SelectHubbard:

qualtran/bloqs/chemistry/sparse/select_bloq.py

@@ -16,11 +16,23 @@
 from functools import cached_property
 from typing import Dict, Optional, Tuple, TYPE_CHECKING
 
+import attrs
 import cirq
 from attrs import frozen
 
-from qualtran import bloq_example, BloqBuilder, BloqDocSpec, BQUInt, QAny, QBit, Register, SoquetT
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
+from qualtran import (
+    AddControlledT,
+    Bloq,
+    bloq_example,
+    BloqBuilder,
+    BloqDocSpec,
+    BQUInt,
+    CtrlSpec,
+    QAny,
+    QBit,
+    Register,
+    SoquetT,
+)
 from qualtran.bloqs.basic_gates import SGate
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
 from qualtran.bloqs.multiplexers.selected_majorana_fermion import SelectedMajoranaFermion
@@ -30,7 +42,7 @@
 
 
 @frozen
-class SelectSparse(SpecializedSingleQubitControlledExtension, SelectOracle):  # type: ignore[misc]
+class SelectSparse(SelectOracle):
     r"""SELECT oracle for the sparse Hamiltonian.
 
     Implements the two applications of Fig. 13.
@@ -157,6 +169,21 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         c_maj_y = SelectedMajoranaFermion(sel_pa, target_gate=cirq.Y)
         return {SGate(): 1, maj_x: 1, c_maj_x: 1, maj_y: 1, c_maj_y: 1}
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(
+            ctrl_spec=ctrl_spec,
+            current_ctrl_bit=self.control_val,
+            bloq_without_ctrl=attrs.evolve(self, control_val=None),
+            get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (
+                attrs.evolve(self, control_val=cv),
+                'control',
+            ),
+        )
+
 
 @bloq_example
 def _sel_sparse() -> SelectSparse:

qualtran/bloqs/chemistry/thc/select_bloq.py

@@ -20,18 +20,19 @@
 from attrs import evolve, frozen
 
 from qualtran import (
+    AddControlledT,
     Bloq,
     bloq_example,
     BloqBuilder,
     BloqDocSpec,
     BQUInt,
+    CtrlSpec,
     QAny,
     QBit,
     Register,
     Signature,
     SoquetT,
 )
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
 from qualtran.bloqs.basic_gates import CSwap, Toffoli, XGate
 from qualtran.bloqs.chemistry.black_boxes import ApplyControlledZs
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
@@ -120,7 +121,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
 
 
 @frozen
-class SelectTHC(SpecializedSingleQubitControlledExtension, SelectOracle):  # type: ignore[misc]
+class SelectTHC(SelectOracle):
     r"""SELECT for THC Hamiltonian.
 
     Args:
@@ -313,6 +314,18 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
 
         return out_soqs
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(
+            ctrl_spec=ctrl_spec,
+            current_ctrl_bit=self.control_val,
+            bloq_without_ctrl=evolve(self, control_val=None),
+            get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (evolve(self, control_val=cv), 'control'),
+        )
+
 
 @bloq_example
 def _thc_sel() -> SelectTHC:

qualtran/bloqs/for_testing/random_select_and_prepare.py

@@ -14,13 +14,13 @@
 from functools import cached_property
 from typing import Iterator, Optional, Tuple
 
+import attrs
 import cirq
 import numpy as np
 from attrs import frozen
 from numpy.typing import NDArray
 
-from qualtran import BloqBuilder, BQUInt, QBit, Register, SoquetT
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
+from qualtran import AddControlledT, Bloq, BloqBuilder, BQUInt, CtrlSpec, QBit, Register, SoquetT
 from qualtran.bloqs.block_encoding.lcu_block_encoding import SelectBlockEncoding
 from qualtran.bloqs.for_testing.matrix_gate import MatrixGate
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
@@ -84,7 +84,7 @@ def alphas(self):
 
 
 @frozen
-class TestPauliSelectOracle(SpecializedSingleQubitControlledExtension, SelectOracle):  # type: ignore[misc]
+class TestPauliSelectOracle(SelectOracle):  # type: ignore[misc]
     r"""Paulis acting on $m$ qubits, controlled by an $n$-qubit register.
 
     Given $2^n$ multi-qubit-Paulis (acting on $m$ qubits) $U_j$,
@@ -149,6 +149,21 @@ def decompose_from_registers(
                 op = op.controlled_by(*quregs['control'], control_values=[self.control_val])
             yield op
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(
+            ctrl_spec=ctrl_spec,
+            current_ctrl_bit=self.control_val,
+            bloq_without_ctrl=attrs.evolve(self, control_val=None),
+            get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (
+                attrs.evolve(self, control_val=cv),
+                'control',
+            ),
+        )
+
 
 def random_qubitization_walk_operator(
     select_bitsize: int, target_bitsize: int, *, random_state: np.random.RandomState