diff --git a/qualtran/bloqs/mcmt/multi_control_pauli.py b/qualtran/bloqs/mcmt/multi_control_pauli.py
index a64bb2cf26..2fc966ddda 100644
--- a/qualtran/bloqs/mcmt/multi_control_pauli.py
+++ b/qualtran/bloqs/mcmt/multi_control_pauli.py
@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
+import abc
 from functools import cached_property
 from typing import Dict, Set, Tuple, TYPE_CHECKING, Union
 
@@ -28,15 +28,17 @@
     CtrlSpec,
     DecomposeTypeError,
     GateWithRegisters,
+    QAny,
     QBit,
     Register,
     Signature,
+    Soquet,
     SoquetT,
 )
 from qualtran.bloqs.basic_gates import XGate
 from qualtran.bloqs.mcmt.and_bloq import _to_tuple_or_has_length, is_symbolic
 from qualtran.bloqs.mcmt.controlled_via_and import ControlledViaAnd
-from qualtran.symbolics import HasLength, SymbolicInt
+from qualtran.symbolics import HasLength, slen, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
@@ -184,26 +186,97 @@ def _ccpauli_symb() -> MultiControlPauli:
 
 
 @frozen
-class MultiControlX(MultiControlPauli):
+class _MultiControlPauli(Bloq):
+    """Abstract base class for a generalized implementation of Multi-Control Pauli"""
+
+    cvs: CtrlSpec
+
+    @property
+    @abc.abstractmethod
+    def _target_gate(self) -> cirq.Pauli:
+        ...
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature([*self.ctrl_regs, Register('target', QBit())])
+
+    @cached_property
+    def ctrl_regs(self) -> Tuple[Register, ...]:
+        ctrl_regs = []
+        for i, (dtype, shape) in enumerate(self.cvs.activation_function_dtypes()):
+            if is_symbolic(dtype.num_qubits) or dtype.num_qubits > 0:
+                ctrl_regs.append(Register(f'ctrl{i}_', dtype=dtype, shape=shape))
+        return tuple(ctrl_regs)
+
+    @cached_property
+    def flat_cvs(self) -> Union[HasLength, Tuple[int, ...]]:
+        return tuple(
+            b
+            for cvs, qdtype in zip(self.cvs.cvs, self.cvs.qdtypes)
+            for cv in tuple(cvs.reshape(-1))
+            for b in qdtype.to_bits(cv)
+        )
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str, 'SoquetT']:
+        flat_ctrls = []
+        for reg in self.ctrl_regs:
+            if reg.shape:
+                for soq in soqs[reg.name].reshape(-1):
+                    flat_ctrls += bb.split(soq).tolist() if reg.bitsize > 1 else [soq]
+            else:
+                soq = soqs[reg.name]
+                assert isinstance(soq, Soquet)
+                flat_ctrls += bb.split(soq).tolist() if reg.bitsize > 1 else [soq]
+        print(self.flat_cvs)
+        print(flat_ctrls)
+        flat_ctrls, target = bb.add(
+            MultiControlPauli(self.flat_cvs, self._target_gate),
+            controls=flat_ctrls,
+            target=soqs['target'],
+        )
+        ctrls = {}
+        st = 0
+        for reg in self.ctrl_regs:
+            if reg.shape:
+                curr_soqs = np.empty(reg.shape, dtype=object)
+                for idx in reg.all_idxs():
+                    if reg.bitsize > 1:
+                        curr_soqs[idx] = bb.join(flat_ctrls[st : st + reg.bitsize], dtype=reg.dtype)
+                    else:
+                        curr_soqs[idx] = flat_ctrls[st]
+                    st += reg.bitsize
+                ctrls[reg.name] = curr_soqs
+            else:
+                if reg.bitsize > 1:
+                    ctrls[reg.name] = bb.join(flat_ctrls[st : st + reg.bitsize], dtype=reg.dtype)
+                else:
+                    ctrls[reg.name] = flat_ctrls[st]
+                st += reg.bitsize
+        return ctrls | {'target': target}
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        return {(MultiControlPauli(self.flat_cvs, self._target_gate), 1)}
+
+
+@frozen
+class MultiControlX(_MultiControlPauli):
     r"""Implements multi-control, single-target X gate.
 
     See :class:`MultiControlPauli` for implementation and costs.
     """
-    target_gate: cirq.Pauli = field(init=False)
 
-    @target_gate.default
-    def _X(self):
+    @property
+    def _target_gate(self) -> cirq.Pauli:
         return cirq.X
 
 
 @frozen
-class MultiControlZ(MultiControlPauli):
+class MultiControlZ(_MultiControlPauli):
     r"""Implements multi-control, single-target Z gate.
 
     See :class:`MultiControlPauli` for implementation and costs.
     """
-    target_gate: cirq.Pauli = field(init=False)
 
-    @target_gate.default
-    def _Z(self):
+    @property
+    def _target_gate(self) -> cirq.Pauli:
         return cirq.Z
diff --git a/qualtran/bloqs/reflections/reflection_using_prepare.py b/qualtran/bloqs/reflections/reflection_using_prepare.py
index 3878edc742..9c3a916f8f 100644
--- a/qualtran/bloqs/reflections/reflection_using_prepare.py
+++ b/qualtran/bloqs/reflections/reflection_using_prepare.py
@@ -13,24 +13,25 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Iterator, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import Dict, Iterator, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
 import numpy as np
 from numpy.typing import NDArray
 
-from qualtran import Bloq, bloq_example, BloqDocSpec, CtrlSpec, QBit, Register, Signature
+from qualtran import Bloq, bloq_example, BloqDocSpec, CtrlSpec, QBit, Register, Side, Signature
 from qualtran._infra.gate_with_registers import GateWithRegisters, merge_qubits, total_bits
 from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
-from qualtran.bloqs.basic_gates.global_phase import GlobalPhase
-from qualtran.bloqs.basic_gates.x_basis import XGate
+from qualtran.bloqs.basic_gates import GlobalPhase, XGate, ZPowGate
 from qualtran.bloqs.mcmt import MultiControlZ
 from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity
+from qualtran.drawing import Circle, ModPlus, TextBox
 from qualtran.resource_counting.generalizers import ignore_split_join
-from qualtran.symbolics import HasLength, is_symbolic, SymbolicInt
+from qualtran.symbolics import HasLength, is_symbolic, pi, sarg, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
+    from qualtran import BloqBuilder, SoquetT
     from qualtran.bloqs.block_encoding.lcu_block_encoding import BlackBoxPrepare
     from qualtran.bloqs.state_preparation.prepare_base import PrepareOracle
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
@@ -78,7 +79,7 @@ class ReflectionUsingPrepare(GateWithRegisters, SpecializedSingleQubitControlled
     prepare_gate: Union['PrepareOracle', 'BlackBoxPrepare']
     control_val: Optional[int] = None
     global_phase: complex = 1
-    eps: float = 1e-11
+    eps: SymbolicFloat = 1e-11
 
     @cached_property
     def control_registers(self) -> Tuple[Register, ...]:
@@ -88,9 +89,13 @@ def control_registers(self) -> Tuple[Register, ...]:
     def selection_registers(self) -> Tuple[Register, ...]:
         return self.prepare_gate.selection_registers
 
+    @cached_property
+    def junk_registers(self) -> Tuple[Register, ...]:
+        return self.prepare_gate.junk_registers
+
     @cached_property
     def signature(self) -> Signature:
-        return Signature([*self.control_registers, *self.selection_registers])
+        return Signature([*self.control_registers, *self.selection_registers, *self.junk_registers])
 
     @classmethod
     def reflection_around_zero(
@@ -117,58 +122,79 @@ def reflection_around_zero(
             prepare_gate, control_val=control_val, global_phase=global_phase, eps=eps
         )
 
-    def decompose_from_registers(
-        self,
-        context: cirq.DecompositionContext,
-        **quregs: NDArray[cirq.Qid],  # type:ignore[type-var]
-    ) -> Iterator[cirq.OP_TREE]:
-        qm = context.qubit_manager
+    def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str, 'SoquetT']:
         # 0. Allocate new ancillas, if needed.
-        phase_target = qm.qalloc(1)[0] if self.control_val is None else quregs.pop('control')[0]
-        state_prep_ancilla = {
-            reg.name: np.array(qm.qalloc(reg.total_bits())).reshape(reg.shape + (reg.bitsize,))
-            for reg in self.prepare_gate.junk_registers
-        }
-        state_prep_selection_regs = quregs
-        prepare_op = self.prepare_gate.on_registers(
-            **state_prep_selection_regs, **state_prep_ancilla
+        phase_target = (
+            bb.allocate(dtype=QBit()) if self.control_val is None else soqs.pop('control')
         )
+        # prep_soqs =
+        # state_prep_junk = {
+        #     reg.name: soqs.pop(reg.name)
+        #     for reg in self.prepare_gate.junk_registers
+        #     if reg.side & Side.LEFT
+        # }
+        # state_prep_sel = soqs
+        # state_prep_ancilla = {
+        #     reg.name: np.array(bb.qalloc(reg.total_bits())).reshape(reg.shape + (reg.bitsize,))
+        #     for reg in self.prepare_gate.junk_registers
+        # }
+        # state_prep_selection_regs = quregs
+        # prepare_op = self.prepare_gate.on_registers(
+        #     **state_prep_selection_regs, **state_prep_ancilla
+        # )
         # 1. PREPARE†
-        yield cirq.inverse(prepare_op)
+        soqs = bb.add_d(self.prepare_gate.adjoint(), **soqs)
+        # yield cirq.inverse(prepare_op)
         # 2. MultiControlled Z, controlled on |000..00> state.
-        phase_control = np.array(
-            merge_qubits(self.selection_registers, **state_prep_selection_regs)
-        )
-        yield cirq.X(phase_target) if not self.control_val else []
-        yield MultiControlZ([0] * len(phase_control)).on_registers(
-            controls=phase_control.reshape(phase_control.shape + (1,)), target=phase_target
-        )
+        # phase_control = np.array(
+        #     merge_qubits(self.selection_registers, **state_prep_selection_regs)
+        # )
+        if self.control_val is None:
+            phase_target = bb.add(XGate(), q=phase_target)
+        multi_ctrl_soqs = {
+            ctrl_reg.name: soqs.pop(sel_reg.name)
+            for ctrl_reg, sel_reg in zip(self._multi_ctrl_z.ctrl_regs, self.selection_registers)
+        }
+        multi_ctrl_soqs = bb.add_d(self._multi_ctrl_z, **multi_ctrl_soqs, target=phase_target)
+        for ctrl_reg, sel_reg in zip(self._multi_ctrl_z.ctrl_regs, self.selection_registers):
+            soqs[sel_reg.name] = multi_ctrl_soqs.pop(ctrl_reg.name)
+        phase_target = multi_ctrl_soqs.pop('target')
+        assert not multi_ctrl_soqs
         if self.global_phase != 1:
+            exponent = sarg(self.global_phase) / pi(self.global_phase)
             if self.control_val is None:
-                yield cirq.global_phase_operation(self.global_phase, atol=self.eps)
+                bb.add(GlobalPhase(exponent=exponent, eps=self.eps))
             else:
-                yield cirq.Z(phase_target) ** (np.angle(self.global_phase) / np.pi)
-        yield cirq.X(phase_target) if not self.control_val else []
+                phase_target = bb.add(ZPowGate(exponent=exponent, eps=self.eps), q=phase_target)
         # 3. PREPARE
-        yield prepare_op
-
-        # 4. Deallocate ancilla.
-        qm.qfree([q for anc in state_prep_ancilla.values() for q in anc.flatten()])
+        soqs = bb.add_d(self.prepare_gate, **soqs)
+        # 4. Deallocate phase_target.
         if self.control_val is None:
-            qm.qfree([phase_target])
-
-    def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
-        wire_symbols = ['@' if self.control_val else '@(0)'] * total_bits(self.control_registers)
-        wire_symbols += ['R_L'] * total_bits(self.selection_registers)
-        return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
+            phase_target = bb.add(XGate(), q=phase_target)
+            bb.free(phase_target)
+        else:
+            soqs['control'] = phase_target
+        return soqs
+
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg.name == 'control':
+            return Circle(filled=True)
+        if reg.name in {reg.name for reg in self.selection_registers}:
+            return f'R_L[{",".join(str(i) for i in idx)}]'
+        return self.prepare_gate.wire_symbol(reg, idx)
+
+    # def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
+    #     wire_symbols = ['@' if self.control_val else '@(0)'] * total_bits(self.control_registers)
+    #     wire_symbols += ['R_L'] * total_bits(self.selection_registers)
+    #     return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
-        n_phase_control = sum(reg.total_bits() for reg in self.selection_registers)
-        cvs = HasLength(n_phase_control) if is_symbolic(n_phase_control) else [0] * n_phase_control
+        # n_phase_control = sum(reg.total_bits() for reg in self.selection_registers)
+        # cvs = HasLength(n_phase_control) if is_symbolic(n_phase_control) else [0] * n_phase_control
         costs: Set['BloqCountT'] = {
             (self.prepare_gate, 1),
             (self.prepare_gate.adjoint(), 1),
-            (MultiControlZ(cvs), 1),
+            (self._multi_ctrl_z, 1),
         }
         if self.control_val is None:
             costs.add((XGate(), 2))
@@ -179,6 +205,14 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
             costs.add((phase_op, 1))
         return costs
 
+    @cached_property
+    def _multi_ctrl_z(self) -> MultiControlZ:
+        ctrl_soqs, dtypes, cvs = {}, [], []
+        for reg in self.selection_registers:
+            dtypes.append(reg.dtype)
+            cvs.append(np.zeros(reg.shape, dtype=int))
+        return MultiControlZ(CtrlSpec(tuple(dtypes), tuple(cvs)))
+
 
 @bloq_example(generalizer=ignore_split_join)
 def _refl_using_prep() -> ReflectionUsingPrepare:
diff --git a/qualtran/bloqs/reflections/reflection_using_prepare_test.py b/qualtran/bloqs/reflections/reflection_using_prepare_test.py
index 30830e6d38..439ca0bf3a 100644
--- a/qualtran/bloqs/reflections/reflection_using_prepare_test.py
+++ b/qualtran/bloqs/reflections/reflection_using_prepare_test.py
@@ -13,17 +13,19 @@
 #  limitations under the License.
 
 import itertools
+from typing import *
 from typing import Optional
 
+import attrs
 import cirq
 import numpy as np
 import pytest
 from numpy.typing import NDArray
 
-from qualtran import Adjoint, Bloq
+from qualtran import Adjoint, Bloq, BloqBuilder, Side, Signature
 from qualtran._infra.gate_with_registers import get_named_qubits
 from qualtran.bloqs.arithmetic import LessThanConstant, LessThanEqual
-from qualtran.bloqs.basic_gates import ZPowGate
+from qualtran.bloqs.basic_gates import Hadamard, OnEach, ZeroState, ZPowGate
 from qualtran.bloqs.basic_gates.swap import CSwap
 from qualtran.bloqs.mcmt import MultiControlPauli, MultiTargetCNOT
 from qualtran.bloqs.mcmt.and_bloq import And
@@ -33,7 +35,10 @@
     _refl_using_prep,
     ReflectionUsingPrepare,
 )
-from qualtran.bloqs.state_preparation import StatePreparationAliasSampling
+from qualtran.bloqs.state_preparation import (
+    PrepareUniformSuperposition,
+    StatePreparationAliasSampling,
+)
 from qualtran.cirq_interop import BloqAsCirqGate
 from qualtran.cirq_interop.testing import GateHelper
 from qualtran.resource_counting.generalizers import (
@@ -112,32 +117,116 @@ def get_3q_uniform_dirac_notation(signs, global_phase: complex = 1):
     return ret
 
 
-@pytest.mark.slow
-@pytest.mark.parametrize('num_ones', [5])
-@pytest.mark.parametrize('eps', [0.05])
-@pytest.mark.parametrize('global_phase', [+1, -1j])
-def test_reflection_using_prepare(num_ones, eps, global_phase):
-    data = [1] * num_ones
-    prepare_gate = StatePreparationAliasSampling.from_probabilities(data, precision=eps)
-
-    gate = ReflectionUsingPrepare(prepare_gate, global_phase=global_phase)
-    assert_valid_bloq_decomposition(gate)
-
-    g, qubit_order, decomposed_circuit = construct_gate_helper_and_qubit_order(gate)
-
-    initial_state_prep = cirq.Circuit(cirq.H.on_each(*g.quregs['selection']))
-    initial_state = cirq.dirac_notation(initial_state_prep.final_state_vector())
-    assert initial_state == get_3q_uniform_dirac_notation('++++++++')
-    result = cirq.Simulator(dtype=np.complex64).simulate(
-        initial_state_prep + decomposed_circuit, qubit_order=qubit_order
-    )
-    selection = g.quregs['selection']
-    prepared_state = result.final_state_vector.reshape(2 ** len(selection), -1).sum(axis=1)
-    if np.sign(global_phase) == 1:
-        signs = '-' * num_ones + '+' * (9 - num_ones)
-    elif np.sign(global_phase) == -1:
-        signs = '+' * num_ones + '-' * (9 - num_ones)
-    assert cirq.dirac_notation(prepared_state) == get_3q_uniform_dirac_notation(signs, global_phase)
+# def _add_classical_kets(bb: BloqBuilder, registers: Iterable['Register']) -> Dict[str, 'SoquetT']:
+#     """Use `bb` to add `IntState(0)` for all the `vals`."""
+#
+#     from qualtran.bloqs.basic_gates import IntState, PlusState, OnEach, ZeroState
+#
+#     soqs: Dict[str, 'SoquetT'] = {}
+#     for reg in registers:
+#         # add_state = PlusState if reg.name == 'selection' else ZeroState
+#         add_state = ZeroState
+#         if reg.shape:
+#             soq = np.empty(reg.shape, dtype=object)
+#             for idx in reg.all_idxs():
+#                 soq[idx] = bb.join(
+#                     np.array([bb.add(add_state()) for _ in range(reg.bitsize)]), dtype=reg.dtype
+#                 )
+#         else:
+#             soq = bb.join(
+#                 np.array([bb.add(add_state()) for _ in range(reg.bitsize)]), dtype=reg.dtype
+#             )
+#
+#         soqs[reg.name] = soq
+#     return soqs
+#
+#
+# def initialize_from_zero(bloq: Bloq):
+#     """Take `bloq` and compose it with initial zero states for each left register.
+#
+#     This can be contracted to a state vector for a given unitary.
+#     """
+#     bb = BloqBuilder()
+#
+#     # Add the initial 'kets' according to the provided values.
+#     in_soqs = _add_classical_kets(bb, bloq.signature.lefts())
+#
+#     in_soqs['selection'] = bb.add(PrepareUniformSuperposition(5), target=in_soqs['selection'])
+#
+#     # Add the bloq itself
+#     out_soqs = bb.add_d(bloq, **in_soqs)
+#     return bb.finalize(**out_soqs)
+#
+#
+# @pytest.mark.slow
+# @pytest.mark.parametrize('num_ones', [5])
+# @pytest.mark.parametrize('eps', [0.05])
+# @pytest.mark.parametrize('global_phase', [+1, -1j])
+# def test_reflection_using_prepare(num_ones, eps, global_phase):
+#     data = [1] * num_ones
+#     prepare_gate = StatePreparationAliasSampling.from_probabilities(data, precision=eps)
+#     print(prepare_gate.signature)
+#     prepared_state = initialize_from_zero(prepare_gate.adjoint()).tensor_contract()
+#     print(cirq.dirac_notation(prepared_state))
+#     L, logL = len(data), 3
+#     state_vector = prepared_state.reshape(2**logL, len(prepared_state) // 2**logL)
+#     prepared_state = np.linalg.norm(state_vector, axis=1)
+#     print(cirq.dirac_notation(prepared_state))
+#
+#     gate = ReflectionUsingPrepare(prepare_gate, global_phase=global_phase)
+#     sel_reg = gate.selection_registers[0]
+#     bb, soqs = BloqBuilder.from_signature(
+#         Signature([attrs.evolve(reg, side=Side.RIGHT) for reg in gate.signature])
+#     )
+#     for reg in gate.signature:
+#         soqs[reg.name] = bb.join(
+#             np.array([bb.add(ZeroState()) for _ in range(reg.bitsize)]), dtype=reg.dtype
+#         )
+#     soqs['selection'] = bb.add(
+#         OnEach(
+#             prepare_gate.selection_bitsize,
+#             Hadamard(),
+#             target_dtype=prepare_gate.selection_registers[0].dtype,
+#         ),
+#         q=soqs['selection'],
+#     )
+#     # for reg in gate.signature:
+#     #     if reg.name != 'selection':
+#     #         soqs[reg.name] = bb.allocate(dtype=reg.dtype)
+#     soqs = bb.add_d(gate, **soqs)
+#     # for reg in gate.signature:
+#     #     if reg.name != 'selection':
+#     #         bb.free(soqs.pop(reg.name))
+#     cbloq = bb.finalize(**soqs)
+#     print(cbloq.flatten_once().to_cirq_circuit())
+#     assert_valid_bloq_decomposition(gate)
+#     prepared_state = cbloq.tensor_contract()
+#     print(prepared_state.shape)
+#     print(cirq.dirac_notation(prepared_state))
+#     L, logL = len(data), 3
+#     state_vector = prepared_state.reshape(2**logL, len(prepared_state) // 2**logL)
+#     prepared_state = np.linalg.norm(state_vector, axis=1)
+#     print(cirq.dirac_notation(prepared_state))
+#     # print(gate.signature)
+#     # print(gate._multi_ctrl_z.signature)
+#     #
+#     # g, qubit_order, decomposed_circuit = construct_gate_helper_and_qubit_order(gate)
+#     # print(decomposed_circuit)
+#     #
+#     # initial_state_prep = cirq.Circuit(cirq.H.on_each(*g.quregs['selection']))
+#     # initial_state = cirq.dirac_notation(initial_state_prep.final_state_vector())
+#     # assert initial_state == get_3q_uniform_dirac_notation('++++++++')
+#     # result = cirq.Simulator(dtype=np.complex128).simulate(
+#     #     initial_state_prep + decomposed_circuit, qubit_order=qubit_order
+#     # )
+#     # selection = g.quregs['selection']
+#     # prepared_state = result.final_state_vector.reshape(2 ** len(selection), -1).sum(axis=1)
+#     if np.sign(global_phase) == 1:
+#         signs = '-' * num_ones + '+' * (9 - num_ones)
+#     elif np.sign(global_phase) == -1:
+#         signs = '+' * num_ones + '-' * (9 - num_ones)
+#     assert cirq.dirac_notation(prepared_state) == get_3q_uniform_dirac_notation(signs, global_phase)
+#
 
 
 def test_reflection_using_prepare_diagram():
@@ -151,25 +240,27 @@ def test_reflection_using_prepare_diagram():
     cirq.testing.assert_has_diagram(
         circuit,
         '''
-ancilla_0: ────X───────────────────────────────Z──────X───────────────────────────────
-                                               │
-ancilla_1: ────sigma_mu────────────────────────┼──────sigma_mu────────────────────────
-               │                               │      │
-ancilla_2: ────alt─────────────────────────────┼──────alt─────────────────────────────
-               │                               │      │
-ancilla_3: ────alt─────────────────────────────┼──────alt─────────────────────────────
-               │                               │      │
-ancilla_4: ────alt─────────────────────────────┼──────alt─────────────────────────────
-               │                               │      │
-ancilla_5: ────keep────────────────────────────┼──────keep────────────────────────────
-               │                               │      │
-ancilla_6: ────less_than_equal─────────────────┼──────less_than_equal─────────────────
-               │                               │      │
-selection0: ───StatePreparationAliasSampling───@(0)───StatePreparationAliasSampling───
-               │                               │      │
-selection1: ───selection───────────────────────@(0)───selection───────────────────────
-               │                               │      │
-selection2: ───selection^-1────────────────────@(0)───selection───────────────────────
+                    ┌──────────────────────────────┐            ┌──────────────────────────────┐
+alt0: ───────────────alt─────────────────────────────────────────alt───────────────────────────────
+                     │                                           │
+alt1: ───────────────alt─────────────────────────────────────────alt───────────────────────────────
+                     │                                           │
+alt2: ───────────────alt─────────────────────────────────────────alt───────────────────────────────
+                     │                                           │
+ancilla_0: ──────────┼────────────────────────────X────target────┼────────────────────────────X────
+                     │                                 │         │
+keep: ───────────────keep──────────────────────────────┼─────────keep──────────────────────────────
+                     │                                 │         │
+less_than_equal: ────less_than_equal───────────────────┼─────────less_than_equal───────────────────
+                     │                                 │         │
+selection0: ─────────StatePreparationAliasSampling─────ctrl0_────StatePreparationAliasSampling─────
+                     │                                 │         │
+selection1: ─────────selection─────────────────────────ctrl0_────selection─────────────────────────
+                     │                                 │         │
+selection2: ─────────selection─────────────────────────ctrl0_────selection─────────────────────────
+                     │                                           │
+sigma_mu: ───────────sigma_mu^-1─────────────────────────────────sigma_mu──────────────────────────
+                    └──────────────────────────────┘            └──────────────────────────────┘
 ''',
     )
 
@@ -179,25 +270,25 @@ def test_reflection_using_prepare_diagram():
     cirq.testing.assert_has_diagram(
         circuit,
         '''
-ancilla_0: ────sigma_mu───────────────────────────────sigma_mu────────────────────────
-               │                                      │
-ancilla_1: ────alt────────────────────────────────────alt─────────────────────────────
-               │                                      │
-ancilla_2: ────alt────────────────────────────────────alt─────────────────────────────
-               │                                      │
-ancilla_3: ────alt────────────────────────────────────alt─────────────────────────────
-               │                                      │
-ancilla_4: ────keep───────────────────────────────────keep────────────────────────────
-               │                                      │
-ancilla_5: ────less_than_equal────────────────────────less_than_equal─────────────────
-               │                                      │
-control: ──────┼───────────────────────────────Z──────┼───────────────────────────────
-               │                               │      │
-selection0: ───StatePreparationAliasSampling───@(0)───StatePreparationAliasSampling───
-               │                               │      │
-selection1: ───selection───────────────────────@(0)───selection───────────────────────
-               │                               │      │
-selection2: ───selection^-1────────────────────@(0)───selection───────────────────────
+alt0: ──────────────alt──────────────────────────────────────alt─────────────────────────────
+                    │                                        │
+alt1: ──────────────alt──────────────────────────────────────alt─────────────────────────────
+                    │                                        │
+alt2: ──────────────alt──────────────────────────────────────alt─────────────────────────────
+                    │                                        │
+control: ───────────┼───────────────────────────────target───┼───────────────────────────────
+                    │                               │        │
+keep: ──────────────keep────────────────────────────┼────────keep────────────────────────────
+                    │                               │        │
+less_than_equal: ───less_than_equal─────────────────┼────────less_than_equal─────────────────
+                    │                               │        │
+selection0: ────────StatePreparationAliasSampling───ctrl0_───StatePreparationAliasSampling───
+                    │                               │        │
+selection1: ────────selection───────────────────────ctrl0_───selection───────────────────────
+                    │                               │        │
+selection2: ────────selection───────────────────────ctrl0_───selection───────────────────────
+                    │                                        │
+sigma_mu: ──────────sigma_mu^-1──────────────────────────────sigma_mu────────────────────────
 ''',
     )
 
@@ -207,27 +298,25 @@ def test_reflection_using_prepare_diagram():
     cirq.testing.assert_has_diagram(
         circuit,
         '''
-               ┌──────────────────────────────┐          ┌──────────────────────────────┐
-ancilla_0: ─────sigma_mu───────────────────────────────────sigma_mu─────────────────────────
-                │                                          │
-ancilla_1: ─────alt────────────────────────────────────────alt──────────────────────────────
-                │                                          │
-ancilla_2: ─────alt────────────────────────────────────────alt──────────────────────────────
-                │                                          │
-ancilla_3: ─────alt────────────────────────────────────────alt──────────────────────────────
-                │                                          │
-ancilla_4: ─────keep───────────────────────────────────────keep─────────────────────────────
-                │                                          │
-ancilla_5: ─────less_than_equal────────────────────────────less_than_equal──────────────────
-                │                                          │
-control: ───────┼────────────────────────────X────Z───────X┼────────────────────────────────
-                │                                 │        │
-selection0: ────StatePreparationAliasSampling─────@(0)─────StatePreparationAliasSampling────
-                │                                 │        │
-selection1: ────selection─────────────────────────@(0)─────selection────────────────────────
-                │                                 │        │
-selection2: ────selection^-1──────────────────────@(0)─────selection────────────────────────
-               └──────────────────────────────┘          └──────────────────────────────┘
+alt0: ──────────────alt──────────────────────────────────────alt─────────────────────────────
+                    │                                        │
+alt1: ──────────────alt──────────────────────────────────────alt─────────────────────────────
+                    │                                        │
+alt2: ──────────────alt──────────────────────────────────────alt─────────────────────────────
+                    │                                        │
+control: ───────────┼───────────────────────────────target───┼───────────────────────────────
+                    │                               │        │
+keep: ──────────────keep────────────────────────────┼────────keep────────────────────────────
+                    │                               │        │
+less_than_equal: ───less_than_equal─────────────────┼────────less_than_equal─────────────────
+                    │                               │        │
+selection0: ────────StatePreparationAliasSampling───ctrl0_───StatePreparationAliasSampling───
+                    │                               │        │
+selection1: ────────selection───────────────────────ctrl0_───selection───────────────────────
+                    │                               │        │
+selection2: ────────selection───────────────────────ctrl0_───selection───────────────────────
+                    │                                        │
+sigma_mu: ──────────sigma_mu^-1──────────────────────────────sigma_mu────────────────────────
 ''',
     )