Use greedy_topological_sort to sort operations in musical scores

qualtran/drawing/musical_score.py

@@ -45,6 +45,7 @@
     Side,
     Signature,
 )
+from qualtran._infra.binst_graph_iterators import greedy_topological_sort
 from qualtran._infra.composite_bloq import _binst_to_cxns
 from qualtran._infra.quantum_graph import _Soquet
 
@@ -351,15 +352,39 @@ def _cbloq_musical_score(
 
     # Bloq-by-bloq application
     seq_x = 0
-    for topo_gen, binsts in enumerate(nx.topological_generations(binst_graph)):
-        for binst in binsts:
-            if isinstance(binst, DanglingT):
-                continue
-            pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph)
-            _binst_assign_line(
-                binst, pred_cxns, soq_assign, seq_x=seq_x, topo_gen=topo_gen, manager=manager
-            )
-            seq_x += 1
+    # Maintain a mapping of the current depth of the rightmost ends of all qubits.
+    y_to_topo_gen = {}
+    for binst in greedy_topological_sort(binst_graph):
+        if isinstance(binst, DanglingT):
+            continue
+        pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph)
+
+        # Compute the topological generation from predecesor nodes.
+        topo_gen = max((soq_assign[pred.left].topo_gen for pred in pred_cxns), default=0) + 1
+
+        # Offload line assignment of bloq instance to LineManager.
+        _binst_assign_line(
+            binst, pred_cxns, soq_assign, seq_x=seq_x, topo_gen=topo_gen, manager=manager
+        )
+
+        # Using the line assignment from earlier, compute the topological generation as being at
+        # least after the furthest node on the line it is placed on. This is generally only
+        # required for allocations or right-side bloqs that may reuse a previous qubit.
+        topo_gen = max(
+            [y_to_topo_gen.get(soq_assign[succ.left].y, 0) + 1 for succ in succ_cxns]
+            + [topo_gen]
+        )
+        for succ in succ_cxns:
+            soq_assign[succ.left] = attrs.evolve(soq_assign[succ.left], topo_gen=topo_gen)
+
+        # Update latest topological generation.
+        reg_idxs = set(
+            [soq_assign[pred.right].y for pred in pred_cxns]
+            + [soq_assign[succ.left].y for succ in succ_cxns]
+        )
+        for idx in reg_idxs:
+            y_to_topo_gen[idx] = topo_gen
+        seq_x += 1
 
     # Track bloq-to-dangle name changes
     if len(list(signature.rights())) > 0:

qualtran/drawing/musical_score_test.py

@@ -32,6 +32,31 @@ def test_dump_json(tmp_path):
     dump_musical_score(msd, name=f'{tmp_path}/musical_score_example')
 
 
+def test_musical_score_aligns_with_qubit_count():
+    from qualtran.bloqs.for_testing.qubit_count_many_alloc import (
+        TestManyAllocAbstracted,
+        TestManyAllocOnce,
+        TestManyAllocMany,
+    )
+    from qualtran.resource_counting import get_cost_cache, get_cost_value, QubitCount
+
+    n = 10
+    for bloq in [TestManyAllocMany(n), TestManyAllocOnce(n), TestManyAllocAbstracted(n)]:
+        expected_qubits = get_cost_value(bloq, QubitCount())
+        msd = get_musical_score_data(bloq.decompose_bloq())
+        # Ensure qubits (horizontal rows) match qubits from cost values in expected count by
+        # counting all rows with at least one non-dangle Soquet.
+        actual_qubits = len(
+            set(soq.rpos.y for soq in msd.soqs if not soq.ident.startswith('dang'))
+        )
+        assert (
+            actual_qubits == expected_qubits
+        ), (
+            f'{type(bloq).__name__} has too many non-dangle lines; '
+            + f'expected {expected_qubits}; got {actual_qubits}'
+        )
+
+
 @pytest.mark.notebook
 def test_notebook():
     execute_notebook('musical_score')