|
33 | 33 |
|
34 | 34 | from qualtran import Bloq, BloqInstance, CompositeBloq |
35 | 35 | from qualtran import Connection as QualtranConnection |
36 | | -from qualtran import Register, Side, Soquet |
| 36 | +from qualtran import DecomposeNotImplementedError, DecomposeTypeError, Register, Side, Soquet |
37 | 37 | from qualtran.cirq_interop import CirqGateAsBloq |
38 | 38 | from qualtran.symbolics import is_symbolic |
39 | 39 |
|
@@ -128,23 +128,92 @@ def _bloq_instance_name(instance: BloqInstance) -> str: |
128 | 128 | return f"{_bloq_type(instance.bloq)}_{instance.i}" |
129 | 129 |
|
130 | 130 |
|
| 131 | +@singledispatch |
131 | 132 | def bloq_to_qref( |
132 | | - obj: Union[Bloq, CompositeBloq, BloqInstance], from_callgraph: bool = False |
| 133 | + obj: Union[Bloq, CompositeBloq, BloqInstance], |
| 134 | + *, |
| 135 | + from_callgraph: bool = False, |
| 136 | + decomposition_rules: Union[bool, Iterable[type[Bloq]]] = False, |
133 | 137 | ) -> SchemaV1: |
134 | | - """Converts Bloq to QREF SchemaV1 object. |
| 138 | + """Converts a Qualtran Bloq into a QREF SchemaV1. |
135 | 139 |
|
136 | 140 | Args: |
137 | | - obj: bloq to be converted |
138 | | - from_callgraph: a flag indicating whether conversion should be performed using only the information |
139 | | - from callgraph. It's useful when a bloq doesn't have a full decomposition, but has a callgraph. |
| 141 | + obj: Bloq, CompositeBloq, or BloqInstance to be converted. |
| 142 | + from_callgraph: if True, import purely from the bloq’s call-graph. |
| 143 | + decomposition_rules: |
| 144 | + • False (default): no extra unfolding. |
| 145 | + • True: attempt to decompose *every* Bloq via decompose_bloq(). |
| 146 | + • Iterable of Bloq classes: decompose only those types. |
140 | 147 |
|
| 148 | + Returns: |
| 149 | + A SchemaV1 whose `.program` is the top-level RoutineV1. |
141 | 150 | """ |
142 | 151 | if from_callgraph: |
143 | 152 | if isinstance(obj, BloqInstance): |
144 | | - raise ValueError("BloqInstance object can't be used to generate a callgraph.") |
| 153 | + raise ValueError("BloqInstance cannot drive a callgraph import.") |
145 | 154 | return SchemaV1(version="v1", program=bloq_to_routine_from_callgraph(obj)) |
146 | 155 | else: |
147 | | - return SchemaV1(version="v1", program=bloq_to_routine(obj)) |
| 156 | + if decomposition_rules is False: |
| 157 | + program = bloq_to_routine(obj) |
| 158 | + else: |
| 159 | + keep = None if decomposition_rules is True else set(decomposition_rules) |
| 160 | + program = _routine_with_decomposition(obj, keep) |
| 161 | + return SchemaV1(version="v1", program=program) |
| 162 | + |
| 163 | + |
| 164 | +def _routine_with_decomposition( |
| 165 | + obj: Union[Bloq, CompositeBloq, BloqInstance], |
| 166 | + preserve: Optional[set[type[Bloq]]], |
| 167 | + *, |
| 168 | + name: Optional[str] = None, |
| 169 | +) -> RoutineV1: |
| 170 | + """Convert a Qualtran bloq into a QREF RoutineV1, selectively decomposing. |
| 171 | +
|
| 172 | + This helper will import `obj` (which may be a Bloq, CompositeBloq or BloqInstance) |
| 173 | + into a nested RoutineV1, but only expand (“decompose”) those Bloq types |
| 174 | + requested in `preserve`. |
| 175 | +
|
| 176 | + Behavior: |
| 177 | + - If `preserve` is None, every bloq that implements `decompose_bloq()` |
| 178 | + will be peeled apart into its CompositeBloq and imported recursively. |
| 179 | + - If `preserve` is a set of Bloq classes, then only instances of those |
| 180 | + classes will be decomposed; all other Bloqs are treated as atomic leaves |
| 181 | + (using the original, default `bloq_to_routine` handler). |
| 182 | +
|
| 183 | + Args: |
| 184 | + obj: The root Bloq, CompositeBloq, or BloqInstance to convert. |
| 185 | + preserve: |
| 186 | + • None ⇒ decompose **all** bloqs where possible. |
| 187 | + • set of Bloq types ⇒ only those classes will be inlined; others remain atomic. |
| 188 | +
|
| 189 | + Returns: |
| 190 | + A RoutineV1 representing the imported bloq hierarchy, with clusters |
| 191 | + for exactly the bloq types you asked to decompose. |
| 192 | + """ |
| 193 | + # CompositeBloq → recurse |
| 194 | + if isinstance(obj, CompositeBloq): |
| 195 | + children = [_routine_with_decomposition(i, preserve) for i in obj.bloq_instances] |
| 196 | + connections = [_import_connection(c) for c in obj.connections] |
| 197 | + return RoutineV1( |
| 198 | + **_extract_common_bloq_attributes(obj, name), children=children, connections=connections |
| 199 | + ) |
| 200 | + |
| 201 | + # BloqInstance → unwrap |
| 202 | + if isinstance(obj, BloqInstance): |
| 203 | + return _routine_with_decomposition(obj.bloq, preserve, name=_bloq_instance_name(obj)) |
| 204 | + |
| 205 | + # Plain Bloq → decide |
| 206 | + bloq = obj |
| 207 | + if preserve is None or type(bloq) in preserve: |
| 208 | + try: |
| 209 | + cb = bloq.decompose_bloq() |
| 210 | + except (DecomposeTypeError, DecomposeNotImplementedError): |
| 211 | + return _default_leaf(bloq, name=name) |
| 212 | + else: |
| 213 | + return _routine_with_decomposition(cb, preserve, name=name or type(bloq).__name__) |
| 214 | + |
| 215 | + # leave leaf |
| 216 | + return _default_leaf(bloq, name=name) |
148 | 217 |
|
149 | 218 |
|
150 | 219 | @singledispatch |
@@ -411,3 +480,6 @@ def bloq_to_routine_from_callgraph(bloq: Union[Bloq, CompositeBloq]) -> RoutineV |
411 | 480 | nodes_to_routine_map[edge[0]] = parent |
412 | 481 |
|
413 | 482 | return nodes_to_routine_map[list(call_graph.nodes)[0]] |
| 483 | + |
| 484 | + |
| 485 | +_default_leaf = bloq_to_routine.registry[Bloq] |
0 commit comments