Cache Qid instances for common types (#6371)

Review: @dstrain115

cirq-core/cirq/devices/grid_qubit.py

@@ -14,6 +14,7 @@
 
 import abc
 import functools
+import weakref
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Set, TYPE_CHECKING, Union
 from typing_extensions import Self
 
@@ -34,14 +35,6 @@ class _BaseGridQid(ops.Qid):
  _dimension: int
  _hash: Optional[int] = None
 
- def __getstate__(self):
- # Don't save hash when pickling; see #3777.
- state = self.__dict__
- if "_hash" in state:
- state = state.copy()
- del state["_hash"]
- return state
-
  def __hash__(self) -> int:
  if self._hash is None:
  self._hash = hash((self._row, self._col, self._dimension))
@@ -50,7 +43,7 @@ def __hash__(self) -> int:
  def __eq__(self, other):
  # Explicitly implemented for performance (vs delegating to Qid).
  if isinstance(other, _BaseGridQid):
- return (
+ return self is other or (
  self._row == other._row
  and self._col == other._col
  and self._dimension == other._dimension
@@ -60,7 +53,7 @@ def __eq__(self, other):
  def __ne__(self, other):
  # Explicitly implemented for performance (vs delegating to Qid).
  if isinstance(other, _BaseGridQid):
- return (
+ return self is not other and (
  self._row != other._row
  or self._col != other._col
  or self._dimension != other._dimension
@@ -178,22 +171,36 @@ class GridQid(_BaseGridQid):
  cirq.GridQid(5, 4, dimension=2)
  """
 
- def __init__(self, row: int, col: int, *, dimension: int) -> None:
- """Initializes a grid qid at the given row, col coordinate
+ # Cache of existing GridQid instances, returned by __new__ if available.
+ # Holds weak references so instances can still be garbage collected.
+ _cache = weakref.WeakValueDictionary[Tuple[int, int, int], 'cirq.GridQid']()
+
+ def __new__(cls, row: int, col: int, *, dimension: int) -> 'cirq.GridQid':
+ """Creates a grid qid at the given row, col coordinate
 
  Args:
  row: the row coordinate
  col: the column coordinate
  dimension: The dimension of the qid's Hilbert space, i.e.
  the number of quantum levels.
  """
- self.validate_dimension(dimension)
- self._row = row
- self._col = col
- self._dimension = dimension
+ key = (row, col, dimension)
+ inst = cls._cache.get(key)
+ if inst is None:
+ cls.validate_dimension(dimension)
+ inst = super().__new__(cls)
+ inst._row = row
+ inst._col = col
+ inst._dimension = dimension
+ cls._cache[key] = inst
+ return inst
+
+ def __getnewargs_ex__(self):
+ """Returns a tuple of (args, kwargs) to pass to __new__ when unpickling."""
+ return (self._row, self._col), {"dimension": self._dimension}
 
  def _with_row_col(self, row: int, col: int) -> 'GridQid':
- return GridQid(row, col, dimension=self.dimension)
+ return GridQid(row, col, dimension=self._dimension)
 
  @staticmethod
  def square(diameter: int, top: int = 0, left: int = 0, *, dimension: int) -> List['GridQid']:
@@ -290,16 +297,16 @@ def from_diagram(diagram: str, dimension: int) -> List['GridQid']:
  return [GridQid(*c, dimension=dimension) for c in coords]
 
  def __repr__(self) -> str:
- return f"cirq.GridQid({self._row}, {self._col}, dimension={self.dimension})"
+ return f"cirq.GridQid({self._row}, {self._col}, dimension={self._dimension})"
 
  def __str__(self) -> str:
- return f"q({self._row}, {self._col}) (d={self.dimension})"
+ return f"q({self._row}, {self._col}) (d={self._dimension})"
 
  def _circuit_diagram_info_(
  self, args: 'cirq.CircuitDiagramInfoArgs'
  ) -> 'cirq.CircuitDiagramInfo':
  return protocols.CircuitDiagramInfo(
- wire_symbols=(f"({self._row}, {self._col}) (d={self.dimension})",)
+ wire_symbols=(f"({self._row}, {self._col}) (d={self._dimension})",)
  )
 
  def _json_dict_(self) -> Dict[str, Any]:
@@ -325,11 +332,31 @@ class GridQubit(_BaseGridQid):
 
  _dimension = 2
 
- def __init__(self, row: int, col: int) -> None:
-  self._row = row
-  self._col = col
+ # Cache of existing GridQubit instances, returned by __new__ if available.
+ # Holds weak references so instances can still be garbage collected.
+ _cache = weakref.WeakValueDictionary[Tuple[int, int], 'cirq.GridQubit']()
 
- def _with_row_col(self, row: int, col: int):
+ def __new__(cls, row: int, col: int) -> 'cirq.GridQubit':
+ """Creates a grid qubit at the given row, col coordinate
+
+ Args:
+ row: the row coordinate
+ col: the column coordinate
+ """
+ key = (row, col)
+ inst = cls._cache.get(key)
+ if inst is None:
+ inst = super().__new__(cls)
+ inst._row = row
+ inst._col = col
+ cls._cache[key] = inst
+ return inst
+
+ def __getnewargs__(self):
+ """Returns a tuple of args to pass to __new__ when unpickling."""
+ return (self._row, self._col)
+
+ def _with_row_col(self, row: int, col: int) -> 'GridQubit':
  return GridQubit(row, col)
 
  def _cmp_tuple(self):

cirq-core/cirq/devices/grid_qubit_test.py

@@ -40,11 +40,29 @@ def test_eq():
  eq.make_equality_group(lambda: cirq.GridQid(0, 0, dimension=3))
 
 
-def test_pickled_hash():
- q = cirq.GridQubit(3, 4)
- q_bad = cirq.GridQubit(3, 4)
+def test_grid_qubit_pickled_hash():
+ # Use a large number that is unlikely to be used by any other tests.
+ row, col = 123456789, 2345678910
+ q_bad = cirq.GridQubit(row, col)
+ cirq.GridQubit._cache.pop((row, col))
+ q = cirq.GridQubit(row, col)
+ _test_qid_pickled_hash(q, q_bad)
+
+
+def test_grid_qid_pickled_hash():
+ # Use a large number that is unlikely to be used by any other tests.
+ row, col = 123456789, 2345678910
+ q_bad = cirq.GridQid(row, col, dimension=3)
+ cirq.GridQid._cache.pop((row, col, 3))
+ q = cirq.GridQid(row, col, dimension=3)
+ _test_qid_pickled_hash(q, q_bad)
+
+
+def _test_qid_pickled_hash(q: 'cirq.Qid', q_bad: 'cirq.Qid') -> None:
+ """Test that hashes are not pickled with Qid instances."""
+ assert q_bad is not q
  _ = hash(q_bad) # compute hash to ensure it is cached.
- q_bad._hash = q_bad._hash + 1
+ q_bad._hash = q_bad._hash + 1 # type: ignore[attr-defined]
  assert q_bad == q
  assert hash(q_bad) != hash(q)
  data = pickle.dumps(q_bad)

cirq-core/cirq/devices/line_qubit.py

@@ -14,7 +14,8 @@
 
 import abc
 import functools
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, TYPE_CHECKING, Union
+import weakref
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
 from typing_extensions import Self
 
 from cirq import ops, protocols
@@ -31,14 +32,6 @@ class _BaseLineQid(ops.Qid):
  _dimension: int
  _hash: Optional[int] = None
 
- def __getstate__(self):
- # Don't save hash when pickling; see #3777.
- state = self.__dict__
- if "_hash" in state:
- state = state.copy()
- del state["_hash"]
- return state
-
  def __hash__(self) -> int:
  if self._hash is None:
  self._hash = hash((self._x, self._dimension))
@@ -47,13 +40,15 @@ def __hash__(self) -> int:
  def __eq__(self, other):
  # Explicitly implemented for performance (vs delegating to Qid).
  if isinstance(other, _BaseLineQid):
- return self._x == other._x and self._dimension == other._dimension
+ return self is other or (self._x == other._x and self._dimension == other._dimension)
  return NotImplemented
 
  def __ne__(self, other):
  # Explicitly implemented for performance (vs delegating to Qid).
  if isinstance(other, _BaseLineQid):
- return self._x != other._x or self._dimension != other._dimension
+ return self is not other and (
+ self._x != other._x or self._dimension != other._dimension
+ )
  return NotImplemented
 
  def _comparison_key(self):
@@ -154,17 +149,31 @@ class LineQid(_BaseLineQid):
 
  """
 
- def __init__(self, x: int, dimension: int) -> None:
+ # Cache of existing LineQid instances, returned by __new__ if available.
+ # Holds weak references so instances can still be garbage collected.
+ _cache = weakref.WeakValueDictionary[Tuple[int, int], 'cirq.LineQid']()
+
+ def __new__(cls, x: int, dimension: int) -> 'cirq.LineQid':
  """Initializes a line qid at the given x coordinate.
 
  Args:
  x: The x coordinate.
  dimension: The dimension of the qid's Hilbert space, i.e.
  the number of quantum levels.
  """
- self.validate_dimension(dimension)
- self._x = x
- self._dimension = dimension
+ key = (x, dimension)
+ inst = cls._cache.get(key)
+ if inst is None:
+ cls.validate_dimension(dimension)
+ inst = super().__new__(cls)
+ inst._x = x
+ inst._dimension = dimension
+ cls._cache[key] = inst
+ return inst
+
+ def __getnewargs__(self):
+ """Returns a tuple of args to pass to __new__ when unpickling."""
+ return (self._x, self._dimension)
 
  def _with_x(self, x: int) -> 'LineQid':
  return LineQid(x, dimension=self._dimension)
@@ -246,13 +255,26 @@ class LineQubit(_BaseLineQid):
 
  _dimension = 2
 
- def __init__(self, x: int) -> None:
- """Initializes a line qubit at the given x coordinate.
+ # Cache of existing LineQubit instances, returned by __new__ if available.
+ # Holds weak references so instances can still be garbage collected.
+ _cache = weakref.WeakValueDictionary[int, 'cirq.LineQubit']()
+
+ def __new__(cls, x: int) -> 'cirq.LineQubit':
+ """Initializes a line qid at the given x coordinate.
 
  Args:
  x: The x coordinate.
  """
- self._x = x
+ inst = cls._cache.get(x)
+ if inst is None:
+ inst = super().__new__(cls)
+ inst._x = x
+ cls._cache[x] = inst
+ return inst
+
+ def __getnewargs__(self):
+ """Returns a tuple of args to pass to __new__ when unpickling."""
+ return (self._x,)
 
  def _with_x(self, x: int) -> 'LineQubit':
  return LineQubit(x)

cirq-core/cirq/devices/line_qubit_test.py

@@ -15,6 +15,7 @@
 import pytest
 
 import cirq
+from cirq.devices.grid_qubit_test import _test_qid_pickled_hash
 
 
 def test_init():
@@ -67,6 +68,24 @@ def test_cmp_failure():
  _ = cirq.LineQid(1, 3) < 0
 
 
+def test_line_qubit_pickled_hash():
+ # Use a large number that is unlikely to be used by any other tests.
+ x = 1234567891011
+ q_bad = cirq.LineQubit(x)
+ cirq.LineQubit._cache.pop(x)
+ q = cirq.LineQubit(x)
+ _test_qid_pickled_hash(q, q_bad)
+
+
+def test_line_qid_pickled_hash():
+ # Use a large number that is unlikely to be used by any other tests.
+ x = 1234567891011
+ q_bad = cirq.LineQid(x, dimension=3)
+ cirq.LineQid._cache.pop((x, 3))
+ q = cirq.LineQid(x, dimension=3)
+ _test_qid_pickled_hash(q, q_bad)
+
+
 def test_is_adjacent():
  assert cirq.LineQubit(1).is_adjacent(cirq.LineQubit(2))
  assert cirq.LineQubit(1).is_adjacent(cirq.LineQubit(0))