Bugfix in comparison_key and nicer with_dimension of _BaseAncillaQid (#6554)

* Bugfix in comparison_key and nicer with_dimension of _BaseAncillaQid

* Fix pylint

Author: tanujkhattar

cirq-core/cirq/ops/qubit_manager.py

@@ -14,7 +14,7 @@
 
 import abc
 import dataclasses
-from typing import Iterable, List, TYPE_CHECKING
+from typing import Iterable, List, TYPE_CHECKING, Tuple
 from cirq.ops import raw_types
 
 if TYPE_CHECKING:
@@ -41,13 +41,16 @@ class _BaseAncillaQid(raw_types.Qid):
     dim: int = 2
     prefix: str = ''
 
-    def _comparison_key(self) -> int:
-        return self.id
+    def _comparison_key(self) -> Tuple[str, int]:
+        return self.prefix, self.id
 
     @property
     def dimension(self) -> int:
         return self.dim
 
+    def with_dimension(self, dimension: int) -> '_BaseAncillaQid':
+        return dataclasses.replace(self, dim=dimension)
+
     def __repr__(self) -> str:
         dim_str = f', dim={self.dim}' if self.dim != 2 else ''
         prefix_str = f', prefix={self.prefix}' if self.prefix != '' else ''

cirq-core/cirq/ops/qubit_manager_test.py

@@ -27,12 +27,23 @@ def test_clean_qubits():
     q = cqi.CleanQubit(2, dim=3)
     assert q.id == 2
     assert q.dimension == 3
+    assert q.with_dimension(4) == cqi.CleanQubit(2, dim=4)
     assert str(q) == '_c(2) (d=3)'
     assert repr(q) == 'cirq.ops.CleanQubit(2, dim=3)'
 
     assert cqi.CleanQubit(1) < cqi.CleanQubit(2)
 
 
+def test_ancilla_qubits_prefix():
+    assert cqi.CleanQubit(1, prefix="1") != cqi.CleanQubit(1, prefix="2")
+    assert cqi.CleanQubit(1, prefix="1") < cqi.CleanQubit(1, prefix="2")
+    assert cqi.CleanQubit(1, prefix="1") < cqi.CleanQubit(2, prefix="1")
+    assert cqi.BorrowableQubit(1, prefix="1") != cqi.BorrowableQubit(1, prefix="2")
+    assert cqi.BorrowableQubit(1, prefix="1") < cqi.BorrowableQubit(1, prefix="2")
+    assert cqi.BorrowableQubit(1, prefix="1") < cqi.BorrowableQubit(2, prefix="1")
+    assert cqi.CleanQubit(1, prefix="1") != cqi.BorrowableQubit(1, prefix="1")
+
+
 def test_borrow_qubits():
     q = cqi.BorrowableQubit(10)
     assert q.id == 10
@@ -43,6 +54,7 @@ def test_borrow_qubits():
     q = cqi.BorrowableQubit(20, dim=4)
     assert q.id == 20
     assert q.dimension == 4
+    assert q.with_dimension(10) == cqi.BorrowableQubit(20, dim=10)
     assert str(q) == '_b(20) (d=4)'
     assert repr(q) == 'cirq.ops.BorrowableQubit(20, dim=4)'