Change return type of build_call_graph to set | dict (#1356)

* Change return type of build_call_graph to set | dict

- Changes build_call_graph to return a Union of
Dict[Bloq, Union[int, Expr]] and Set[BloqCountT]
- Adds checks in main call sites to check both cases.

Author: dstrain115

qualtran/_infra/adjoint.py

@@ -12,8 +12,9 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from collections import Counter
 from functools import cached_property
-from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING
+from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
 
 import cirq
 from attrs import frozen
@@ -26,7 +27,7 @@
 if TYPE_CHECKING:
     from qualtran import Bloq, CompositeBloq, Register, Signature, SoquetT
     from qualtran.drawing import WireSymbol
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 
 
 def _adjoint_final_soqs(cbloq: 'CompositeBloq', new_signature: Signature) -> Dict[str, 'SoquetT']:
@@ -158,9 +159,17 @@ def adjoint(self) -> 'Bloq':
         """The 'double adjoint' brings you back to the original bloq."""
         return self.subbloq
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         """The call graph takes the adjoint of each of the bloqs in `subbloq`'s call graph."""
-        return {(bloq.adjoint(), n) for bloq, n in self.subbloq.build_call_graph(ssa=ssa)}
+        sub_cg = self.subbloq.build_call_graph(ssa=ssa)
+        counts = Counter['Bloq']()
+        if isinstance(sub_cg, set):
+            for bloq, n in sub_cg:
+                counts[bloq.adjoint()] += n
+        else:
+            for bloq, n in sub_cg.items():
+                counts[bloq.adjoint()] += n
+        return counts
 
     def pretty_name(self) -> str:
         """The subbloq's pretty_name with a dagger."""

qualtran/_infra/bloq.py

@@ -39,7 +39,13 @@
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
     from qualtran.drawing import WireSymbol
-    from qualtran.resource_counting import BloqCountT, CostKey, GeneralizerT, SympySymbolAllocator
+    from qualtran.resource_counting import (
+        BloqCountDictT,
+        BloqCountT,
+        CostKey,
+        GeneralizerT,
+        SympySymbolAllocator,
+    )
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -279,7 +285,9 @@ def my_tensors(
         """
         raise NotImplementedError(f"{self} does not support tensor simulation.")
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         """Override this method to build the bloq call graph.
 
         This method must return a set of `(bloq, n)` tuples where `bloq` is called `n` times in

qualtran/_infra/controlled.py

@@ -11,6 +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.
+from collections import Counter
 from functools import cached_property
 from typing import (
     Any,
@@ -20,7 +21,6 @@
     Optional,
     Protocol,
     Sequence,
-    Set,
     Tuple,
     TYPE_CHECKING,
     Union,
@@ -39,10 +39,10 @@
 if TYPE_CHECKING:
     import quimb.tensor as qtn
 
-    from qualtran import BloqBuilder, CompositeBloq, ConnectionT, SoquetT
+    from qualtran import Bloq, BloqBuilder, CompositeBloq, ConnectionT, SoquetT
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.drawing import WireSymbol
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -386,7 +386,7 @@ def build_composite_bloq(
         fsoqs |= dict(zip(self.ctrl_reg_names, ctrl_soqs))
         return fsoqs
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         try:
             sub_cg = self.subbloq.build_call_graph(ssa=ssa)
         except DecomposeTypeError as e1:
@@ -396,7 +396,14 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
                 f"Could not build call graph for {self}: {e2}"
             ) from e2
 
-        return {(bloq.controlled(self.ctrl_spec), n) for bloq, n in sub_cg}
+        counts = Counter['Bloq']()
+        if isinstance(sub_cg, set):
+            for bloq, n in sub_cg:
+                counts[bloq.controlled(self.ctrl_spec)] += n
+        else:
+            for bloq, n in sub_cg.items():
+                counts[bloq.controlled(self.ctrl_spec)] += n
+        return counts
 
     def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
         ctrl_vals = [vals[reg_name] for reg_name in self.ctrl_reg_names]

qualtran/bloqs/arithmetic/addition.py

@@ -62,7 +62,7 @@
 
 if TYPE_CHECKING:
     from qualtran.drawing import WireSymbol
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
     from qualtran.symbolics import SymbolicInt
 
@@ -500,7 +500,9 @@ def build_composite_bloq(
         else:
             return {'x': x}
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         loading_cost: Tuple[Bloq, SymbolicInt]
         if len(self.cvs) == 0:
             loading_cost = (XGate(), self.bitsize)  # upper bound; depends on the data.

qualtran/bloqs/arithmetic/permutation.py

@@ -52,7 +52,7 @@
     import sympy
 
     from qualtran import BloqBuilder, SoquetT
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
 
 SymbolicCycleT: TypeAlias = Union[CycleT, Shaped]
 
@@ -122,7 +122,9 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'SoquetT') -> dict[str, 'So
 
         return {'x': x}
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         if is_symbolic(self.cycle):
             x = ssa.new_symbol('x')
             cycle_len = slen(self.cycle)
@@ -267,7 +269,9 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'Soquet') -> dict[str, 'Soq
 
         return {'x': x}
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         if is_symbolic(self.cycles):
             # worst case cost: single cycle of length N
             cycle = Shaped((self.N,))

qualtran/bloqs/data_loading/qrom.py

@@ -33,7 +33,7 @@
 from qualtran.symbolics import prod, SymbolicInt
 
 if TYPE_CHECKING:
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
 
 
 def _to_tuple(x: Iterable[NDArray]) -> Sequence[NDArray]:
@@ -208,7 +208,9 @@ def nth_operation_callgraph(self, **kwargs: int) -> Set['BloqCountT']:
                 ret += data_to_load.bit_count()
         return {(CNOT(), ret)}
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         if self.has_data():
             return super().build_call_graph(ssa=ssa)
         n_and = prod(self.data_shape) - 2 + self.num_controls

qualtran/bloqs/multiplexers/unary_iteration_bloq.py

@@ -30,7 +30,7 @@
     import sympy
 
     from qualtran import Bloq
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
     from qualtran.symbolics import SymbolicInt
 
 
@@ -591,7 +591,9 @@ def nth_operation_callgraph(self, **selection_regs_name_to_val) -> Set['BloqCoun
             f"Derived class {type(self)} does not implement `nth_operation_callgraph`."
         )
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         if total_bits(self.selection_registers) == 0 or self._break_early(
             (), 0, self.selection_registers[0].dtype.iteration_length_or_zero()
         ):

qualtran/bloqs/state_preparation/prepare_uniform_superposition.py

@@ -37,7 +37,7 @@
 )
 
 if TYPE_CHECKING:
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
 
 
 @attrs.frozen
@@ -148,7 +148,9 @@ def decompose_from_registers(
         yield cirq.H.on_each(*logL_qubits)
         context.qubit_manager.qfree([*and_target, *and_ancilla])
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
         if not is_symbolic(self.n, self.cvs):
             # build from decomposition
             return super().build_call_graph(ssa)

qualtran/cirq_interop/t_complexity_protocol.py

@@ -19,7 +19,7 @@
 import cirq
 
 from qualtran import Bloq, DecomposeNotImplementedError, DecomposeTypeError
-from qualtran.resource_counting import SympySymbolAllocator
+from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
 from qualtran.symbolics import ceil, log2, SymbolicFloat, SymbolicInt
 
 from .decompose_protocol import _decompose_once_considering_known_decomposition
@@ -166,7 +166,11 @@ def _from_bloq_build_call_graph(bloq: Bloq) -> Optional[TComplexity]:
         return None
 
     ret = TComplexity()
-    for callee, n in callee_counts:
+    if isinstance(callee_counts, set):
+        callee_iterator: Iterable[BloqCountT] = callee_counts
+    else:
+        callee_iterator = callee_counts.items()
+    for callee, n in callee_iterator:
         r = t_complexity(callee)
         if r is None:
             return None

qualtran/resource_counting/__init__.py

@@ -20,6 +20,7 @@
 from ._generalization import GeneralizerT
 
 from ._call_graph import (
+    BloqCountDictT,
     BloqCountT,
     big_O,
     SympySymbolAllocator,