fix: correct rare-case bugs from adversarial audit across deadlocks, caching, serialization, memory safety, and Python frontend

.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
     hooks:
       - id: cpplint
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.16.0
+    rev: v0.16.1
     hooks:
       - id: ruff-check
         args: [--fix, --exit-non-zero-on-fix]

CHANGELOG.md

@@ -22,6 +22,49 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Define `Py_GIL_DISABLED` for free-threaded debug builds on Windows when building the C extension to work around an upstream CMake `FindPython` bug by [@XuehaiPan](https://github.com/XuehaiPan) in [#285](https://github.com/metaopt/optree/pull/285).
+- Fix a deadlock when registering or unregistering a pytree node concurrently with flattening, caused by releasing the GIL while holding the registry lock by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a deadlock when a duplicate registration or an unregistration of an unregistered type formats its error message, which runs the type's `__repr__` while the registry lock is held and wedges any thread that is flattening by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix registering an already-registered `collections.namedtuple` subclass or `PyStructSequence` emitting the override warning even though the registration is rejected, which raised `UserWarning` instead of `ValueError` under warnings-as-errors by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a deadlock when unregistering a pytree node drops the last reference to its flatten or unflatten function, so a `__del__` or weakref callback re-entered the registry while its lock was held by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a deadlock while flattening on free-threading builds, caused by acquiring the non-recursive dictionary insertion order lock twice by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a spurious `SystemError` from `optree._C.get_registry_size()` when a concurrent registration slipped between its two reads by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `optree.tree_iter()` hanging uninterruptibly when the `is_leaf` predicate or a custom flatten function advances the same iterator, now reported as a `RuntimeError` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the type caches handing an interpreter a value owned by another one, by keying them on the interpreter in addition to the type address by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the type caches retaining entries for a finalized interpreter, which leaked immortal keys and could be inherited by a fresh interpreter reusing the same ID by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a deadlock in the `PyStructSequence` field cache, caused by re-acquiring the GIL while still holding the cache lock by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the type caches publishing an entry that a failed cleanup registration would orphan, leaving a value owned by an interpreter with no callback able to evict it, whether the publish preceded that registration or raced another thread still performing it by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `optree.structseq_fields()` and the struct sequence accessors reporting the trailing hidden field names for unnamed sequence slots, by mapping each field to the slot its member offset names; the pure Python fallback, which cannot read those offsets, recovers the positions from a probe instance instead of assuming the unnamed slots trail the named ones by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the treespec repr printing the raw unnamed-field marker as if it were a keyword argument, now rendered as `<unnamed@N>` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `StructSequenceEntry` reporting an unnamed sequence slot as `field='unnamed field'`, which reads as a real field name, now reported by its index by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeSpec.__setstate__()` accepting a malformed state, which could read out of bounds or abort the interpreter when the treespec was later used by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeSpec.__getstate__()` returning the treespec's internal mutable containers, so mutating the pickled state corrupted an otherwise immutable treespec by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeSpec.__setstate__()` borrowing the key lists from the state it was given rather than copying them, so mutating that state afterwards silently reordered the restored treespec by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix pickling a `PyTreeSpec` with protocol 0 or 1 aborting the interpreter, by reducing through `copyreg.__newobj__` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeSpec.compose()`, `PyTreeSpec.broadcast_to_common_suffix()`, `treespec_transform()`, and `treespec_from_collection()` silently rebinding a custom node to a different registration when an empty namespace adopted a non-empty one by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeSpec.broadcast_to_common_suffix()` sorting the argument treespec's dictionary keys in place while building its key-mismatch error message, corrupting a treespec the caller still holds by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the dictionary key order being lost when a treespec built under the global namespace is promoted to an insertion-ordered namespace by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `treespec_from_collection()` on a leaf reporting an escalated `UserWarning` as a confusing `SystemError`, by checking the return value of `PyErr_WarnEx()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `treespec_from_collection()` keeping the caller's namespace on a leaf or `None` root, which made otherwise identical treespecs compare unequal by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `treespec_is_prefix()` and `treespec_is_suffix()` comparing against a stale subtree when a dictionary node's keys had been reordered by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the tree iterator never reporting or clearing its `is_leaf` predicate to the garbage collector, leaking any reference cycle that passes through it by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a list shrunk by an `is_leaf` predicate part way through flattening being read out of bounds on Python versions before 3.13.0a4, now raising `IndexError` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a reference cycle passing through a registered custom type not being collectable once the registry no longer holds the registration, by reporting the registration's members to the garbage collector when a treespec's own nodes hold every reference to it; a cycle spanning several treespecs remains uncollectable by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a deeply nested treespec overflowing the native stack and crashing in `treespec_paths()`, `treespec_accessors()`, and `PyTreeSpec.broadcast_to_common_suffix()` instead of raising `RecursionError` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `DataclassEntry` and `AttrsEntry` resolving an integer entry against every declared field rather than the children the registration emitted, which returned the wrong attribute for a class holding a metadata or non-`init` field, now read from the record `register_node()` leaves on the class itself by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `optree.dataclasses.register_node()` and `optree.integrations.attrs.register_node()` leaving a class in a half-registered state that could never be registered again: the registry entry and the field marker are now committed together, and a failure of either rolls the other back by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `optree.dataclasses.register_node()` silently dropping `InitVar` pseudo-fields, which are neither children nor metadata and cannot round-trip, now rejected with a pointer to the generic API by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `tree_broadcast_common()` and `broadcast_common()` applying the caller's `is_leaf` predicate to an internal sentinel tree, which could raise from the predicate or collapse a filled subtree and under-replicate by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeEntry` equality and hashing comparing the bytecode of `__call__()` and `codify()` rather than the methods themselves, so two entry classes that happened to share an implementation compared equal by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `PyTreeAccessor` overriding `__eq__()` while inheriting `tuple.__ne__()`, so comparing one with a plain tuple of equal entries reported `False` for both `==` and `!=` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `GetAttrEntry.codify()` emitting invalid attribute access for a field name that is not an identifier, now rendered as a `getattr()` call by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `prefix_errors()` raising `AssertionError` instead of reporting an error for a custom node whose per-instance entries differ while its metadata does not, a pair `broadcast_prefix()` accepts by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `prefix_errors()` raising `TypeError` while formatting a dictionary key mismatch when the keys have different types, now ordered with `total_order_sorted()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `tree_transpose_map()`, `tree_transpose_map_with_path()`, `tree_transpose_map_with_accessor()`, and `tree_partition()` rejecting a leafless outer structure when an explicit `inner_treespec` leaves nothing to infer by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix the `tree_broadcast_map()` family flattening a custom node twice when called with a single input tree, which broke a one-shot flatten function and contradicted the documented equivalence with `tree_map()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `treespec_entry()` and `treespec_child()` annotating their index as `int` while the runtime accepts any `SupportsInt` or `SupportsIndex` object by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `optree.utils.total_order_sorted()` mistaking a `TypeError` raised by the caller's `key` callback for a comparison failure and silently returning the sequence unsorted, and calling the callback twice per element by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix a dictionary whose keys cannot be compared being flattened in a partially sorted order instead of insertion order, when a comparison raised part way through the sort by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
+- Fix `tree_broadcast_common()` documenting that its results share one structure and `broadcast_common()` documenting that it returns two pytrees rather than two lists of leaves by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290).
 
 ### Removed
 

docs/source/spelling_wordlist.txt

@@ -2,6 +2,7 @@ abc
 ABCMeta
 accessor
 accessors
+antisymmetric
 arg
 args
 arity
@@ -70,7 +71,9 @@ optree
 OrderedDict
 param
 params
+picklable
 pragma
+preorder
 py
 pypy
 pytree

include/optree/hashing.h

@@ -24,7 +24,7 @@ limitations under the License.
 
 #include <pybind11/pybind11.h>
 
-#include "optree/pymacros.h"  // Py_ALWAYS_INLINE
+#include "optree/pymacros.h"  // Py_ALWAYS_INLINE, interpid_t
 
 namespace py = pybind11;
 
@@ -89,6 +89,26 @@ struct std::not_equal_to<std::pair<std::string, py::handle>> {
                std::not_equal_to<py::handle>{}(lhs.second, rhs.second);
     }
 };
+template <>
+struct std::equal_to<std::pair<interpid_t, py::handle>> {
+    using is_transparent = void;
+    inline constexpr Py_ALWAYS_INLINE bool operator()(const std::pair<interpid_t, py::handle> &lhs,
+                                                      const std::pair<interpid_t, py::handle> &rhs)
+        const noexcept(noexcept(std::equal_to<interpid_t>{}(lhs.first, rhs.first))) {
+        return std::equal_to<interpid_t>{}(lhs.first, rhs.first) &&
+               std::equal_to<py::handle>{}(lhs.second, rhs.second);
+    }
+};
+template <>
+struct std::not_equal_to<std::pair<interpid_t, py::handle>> {
+    using is_transparent = void;
+    inline constexpr Py_ALWAYS_INLINE bool operator()(const std::pair<interpid_t, py::handle> &lhs,
+                                                      const std::pair<interpid_t, py::handle> &rhs)
+        const noexcept(noexcept(std::not_equal_to<interpid_t>{}(lhs.first, rhs.first))) {
+        return std::not_equal_to<interpid_t>{}(lhs.first, rhs.first) ||
+               std::not_equal_to<py::handle>{}(lhs.second, rhs.second);
+    }
+};
 template <class T, class U>
 struct std::hash<std::pair<T, U>> {
     using is_transparent = void;

include/optree/pymacros.h

@@ -71,6 +71,19 @@ inline constexpr Py_ALWAYS_INLINE bool Py_IsConstant(PyObject *x) noexcept {
 }
 #define Py_IsConstant(x) Py_IsConstant(x)
 
+// `PyStructSequence_UnnamedField` is declared `extern` with hidden visibility (so it is not an
+// exported dynamic symbol for extension modules) before Python 3.11.0a2, where it became
+// `PyAPI_DATA`. Referencing it directly leaves an undefined symbol that makes the module fail to
+// import on those versions. Its value is the stable marker "unnamed field", and callers only ever
+// use it by value (never by pointer identity), so fall back to that literal there.
+inline const char *PyStructSequenceUnnamedField() noexcept {
+#if PY_VERSION_HEX >= 0x030B00A2  // Python 3.11.0a2
+    return PyStructSequence_UnnamedField;
+#else
+    return "unnamed field";
+#endif
+}
+
 using interpid_t = decltype(PyInterpreterState_GetID(nullptr));
 
 #if defined(PYBIND11_HAS_SUBINTERPRETER_SUPPORT) &&                                                \
@@ -81,15 +94,17 @@ using interpid_t = decltype(PyInterpreterState_GetID(nullptr));
 }
 
 [[nodiscard]] inline interpid_t GetCurrentPyInterpreterID() {
+    // This sits on the flatten path (twice per leaf, via the `Is*Instance` caches), so it tests the
+    // documented failure returns rather than probing the global error state.
     PyInterpreterState *interp = PyInterpreterState_Get();
-    if (PyErr_Occurred() != nullptr) [[unlikely]] {
-        throw py::error_already_set();
-    }
     if (interp == nullptr) [[unlikely]] {
+        if (PyErr_Occurred() != nullptr) [[unlikely]] {
+            throw py::error_already_set();
+        }
         throw std::runtime_error("Failed to get the current Python interpreter state.");
     }
     const interpid_t interpid = PyInterpreterState_GetID(interp);
-    if (PyErr_Occurred() != nullptr) [[unlikely]] {
+    if (interpid < 0) [[unlikely]] {
         throw py::error_already_set();
     }
     return interpid;