Fix qubit counting counting classical bits (#1816)

Fixes #1811 

 - adds new helper methods that are less misleading
 - uses them in the qubit counting code

@NoureldinYosri

Author: mpharrigan

qualtran/_infra/quantum_graph.py

@@ -139,14 +139,47 @@ class Connection:
     left: Soquet
     right: Soquet
 
+    @cached_property
+    def num_qubits(self) -> int:
+        """The number of qubits in the connection.
+
+        This excludes classical bits.
+        """
+        lq = self.left.reg.dtype.num_qubits
+        rq = self.right.reg.dtype.num_qubits
+
+        if lq != rq:
+            raise ValueError(f"Invalid Connection {self}: num_qubits mismatch: {lq} != {rq}")
+        return lq
+
+    @cached_property
+    def num_cbits(self) -> int:
+        """The number of classical bits in the connection."""
+        lc = self.left.reg.dtype.num_cbits
+        rc = self.right.reg.dtype.num_cbits
+
+        if lc != rc:
+            raise ValueError(f"Invalid Connection {self}: num_cbits mismatch: {lc} != {rc}")
+        return lc
+
+    @cached_property
+    def num_bits(self) -> int:
+        """The number of bits in the connection (quantum + classical)."""
+        lb = self.left.reg.dtype.num_bits
+        rb = self.right.reg.dtype.num_bits
+
+        if lb != rb:
+            raise ValueError(f"Invalid Connection {self}: shape mismatch: {lb} != {rb}")
+        return lb
+
     @cached_property
     def shape(self) -> int:
-        ls = self.left.reg.bitsize
-        rs = self.right.reg.bitsize
+        """The number of bits in the connection (quantum + classical).
 
-        if ls != rs:
-            raise ValueError(f"Invalid Connection {self}: shape mismatch: {ls} != {rs}")
-        return ls
+        This is a misleading name for this property kept for backwards compatibility.
+        Please prefer `.num_bits`.
+        """
+        return self.num_bits
 
     def __str__(self) -> str:
         return f'{self.left} -> {self.right}'

qualtran/drawing/graphviz.py

@@ -311,7 +311,7 @@ def add_binst(self, graph: pydot.Graph, binst: BloqInstance) -> pydot.Graph:
 
     def cxn_label(self, cxn: Connection) -> str:
         """Overridable method to return labels for connections."""
-        return str(cxn.shape)
+        return str(cxn.num_bits)
 
     def cxn_edge(self, left_id: str, right_id: str, cxn: Connection) -> pydot.Edge:
         """Overridable method to style a pydot.Edge for connecionts."""

qualtran/resource_counting/_qubit_counts.py

@@ -65,12 +65,12 @@ def _cbloq_max_width(
                 # During the application of the binst, we have "observer" connections that have
                 # width as well as the width from the binst itself. We consider the case where
                 # the bloq may have a max_width greater than the max of its left/right registers.
-                during_size = _bloq_max_width(binst.bloq) + sum(s.shape for s in in_play)
+                during_size = _bloq_max_width(binst.bloq) + sum(s.num_qubits for s in in_play)
                 max_width = smax(max_width, during_size)
 
             # After the binst, its successor connections are 'in play'.
             in_play.update(succ_cxns)
-            after_size = sum(s.shape for s in in_play)
+            after_size = sum(s.num_qubits for s in in_play)
             max_width = smax(max_width, after_size)
 
     return max_width

qualtran/resource_counting/_qubit_counts_test.py

@@ -95,6 +95,35 @@ def test_many_alloc():
         assert get_cost_value(bloq, QubitCount()) == 11
 
 
+def test_qubit_count_ignores_cbits():
+    import functools
+
+    import attr
+
+    from qualtran import Bloq, BloqBuilder, CBit, QBit, Register, Signature
+
+    @attr.frozen
+    class MyBloq(Bloq):
+        n: int
+        m: int
+
+        @functools.cached_property
+        def signature(self):
+            return Signature(
+                [Register('qs', QBit(), shape=(self.n,)), Register('cs', CBit(), shape=(self.m,))]
+            )
+
+        def build_call_graph(self, ssa):
+            return {}
+
+        def build_composite_bloq(self, bb: BloqBuilder, qs, cs):
+            return {'qs': qs, 'cs': cs}
+
+    blq = MyBloq(5, 100)
+    assert get_cost_value(blq, QubitCount()) == 5
+    assert get_cost_value(blq.decompose_bloq(), QubitCount()) == 5
+
+
 @pytest.mark.notebook
 def test_notebook():
     qlt_testing.execute_notebook("qubit_counts")