FIX: stabilize caching hashes (#515)

* DOC: document hash stability on the public interface
* FIX: make expression hashes independent of SymPy cache state
* FIX: put sets and dicts in a fixed order in make_hashable
* FIX: sort substitution mappings in cached wrappers
* MAINT: hash pickled bytes directly into the digest
* MAINT: use positional-only arguments and type variables

Author: redeboer

src/ampform/sympy/_cache.py

@@ -7,26 +7,30 @@
    import ampform.sympy._cache
 """
 
+# cspell:ignore pickler
 from __future__ import annotations
 
 import hashlib
 import inspect
+import io
 import logging
 import os
 import pickle  # ruff: ignore[suspicious-pickle-import]
 import re
 import sys
 import tempfile
 from collections import abc
+from contextlib import suppress
 from functools import cache, wraps
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
-from typing import TYPE_CHECKING, overload
+from typing import TYPE_CHECKING, NamedTuple, overload
 
+import sympy as sp
 from frozendict import frozendict
 
 if TYPE_CHECKING:
-    from collections.abc import Hashable
+    from collections.abc import Hashable, Iterable
     from io import BufferedReader
 
     from _typeshed import SupportsWrite
@@ -209,7 +213,7 @@ def _get_cache_dir() -> Path:
 
 
 @cache
-def _warn_once(msg):
+def _warn_once(msg, /):
     _LOGGER.warning(msg)
 
 
@@ -239,31 +243,92 @@ def get_system_cache_directory() -> str:
 
 
 @cache
-def get_readable_hash(obj: Hashable) -> str:
+def get_readable_hash(obj: Hashable, /) -> str:
     """Get a human-readable hash of any hashable Python object.
 
+    The hash follows from the value of the object, not from the identity of the parts it
+    is built from. Two SymPy expressions that compare equal therefore hash the same, no
+    matter what was constructed before them or in which process, which is what makes the
+    hash usable as a cache key in :func:`.cache_to_disk`.
+
+    A `set` or `dict` handed to this function directly is serialized in iteration order,
+    which depends on :code:`PYTHONHASHSEED` when its elements or keys are `str`. Pass it
+    through :func:`.make_hashable` first, as :func:`.cache_to_disk` does, to put it in a
+    fixed order.
+
     Args:
         obj: Any hashable object, mutable or immutable, to be hashed.
     """
-    b = to_bytes(obj)
-    h = hashlib.md5(b, usedforsecurity=False)
-    return h.hexdigest()
+    return hashlib.md5(to_bytes(obj), usedforsecurity=False).hexdigest()
+
 
+def to_bytes(obj, /) -> bytes:
+    """Convert any Python object to `bytes` with :mod:`pickle`.
 
-def to_bytes(obj) -> bytes:
-    """Convert any Python object to `bytes` using :func:`pickle.dumps`."""
+    The bytes depend on the value of the object rather than on which of its parts happen
+    to be the same instance. See :func:`.get_readable_hash` for what that guarantees.
+    """
     if isinstance(obj, (bytes, bytearray)):
         return obj
-    return pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+    stream = io.BytesIO()
+    _dump_deterministically(obj, stream)
+    return stream.getvalue()
+
+
+def _dump_deterministically(obj, /, stream: SupportsWrite[bytes]) -> None:
+    """Pickle an object so that the bytes depend on its value alone."""
+    _DeterministicPickler(stream).dump(obj)
+
+
+class _DeterministicPickler(pickle._Pickler):  # ruff: ignore[private-member-access]
+    """Pickler whose output does not depend on object identity.
+
+    A pickle stores a repeated object as a back-reference to the first time it was
+    written, keyed on identity. SymPy hands out a cached instance for equal expressions,
+    but its cache is a bounded LRU (:code:`SYMPY_CACHE_SIZE`), so which sub-expressions
+    are the *same* instance depends on what was constructed before and in which order.
+    Equal expressions are therefore mapped to one representative here, which makes the
+    back-references depend on the expression rather than on the cache state.
+
+    Two SymPy objects that compare equal always carry the same
+    :code:`_hashable_content()`, which for `.unevaluated` classes includes the arguments
+    that are not sympified, so equal objects also pickle to the same bytes.
+
+    Only those canonicalized SymPy objects are memoized. Anything else is written out in
+    full on each occurrence, because whether two equal non-SymPy objects are one shared
+    instance depends on caches and string interning elsewhere.
+
+    The pure-Python pickler is subclassed because the C accelerator does not dispatch to
+    these overrides.
+    """
+
+    def __init__(self, stream: SupportsWrite[bytes]) -> None:
+        super().__init__(stream, protocol=pickle.HIGHEST_PROTOCOL)
+        self._representatives: dict[tuple[type, Any], Any] = {}
+
+    def save(self, obj, *args, **kwargs) -> None:
+        if isinstance(obj, sp.Basic):
+            with suppress(TypeError):
+                obj = self._representatives.setdefault((type(obj), obj), obj)
+        super().save(obj, *args, **kwargs)
+
+    def memoize(self, obj) -> None:
+        if isinstance(obj, sp.Basic):
+            super().memoize(obj)
 
 
 def make_hashable(*args) -> Hashable:
     """Make a hashable object from any Python object.
 
+    Sets and dictionaries are put in a fixed order, because a `set` iterates in an order
+    that depends on :code:`PYTHONHASHSEED` for `str` elements, and a `dict` built by
+    iterating one inherits that order. A cache key built from them would otherwise
+    differ in every process.
+
     >>> make_hashable("a", 1, {"b": 2}, {3, 4})
-    ('a', 1, frozendict.frozendict({'b': 2}), frozenset({3, 4}))
+    ('a', 1, frozendict.frozendict({'b': 2}), _SortedSet(items=(3, 4)))
     >>> make_hashable({"a": {"sub-key": {1, 2, 3}, "b": [4, 5]}})
-    frozendict.frozendict({'a': frozendict.frozendict({'sub-key': frozenset({1, 2, 3}), 'b': (4, 5)})})
+    frozendict.frozendict({'a': frozendict.frozendict({'b': (4, 5), 'sub-key': _SortedSet(items=(1, 2, 3))})})
     >>> make_hashable("already-hashable")
     'already-hashable'
     """
@@ -272,15 +337,30 @@ def make_hashable(*args) -> Hashable:
     return tuple(_make_hashable_impl(x) for x in args)
 
 
-def _make_hashable_impl(obj) -> Hashable:
+def _make_hashable_impl(obj, /) -> Hashable:
     if isinstance(obj, abc.Mapping):
-        return frozendict({k: _make_hashable_impl(v) for k, v in obj.items()})
+        keys = _sorted_deterministically(obj.keys())
+        return frozendict({k: _make_hashable_impl(obj[k]) for k in keys})
     if isinstance(obj, str):
         return obj
     if isinstance(obj, abc.Iterable):
         hashable_items = (_make_hashable_impl(x) for x in obj)
         if isinstance(obj, abc.Sequence):
             return tuple(hashable_items)
         if isinstance(obj, set):
-            return frozenset(hashable_items)
+            return _SortedSet(_sorted_deterministically(hashable_items))
     return obj
+
+
+class _SortedSet(NamedTuple):
+    """Order-normalized stand-in for a `set`, tagged so it cannot pass for a `tuple`."""
+
+    items: tuple
+
+
+def _sorted_deterministically(items: Iterable[T], /) -> tuple[T, ...]:
+    """Sort items, falling back to their serialization when they cannot be compared."""
+    try:
+        return tuple(sorted(items))
+    except TypeError:
+        return tuple(sorted(items, key=to_bytes))

src/ampform/sympy/cached.py

@@ -3,6 +3,7 @@
 .. autofunction:: doit
 """
 
+# cspell:ignore srepr
 from __future__ import annotations
 
 from functools import cache
@@ -18,6 +19,7 @@
     from typing import TypeVar
 
     SympyObject = TypeVar("SympyObject", bound=sp.Basic)
+    V = TypeVar("V")
 
 
 @cache
@@ -59,8 +61,11 @@ def trigsimp(expr: sp.Expr, *args, **kwargs) -> sp.Expr:
 
 
 def subs(expr: sp.Expr, substitutions: Mapping[sp.Basic, Any]) -> sp.Expr:
-    """Call :meth:`~sympy.core.basic.Basic.subs` and cache the result to disk."""
-    return _subs_impl(expr, frozendict(substitutions))
+    """Call :meth:`~sympy.core.basic.Basic.subs` and cache the result to disk.
+
+    The order of the substitutions does not affect the cache key.
+    """
+    return _subs_impl(expr, _sorted_frozendict(substitutions))
 
 
 @cache
@@ -70,8 +75,11 @@ def _subs_impl(expr: sp.Expr, substitutions: frozendict[sp.Basic, Any]) -> sp.Ex
 
 
 def xreplace(expr: sp.Expr, substitutions: Mapping[sp.Basic, Any]) -> sp.Expr:
-    """Call :meth:`~sympy.core.basic.Basic.xreplace` and cache the result to disk."""
-    return _xreplace_impl(expr, frozendict(substitutions))
+    """Call :meth:`~sympy.core.basic.Basic.xreplace` and cache the result to disk.
+
+    The order of the substitutions does not affect the cache key.
+    """
+    return _xreplace_impl(expr, _sorted_frozendict(substitutions))
 
 
 @cache
@@ -104,7 +112,7 @@ def amplitudes(self) -> Mapping[sp.Basic, sp.Basic]: ...
 
 
 def _unfold_impl(expr: sp.Expr, substitutions: Mapping[sp.Basic, Any]) -> sp.Expr:
-    substitutions = _unfold_substitutions(frozendict(substitutions))
+    substitutions = _unfold_substitutions(_sorted_frozendict(substitutions))
     expr = doit(expr)
     return xreplace(expr, substitutions)
 
@@ -114,3 +122,25 @@ def _unfold_substitutions(
     substitutions: frozendict[sp.Basic, Any],
 ) -> frozendict[sp.Basic, Any]:
     return frozendict({k: doit(v) for k, v in substitutions.items()})
+
+
+def _sorted_frozendict(
+    substitutions: Mapping[SympyObject, V], /
+) -> frozendict[SympyObject, V]:
+    """Freeze a substitution mapping in an order that does not depend on the caller.
+
+    The disk cache is keyed on a pickle of the mapping, and a pickled mapping is written
+    in iteration order. A mapping built by iterating over
+    :attr:`~sympy.core.basic.Basic.free_symbols` inherits the order of that `set`, which
+    depends on :code:`PYTHONHASHSEED`, so the same substitutions would otherwise get a
+    different cache key in every process.
+
+    `~sympy.core.sorting.default_sort_key` does not distinguish assumptions, which would
+    leave the order of two otherwise identical symbols to the caller again, so
+    :func:`~sympy.printing.repr.srepr` breaks those ties.
+    """
+    return frozendict(sorted(substitutions.items(), key=lambda kv: _sort_key(kv[0])))
+
+
+def _sort_key(obj: Any) -> tuple[Any, str]:
+    return sp.default_sort_key(obj), sp.srepr(obj)