diff --git a/CHANGELOG.md b/CHANGELOG.md index 76b6bb6b..59f650d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,30 @@ 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 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 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 `optree.structseq_fields()` and the struct sequence accessors reporting the trailing hidden field names for unnamed sequence slots (e.g. `os.stat_result` slots 7-9 reported as `st_atime` / `st_mtime` / `st_ctime`), by mapping `tp_members` by offset instead of by position 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 `` 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 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 `PyTreeSpec.broadcast_to_common_suffix()` dropping a custom node's entries and falling back to positional ones 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_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 reference cycle passing through a registered custom type not being collectable once the registry no longer holds the registration 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 field rather than only the tree children, which returned the wrong attribute for a class holding a metadata or non-`init` field 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()` marking a class as registered before the registration succeeded, leaving the class permanently unusable after a failure 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). ### Removed diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index 3687fbbd..c7f534ce 100644 --- a/docs/source/spelling_wordlist.txt +++ b/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 diff --git a/include/optree/hashing.h b/include/optree/hashing.h index 862612ee..2d373411 100644 --- a/include/optree/hashing.h +++ b/include/optree/hashing.h @@ -19,7 +19,6 @@ limitations under the License. #include // std::size_t #include // std::hash, std::{not_,}equal_to -#include // std::string #include // std::pair #include @@ -89,6 +88,26 @@ struct std::not_equal_to> { std::not_equal_to{}(lhs.second, rhs.second); } }; +template <> +struct std::equal_to> { + using is_transparent = void; + inline constexpr Py_ALWAYS_INLINE bool operator()(const std::pair &lhs, + const std::pair &rhs) + const noexcept(noexcept(std::equal_to{}(lhs.first, rhs.first))) { + return std::equal_to{}(lhs.first, rhs.first) && + std::equal_to{}(lhs.second, rhs.second); + } +}; +template <> +struct std::not_equal_to> { + using is_transparent = void; + inline constexpr Py_ALWAYS_INLINE bool operator()(const std::pair &lhs, + const std::pair &rhs) + const noexcept(noexcept(std::not_equal_to{}(lhs.first, rhs.first))) { + return std::not_equal_to{}(lhs.first, rhs.first) || + std::not_equal_to{}(lhs.second, rhs.second); + } +}; template struct std::hash> { using is_transparent = void; diff --git a/include/optree/pymacros.h b/include/optree/pymacros.h index e0469047..f326a94f 100644 --- a/include/optree/pymacros.h +++ b/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) && \ diff --git a/include/optree/pytypes.h b/include/optree/pytypes.h index a7cac3ea..57ba40b2 100644 --- a/include/optree/pytypes.h +++ b/include/optree/pytypes.h @@ -17,11 +17,15 @@ limitations under the License. #pragma once +#include // std::size_t #include // std::rethrow_exception, std::current_exception +#include // std::optional #include // std::string -#include // std::enable_if_t, std::is_base_of_v +#include // std::enable_if_t, std::is_same_v, std::is_base_of_v, std::conditional_t #include // std::unordered_map -#include // std::move, std::pair, std::make_pair +#include // std::unordered_set +#include // std::forward, std::pair, std::make_pair, std::move +#include // std::vector #include @@ -34,6 +38,7 @@ limitations under the License. #include "optree/hashing.h" #include "optree/pymacros.h" +#include "optree/stdutils.h" #include "optree/synchronization.h" namespace py = pybind11; @@ -49,9 +54,6 @@ inline Py_ALWAYS_INLINE std::string PyRepr(const std::string &string) { return static_cast(py::repr(py::str(string))); } -// The maximum size of the type cache. -constexpr py::ssize_t MAX_TYPE_CACHE_SIZE = 4096; - #define PyNoneTypeObject \ (py::reinterpret_borrow(reinterpret_cast(Py_TYPE(Py_None)))) #define PyTupleTypeObject \ @@ -213,6 +215,149 @@ inline Py_ALWAYS_INLINE void AssertExactDeque(const py::handle &object) { } } +// A process-global cache mapping a Python object (in practice a type) to a value computed from it, +// e.g. whether a type is a namedtuple. It is a function-local static shared by every interpreter +// and outlives the Python runtime. +// +// Entries are keyed by `(interpreter_id, object address)` rather than the address alone, because a +// shared key can map to a per-interpreter-owned result: `int` is immortal and lives at the same +// address in every interpreter, while a computed `py::tuple` belongs to the interpreter that made +// it. An address-only key would hand that value to another interpreter to use after the owner frees +// it on finalization. +// +// A per-entry weakref evicts an entry when its key is collected, so a later key reusing that +// address cannot read a stale value. A per-interpreter `atexit` callback clears that interpreter's +// entries, covering what the weakref cannot: an immortal key is never collected, and interpreter +// ids restart from 0 after a `Py_Finalize`/`Py_Initialize` cycle, so a fresh interpreter must not +// inherit a finalized one's entries. +// +// `ValueType` may be a value such as `bool`, or a pybind11 reference whose entry owns one reference +// that is dropped on eviction. +template +class WeakKeyCache { +public: + explicit WeakKeyCache(const std::size_t &max_size) : m_max_size{max_size} {} + ~WeakKeyCache() = default; + + WeakKeyCache(const WeakKeyCache &) = delete; + WeakKeyCache(WeakKeyCache &&) = delete; + WeakKeyCache &operator=(const WeakKeyCache &) = delete; + WeakKeyCache &operator=(WeakKeyCache &&) = delete; + + // Return the value cached for `key`, computing and inserting it via `compute` on a miss. + // `compute` is a nullary callable returning `ValueType`, invoked with the GIL held and the + // cache lock NOT held. + template + [[nodiscard]] ValueType LookupOrInsert(const py::handle &key, Compute &&compute) { + // Read the interpreter id (part of the cache key) while the GIL is still held, before the + // read lock below releases it: `GetCurrentPyInterpreterID()` needs a valid thread state. + const interpid_t interpreter_id = GetCurrentPyInterpreterID(); + const CacheKey cache_key{interpreter_id, key}; + std::optional cached_value{}; + { +#if !defined(Py_GIL_DISABLED) + const py::gil_scoped_release_simple gil_release{}; +#endif + const scoped_read_lock lock{m_mutex}; + const auto it = m_cache.find(cache_key); + if (it != m_cache.end()) [[likely]] { + cached_value = it->second; + } + } + // The read lock is released and the GIL re-acquired (in that destruction order) BEFORE the + // borrowed object is touched, so the GIL is never (re-)acquired while the lock is held. + // Doing so would invert the lock order against the weakref eviction callback (which holds + // the GIL, then takes the write lock) and could deadlock. `key` stays alive for the whole + // call, so its entry cannot be evicted and `cached_value` stays valid. + if (cached_value.has_value()) [[likely]] { + if constexpr (std::is_same_v) { + // A value or a `py::handle`: the stored type is the value type. + return *cached_value; + } else { + // An owning object stored as a borrowed `py::handle`: return a fresh owning borrow. + return py::reinterpret_borrow(*cached_value); + } + } + + ValueType value = std::forward(compute)(); + bool inserted = false; + bool register_cleanup = false; + { +#if !defined(Py_GIL_DISABLED) + const py::gil_scoped_release_simple gil_release{}; +#endif + const scoped_write_lock lock{m_mutex}; + if (m_cache.size() < m_max_size) [[likely]] { + // The GIL is released here, so store the value without touching any refcount (a + // reference value is stored as a borrowed `py::handle` and owned by the `inc_ref()` + // below). + inserted = m_cache.emplace(cache_key, StoredType{value}).second; + } + register_cleanup = m_registered_interpids.insert(interpreter_id).second; + } + { + // The GIL is held here, so we can safely register the `atexit` callback and increment + // the reference count and create the weakref. + if (register_cleanup) [[unlikely]] { + RegisterInterpreterCleanup(interpreter_id); + } + if (inserted) [[likely]] { + if constexpr (kValueIsPyReference) { + value.inc_ref(); + } + (void)py::weakref(key, + py::cpp_function([this, cache_key](py::handle weakref) -> void { + const scoped_write_lock lock{m_mutex}; + const auto it = m_cache.find(cache_key); + if (it != m_cache.end()) [[likely]] { + if constexpr (kValueIsPyReference) { + it->second.dec_ref(); + } + m_cache.erase(it); + } + weakref.dec_ref(); + })) + .release(); + } + } + return value; + } + +private: + static constexpr bool kValueIsPyReference = std::is_base_of_v; + using StoredType = std::conditional_t; + using CacheKey = std::pair; + + // Register (once per interpreter, with the GIL held and WITHOUT the cache lock, mirroring + // `PyTreeTypeRegistry::Init`) an `atexit` callback that evicts this interpreter's entries on + // shutdown. + void RegisterInterpreterCleanup(const interpid_t &interpreter_id) { + auto atexit_register = py::getattr(py::module_::import("atexit"), "register"); + atexit_register(py::cpp_function([this, interpreter_id]() -> void { + const scoped_write_lock lock{m_mutex}; + for (auto it = m_cache.begin(); it != m_cache.end();) { + if (it->first.first == interpreter_id) [[likely]] { + if constexpr (kValueIsPyReference) { + it->second.dec_ref(); + } + it = m_cache.erase(it); + } else [[unlikely]] { + ++it; + } + } + m_registered_interpids.erase(interpreter_id); + })); + } + + std::unordered_map m_cache{}; + std::unordered_set m_registered_interpids{}; + const std::size_t m_max_size{}; + mutable read_write_mutex m_mutex{}; +}; + +// The maximum size of a type cache. +constexpr std::size_t MAX_TYPE_CACHE_SIZE = 4096; + // NOLINTNEXTLINE[readability-function-cognitive-complexity] inline bool IsNamedTupleClassImpl(const py::handle &type) { // We can only identify namedtuples heuristically, here by the presence of a _fields attribute. @@ -257,40 +402,10 @@ inline bool IsNamedTupleClass(const py::handle &type) { return false; } - static auto cache = std::unordered_map{}; - static read_write_mutex mutex{}; - bool cache_inserted = false; - - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_read_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { - return it->second; - } - } - - const bool result = EVALUATE_WITH_LOCK_HELD(IsNamedTupleClassImpl(type), type); - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_write_lock lock{mutex}; - if (cache.size() < MAX_TYPE_CACHE_SIZE) [[likely]] { - cache_inserted = cache.emplace(type, result).second; - } - } - if (cache_inserted) [[likely]] { - (void)py::weakref(type, py::cpp_function([type](py::handle weakref) -> void { - const scoped_write_lock lock{mutex}; - cache.erase(type); - weakref.dec_ref(); - })) - .release(); - } - return result; + static WeakKeyCache cache{MAX_TYPE_CACHE_SIZE}; + return cache.LookupOrInsert(type, [&type]() -> bool { + return EVALUATE_WITH_LOCK_HELD(IsNamedTupleClassImpl(type), type); + }); } inline Py_ALWAYS_INLINE bool IsNamedTupleInstance(const py::handle &object) { return IsNamedTupleClass(py::type::handle_of(object)); @@ -368,40 +483,10 @@ inline bool IsStructSequenceClass(const py::handle &type) { return false; } - static auto cache = std::unordered_map{}; - static read_write_mutex mutex{}; - bool cache_inserted = false; - - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_read_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { - return it->second; - } - } - - const bool result = EVALUATE_WITH_LOCK_HELD(IsStructSequenceClassImpl(type), type); - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_write_lock lock{mutex}; - if (cache.size() < MAX_TYPE_CACHE_SIZE) [[likely]] { - cache_inserted = cache.emplace(type, result).second; - } - } - if (cache_inserted) [[likely]] { - (void)py::weakref(type, py::cpp_function([type](py::handle weakref) -> void { - const scoped_write_lock lock{mutex}; - cache.erase(type); - weakref.dec_ref(); - })) - .release(); - } - return result; + static WeakKeyCache cache{MAX_TYPE_CACHE_SIZE}; + return cache.LookupOrInsert(type, [&type]() -> bool { + return EVALUATE_WITH_LOCK_HELD(IsStructSequenceClassImpl(type), type); + }); } inline Py_ALWAYS_INLINE bool IsStructSequenceInstance(const py::handle &object) { return IsStructSequenceClass(py::type::handle_of(object)); @@ -424,23 +509,50 @@ inline py::tuple StructSequenceGetFieldsImpl(const py::handle &type) { import sys StructSequenceFieldType = type(type(sys.version_info).major) - indices_by_name = { - name: member.index + # PyPy has no unnamed fields; a descriptor's `.index` is its sequence position. + # Map index -> name and defensively fill any missing (unnamed) slot with the marker. + names_by_index = { + member.index: name for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) } - fields.extend(sorted(indices_by_name, key=indices_by_name.get)[:cls.n_sequence_fields]) + fields.extend( + names_by_index.get(index, unnamed_field) for index in range(cls.n_sequence_fields) + ) )py", - py::dict(py::arg("cls") = type, py::arg("fields") = fields)); + py::dict(py::arg("cls") = type, + py::arg("fields") = fields, + py::arg("unnamed_field") = py::str(PyStructSequenceUnnamedField()))); return py::tuple{fields}; #else const auto n_sequence_fields = thread_safe_cast( EVALUATE_WITH_LOCK_HELD(py::getattr(type, "n_sequence_fields"), type)); const auto * const members = reinterpret_cast(type.ptr())->tp_members; + // `tp_members` lists only the NAMED fields, but each carries a byte `offset` encoding its + // sequence index, relative to `offsetof(PyTupleObject, ob_item)`. Map each member back to its + // slot by offset: indexing `members[i]` by position mislabels every slot after the first + // unnamed one (e.g. `os.stat_result` slots 7/8/9 reported as `st_atime`/`st_mtime`/`st_ctime`). py::tuple fields{n_sequence_fields}; + // Fill the named slots first, then default the remaining (unnamed) slots to the marker. + // Pre-filling every slot with the marker and overwriting the named ones would leak each + // overwritten marker: `TupleSetItem` uses `PyTuple_SET_ITEM`, which does not decref what it + // replaces. + std::vector named(n_sequence_fields, false); + for (const PyMemberDef *member = members; member != nullptr && member->name != nullptr; + // NOLINTNEXTLINE[cppcoreguidelines-pro-bounds-pointer-arithmetic] + ++member) { + const py::ssize_t index = + (member->offset - py::ssize_t_cast(offsetof(PyTupleObject, ob_item))) / + py::ssize_t_cast(sizeof(PyObject *)); + if (index >= 0 && index < n_sequence_fields) [[likely]] { + TupleSetItem(fields, index, py::str(member->name)); + named[index] = true; + } + } for (py::ssize_t i = 0; i < n_sequence_fields; ++i) { - // NOLINTNEXTLINE[cppcoreguidelines-pro-bounds-pointer-arithmetic] - TupleSetItem(fields, i, py::str(members[i].name)); + if (!named[i]) { + TupleSetItem(fields, i, py::str(PyStructSequenceUnnamedField())); + } } return fields; #endif @@ -460,48 +572,10 @@ inline py::tuple StructSequenceGetFields(const py::handle &object) { } } - static auto cache = std::unordered_map{}; - static read_write_mutex mutex{}; - bool cache_inserted = false; - - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_read_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_acquire_simple gil_acquire{}; -#endif - return py::reinterpret_borrow(it->second); - } - } - - const py::tuple fields = EVALUATE_WITH_LOCK_HELD(StructSequenceGetFieldsImpl(type), type); - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_write_lock lock{mutex}; - if (cache.size() < MAX_TYPE_CACHE_SIZE) [[likely]] { - cache_inserted = cache.emplace(type, fields).second; - } - } - if (cache_inserted) [[likely]] { - fields.inc_ref(); - (void)py::weakref(type, py::cpp_function([type](py::handle weakref) -> void { - const scoped_write_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { - it->second.dec_ref(); - cache.erase(it); - } - weakref.dec_ref(); - })) - .release(); - } - return fields; + static WeakKeyCache cache{MAX_TYPE_CACHE_SIZE}; + return cache.LookupOrInsert(type, [&type]() -> py::tuple { + return EVALUATE_WITH_LOCK_HELD(StructSequenceGetFieldsImpl(type), type); + }); } inline void TotalOrderSort(py::list &list) { // NOLINT[runtime/references] @@ -600,3 +674,12 @@ inline std::pair DictKeysDifference(const py::list & /*uniqu TotalOrderSort(extra_keys); return std::make_pair(std::move(missing_keys), std::move(extra_keys)); } + +inline py::ssize_t DistinctCount(const py::handle &iterable) { + const scoped_critical_section cs{iterable}; + const auto set = py::reinterpret_steal(PySet_New(iterable.ptr())); + if (!set) [[unlikely]] { // a non-iterable or an unhashable element raised here + throw py::error_already_set(); + } + return PySet_GET_SIZE(set.ptr()); +} diff --git a/include/optree/registry.h b/include/optree/registry.h index 5d2be589..2f5581c6 100644 --- a/include/optree/registry.h +++ b/include/optree/registry.h @@ -17,7 +17,7 @@ limitations under the License. #pragma once -#include // std::uint8_t +#include // std::uint8_t, UINT8_MAX #include // std::shared_ptr #include // std::optional, std::nullopt #include // std::string @@ -56,6 +56,12 @@ enum class PyTreeKind : std::uint8_t { NumKinds, // Number of kinds (placed at the end) }; +// A new pytree kind must keep `NumKinds` within the `std::uint8_t` underlying type; otherwise the +// enum cannot represent every value and code that narrows a kind to `std::uint8_t` (e.g. pickle +// deserialization in `PyTreeSpec::FromPicklable`) would silently wrap. +static_assert(static_cast(PyTreeKind::NumKinds) <= UINT8_MAX, + "PyTreeKind::NumKinds overflows its std::uint8_t underlying type."); + constexpr PyTreeKind kCustom = PyTreeKind::Custom; constexpr PyTreeKind kLeaf = PyTreeKind::Leaf; constexpr PyTreeKind kNone = PyTreeKind::None; @@ -133,8 +139,12 @@ class PyTreeTypeRegistry { auto ®istry1 = GetSingleton(); auto ®istry2 = GetSingleton(); - const ssize_t count1 = registry1.Size(registry_namespace); - const ssize_t count2 = registry2.Size(registry_namespace); + // Read both registries under a single lock so the two counts form a consistent snapshot. + // Two separate `Size()` calls each drop the lock, letting a concurrent (un)registration + // slip between them and spuriously trip the invariant check below. + const scoped_read_lock lock{sm_mutex}; + const ssize_t count1 = registry1.SizeImpl(registry_namespace); + const ssize_t count2 = registry2.SizeImpl(registry_namespace); EXPECT_EQ(count1, count2 + 1, "The number of registered types in the two registries should match " @@ -165,12 +175,36 @@ class PyTreeTypeRegistry { [[nodiscard]] static inline Py_ALWAYS_INLINE bool IsDictInsertionOrdered( const std::string ®istry_namespace, const bool &inherit_global_namespace = true) { + const auto flags = GetDictInsertionOrderedFlags(registry_namespace); + return inherit_global_namespace ? flags.with_inherited_global_namespace + : flags.in_current_namespace; + } + + // Whether dictionary key insertion order is preserved during flattening. Both flags are + // computed together under a single lock so callers see a consistent snapshot; this also avoids + // recursively read-locking `sm_dict_order_mutex` (which is not a recursive mutex) that would + // happen if a caller held the lock while calling `IsDictInsertionOrdered`. + struct DictInsertionOrderedFlags { + // Whether the given namespace itself preserves insertion order. + bool in_current_namespace; + // Whether the given namespace, or the inherited global namespace, preserves insertion + // order. + bool with_inherited_global_namespace; + }; + + [[nodiscard]] static inline Py_ALWAYS_INLINE DictInsertionOrderedFlags + GetDictInsertionOrderedFlags(const std::string ®istry_namespace) { const scoped_read_lock lock{sm_dict_order_mutex}; const auto interpid = GetCurrentPyInterpreterID(); const auto &namespaces = sm_dict_insertion_ordered_namespaces; - return (namespaces.find({interpid, registry_namespace}) != namespaces.end()) || - (inherit_global_namespace && namespaces.find({interpid, ""}) != namespaces.end()); + const bool in_current_namespace = + namespaces.find({interpid, registry_namespace}) != namespaces.end(); + return { + .in_current_namespace = in_current_namespace, + .with_inherited_global_namespace = + in_current_namespace || namespaces.find({interpid, ""}) != namespaces.end(), + }; } // Set the namespace to preserve the insertion order of the dictionary keys during flattening. @@ -194,16 +228,19 @@ class PyTreeTypeRegistry { template [[nodiscard]] static PyTreeTypeRegistry &GetSingleton(); - template - static void RegisterImpl(const py::object &cls, - const py::function &flatten_func, - const py::function &unflatten_func, - const py::object &path_entry_type, - const std::string ®istry_namespace); + void RegisterImpl(const py::object &cls, + const py::function &flatten_func, + const py::function &unflatten_func, + const py::object &path_entry_type, + const std::string ®istry_namespace); - template - [[nodiscard]] static RegistrationPtr UnregisterImpl(const py::object &cls, - const std::string ®istry_namespace); + [[nodiscard]] RegistrationPtr UnregisterImpl(const py::object &cls, + const std::string ®istry_namespace, + const bool &is_structsequence_class, + const bool &is_namedtuple_class); + + // Get the number of registered types without locking. The caller must hold `sm_mutex`. + [[nodiscard]] ssize_t SizeImpl(const std::optional ®istry_namespace) const; // Initialize the registry for the current interpreter. static void Init(); diff --git a/include/optree/treespec.h b/include/optree/treespec.h index 1e132a51..3d8fb2dd 100644 --- a/include/optree/treespec.h +++ b/include/optree/treespec.h @@ -39,7 +39,15 @@ using size_t = py::size_t; using ssize_t = py::ssize_t; // The maximum depth of a pytree. -#if defined(Py_DEBUG) || defined(PYPY_VERSION) || defined(MS_WINDOWS) || \ +#if defined(MS_WINDOWS) && (defined(Py_DEBUG) || defined(Py_GIL_DISABLED)) +// A debug or free-threading build on Windows combines large frames with a 1MB default stack, and +// the walkers still need room to unwind the `RecursionError` thrown at the guard. +# if PY_VERSION_HEX < 0x030A0000 // Python 3.10.0 +constexpr ssize_t MAX_RECURSION_DEPTH = 100; +# else +constexpr ssize_t MAX_RECURSION_DEPTH = 250; +# endif +#elif defined(Py_DEBUG) || defined(PYPY_VERSION) || defined(MS_WINDOWS) || \ (defined(__wasm__) || defined(__wasm32__) || defined(__wasm64__) || defined(__wasi__) || \ defined(__EMSCRIPTEN__)) constexpr ssize_t MAX_RECURSION_DEPTH = 500; @@ -364,7 +372,15 @@ class PyTreeSpec { const std::vector &traversal, const ssize_t &pos, const std::vector &other_traversal, - const ssize_t &other_pos); + const ssize_t &other_pos, + const ssize_t &depth); + + // Return the type of the first custom node that resolves to a different registration under + // `target_namespace` than it currently holds, or a null object if all custom nodes are + // consistent. Used by namespace merges (compose / broadcast) to detect an unsafe re-tagging; + // the caller raises the error. + [[nodiscard]] std::optional FindReregisteredCustomType( + const std::string &target_namespace) const; template [[nodiscard]] py::object WalkImpl( @@ -431,7 +447,7 @@ class PyTreeIter { private: py::object m_root; std::vector> m_agenda; - const std::optional m_leaf_predicate; + std::optional m_leaf_predicate; const bool m_none_is_leaf; const std::string m_namespace; const bool m_is_dict_insertion_ordered; diff --git a/optree/_C.pyi b/optree/_C.pyi index 2643289c..a35ab3af 100644 --- a/optree/_C.pyi +++ b/optree/_C.pyi @@ -37,6 +37,9 @@ from optree.typing import ( # Set if the type allows subclassing (see CPython's Include/object.h) Py_TPFLAGS_BASETYPE: Final[int] # (1UL << 10) +# The name reported for an unnamed PyStructSequence slot (see CPython's Include/structseq.h) +PyStructSequence_UnnamedField: Final[str] # 'unnamed field' + # Meta-information during build-time BUILDTIME_METADATA: Final[MappingProxyType[str, Any]] PY_VERSION: Final[str] diff --git a/optree/accessors.py b/optree/accessors.py index 181b3df8..1380e14e 100644 --- a/optree/accessors.py +++ b/optree/accessors.py @@ -325,7 +325,10 @@ def __repr__(self, /) -> str: def codify(self, /, node: str = '') -> str: """Generate code for accessing the path entry.""" - return f'{node}.{self.field}' + field = self.field + if not field.isidentifier(): # an unnamed PyStructSequence slot -> access by index + return f'{node}[{self.entry}]' + return f'{node}.{field}' class DataclassEntry(GetAttrEntry): @@ -345,11 +348,27 @@ def init_fields(self, /) -> tuple[str, ...]: """Get the init field names.""" return tuple(f.name for f in dataclasses.fields(self.type) if f.init) + @property + def children_fields(self, /) -> tuple[str, ...]: + """Get the child (pytree-node) field names.""" + # pylint: disable-next=import-outside-toplevel,cyclic-import + from optree.dataclasses import _PYTREE_NODE_DEFAULT + + # The children are the `pytree_node=True` fields (a subset of the init fields); an integer + # entry indexes these. Computed from the fields rather than the integration-only + # `__optree_dataclass_fields__` attribute, so it also works for classes registered via the + # generic `register_pytree_node`. + return tuple( + f.name + for f in dataclasses.fields(self.type) + if f.init and f.metadata.get('pytree_node', _PYTREE_NODE_DEFAULT) + ) + @property def field(self, /) -> str: """Get the field name.""" if isinstance(self.entry, int): - return self.init_fields[self.entry] + return self.children_fields[self.entry] return self.entry @property diff --git a/optree/dataclasses.py b/optree/dataclasses.py index 298397e9..4513d3b8 100644 --- a/optree/dataclasses.py +++ b/optree/dataclasses.py @@ -268,6 +268,9 @@ def dataclass( # noqa: C901,D417 # pylint: disable=function-redefined ) -> _TypeT | Callable[[_TypeT], _TypeT]: """Dataclass decorator with PyTree integration. + See also :func:`register_node` for how instances are reconstructed and when a class needs an + explicit registration with custom flatten/unflatten functions instead. + Args: cls (type or None, optional): The class to decorate. If :data:`None`, return a decorator. namespace (str): The registry namespace used for the PyTree registration. @@ -502,6 +505,14 @@ def register_node( # noqa: C901 # pylint: disable=function-redefined,too-many-b :data:`True`) are treated as children, while init fields with ``metadata['pytree_node']`` set to :data:`False` are treated as metadata. + .. note:: + These generated functions cover the straightforward case where a class merely stores its + fields. Instances are rebuilt with ``cls(**fields)``, which re-runs ``__init__`` and any + ``__post_init__``. The :func:`tree_unflatten` round-trip is exact only when they leave each + field value unchanged from what it is given. Register a class that needs any other + reconstruction behavior explicitly with :func:`optree.register_pytree_node` or + :func:`optree.register_pytree_node_class` using custom flatten/unflatten functions. + Usage:: # Direct function call @@ -549,7 +560,10 @@ def decorator(cls: _TypeT, /) -> _TypeT: raise TypeError(f'{cls!r} is not a dataclass.') if _FIELDS in cls.__dict__: raise TypeError( - f'Cannot register {cls.__name__} as a pytree node more than once.', + f'Cannot register {cls.__name__} as a pytree node more than once with ' + f'`{__name__}.register_node()`. ' + 'Use `optree.register_pytree_node()` or `optree.register_pytree_node_class()` ' + 'with explicit flatten/unflatten functions to register it in a different namespace.', ) if namespace is not GLOBAL_NAMESPACE and not isinstance(namespace, str): raise TypeError(f'The namespace must be a string, got {namespace!r}.') @@ -564,6 +578,25 @@ def decorator(cls: _TypeT, /) -> _TypeT: stacklevel=2, ) + # `InitVar` pseudo-fields are excluded from `dataclasses.fields()` (so they are neither children + # nor metadata) but remain required `__init__` parameters that `cls(**kwargs)` cannot restore. + # `__dataclass_fields__` includes them, tagged with the private `_FIELD_INITVAR` field type. + init_var_names = [ + name + # pylint: disable-next=protected-access + for name, dc_field in getattr(cls, dataclasses._FIELDS).items() # type: ignore[attr-defined] + # pylint: disable-next=protected-access + if dc_field._field_type is dataclasses._FIELD_INITVAR # type: ignore[attr-defined] + ] + if init_var_names: + raise TypeError( + f'Dataclass {cls.__name__!r} has `InitVar` field(s) {init_var_names!r}, which the ' + 'auto-generated flatten/unflatten functions cannot round-trip ' + f'(`{__name__}.register_node()` reconstructs instances with `cls(**kwargs)` and cannot ' + 'restore `InitVar` values). Use `optree.register_pytree_node()` or ' + '`optree.register_pytree_node_class()` with explicit flatten/unflatten functions instead.', + ) + children_fields = {} metadata_fields = {} for f in dataclasses.fields(cls): @@ -578,9 +611,6 @@ def decorator(cls: _TypeT, /) -> _TypeT: metadata_fields[f.name] = f children_field_names = tuple(children_fields) - children_fields_proxy = MappingProxyType(children_fields) - metadata_fields_proxy = MappingProxyType(metadata_fields) - setattr(cls, _FIELDS, (children_fields_proxy, metadata_fields_proxy)) def flatten_func( obj: _T, @@ -609,4 +639,8 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, .. path_entry_type=DataclassEntry, namespace=namespace, ) + # Mark the class as registered only AFTER `register_pytree_node()` succeeds: `_FIELDS` is the + # "already registered" guard, so setting it before a failed registration would leave the class + # impossible to register ever again. + setattr(cls, _FIELDS, (MappingProxyType(children_fields), MappingProxyType(metadata_fields))) return cls diff --git a/optree/integrations/attrs.py b/optree/integrations/attrs.py index 8a43d582..141e0604 100644 --- a/optree/integrations/attrs.py +++ b/optree/integrations/attrs.py @@ -143,11 +143,24 @@ def init_fields(self, /) -> tuple[str, ...]: """Get the init field names.""" return tuple(a.name for a in attrs.fields(self.type) if a.init) + @property + def children_fields(self, /) -> tuple[str, ...]: + """Get the child (pytree-node) field names.""" + # The children are the `pytree_node=True` fields (a subset of the init fields); an integer + # entry indexes these. Computed from the fields rather than read from the `_FIELDS` + # attribute, which is only set by the `optree.integrations.attrs` integration: an attrs + # class registered directly via the generic `register_pytree_node` has no `_FIELDS`. + return tuple( + a.name + for a in attrs.fields(self.type) + if a.init and a.metadata.get('pytree_node', _PYTREE_NODE_DEFAULT) + ) + @property def field(self, /) -> str: """Get the field name.""" if isinstance(self.entry, int): - return self.init_fields[self.entry] + return self.children_fields[self.entry] return self.entry @property @@ -232,6 +245,9 @@ def define( # pylint: disable=function-redefined This is a wrapper around :func:`attrs.define` that also registers the class as a pytree node. + See also :func:`register_node` for how instances are reconstructed and when a class needs an + explicit registration with custom flatten/unflatten functions instead. + Args: cls (type or None, optional): The class to decorate. If :data:`None`, return a decorator. namespace (str): The registry namespace used for the PyTree registration. @@ -365,6 +381,15 @@ def register_node( # noqa: C901 # pylint: disable=function-redefined,too-many-b :data:`True`) are treated as children, while init fields with ``metadata['pytree_node']`` set to :data:`False` are treated as metadata. + .. note:: + These generated functions cover the straightforward case where a class merely stores its + fields. Instances are rebuilt with ``cls(**fields)``, which re-runs the attrs-generated + ``__init__`` and hence its field converters, validation, and ``__attrs_post_init__``. The + :func:`tree_unflatten` round-trip is exact only when ``__init__`` returns each field value + unchanged from what it is given. Register a class that needs any other reconstruction + behavior explicitly with :func:`optree.register_pytree_node` or + :func:`optree.register_pytree_node_class` using custom flatten/unflatten functions. + Usage:: # Direct function call @@ -413,7 +438,10 @@ def decorator(cls: _TypeT, /) -> _TypeT: raise TypeError(f'{cls!r} is not an attrs-decorated class.') if _FIELDS in cls.__dict__: raise TypeError( - f'Cannot register {cls.__name__} as a pytree node more than once.', + f'Cannot register {cls.__name__} as a pytree node more than once with ' + f'`{__name__}.register_node()`. ' + 'Use `optree.register_pytree_node()` or `optree.register_pytree_node_class()` ' + 'with explicit flatten/unflatten functions to register it in a different namespace.', ) if namespace is not GLOBAL_NAMESPACE and not isinstance(namespace, str): raise TypeError(f'The namespace must be a string, got {namespace!r}.') @@ -444,9 +472,6 @@ def decorator(cls: _TypeT, /) -> _TypeT: children_field_names = tuple(children_fields) children_aliases = tuple(a.alias for a in children_fields.values()) - children_fields_proxy = MappingProxyType(children_fields) - metadata_fields_proxy = MappingProxyType(metadata_fields) - setattr(cls, _FIELDS, (children_fields_proxy, metadata_fields_proxy)) def flatten_func( obj: _T, @@ -475,4 +500,8 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, .. path_entry_type=AttrsEntry, namespace=namespace, ) + # Mark the class as registered only AFTER `register_pytree_node()` succeeds: `_FIELDS` is the + # "already registered" guard, so setting it before a failed registration would leave the class + # impossible to register ever again. + setattr(cls, _FIELDS, (MappingProxyType(children_fields), MappingProxyType(metadata_fields))) return cls diff --git a/optree/ops.py b/optree/ops.py index 8db8183c..fb3bdc7d 100644 --- a/optree/ops.py +++ b/optree/ops.py @@ -1727,6 +1727,13 @@ def tree_broadcast_common( ``other_tree``. The number of replicas is determined by the corresponding subtree in the suffix structure. + .. note:: + If ``is_leaf`` classifies nodes by their leaf **value** rather than by type or structure + (e.g., treating an integer tuple as a leaf), broadcasting may replicate a value into a slot + whose filled form the predicate then re-classifies, yielding two trees that re-flatten to + different structures under the same ``is_leaf``. Prefer type- or structure-based predicates + when broadcasting. + >>> tree_broadcast_common(1, [2, 3, 4]) ([1, 1, 1], [2, 3, 4]) >>> tree_broadcast_common([1, 2, 3], [4, 5, 6]) @@ -1775,9 +1782,13 @@ def tree_broadcast_common( ) def broadcast_leaves(x: T, subtree: PyTree[T]) -> PyTree[T]: + # `subtree` is a slice of `common_suffix_tree`, whose leaves are the private `sentinel` + # placeholder rather than real input values. Do not run the user's `is_leaf` on it: the + # subtree structure is already fixed by `common_suffix_treespec`, and `is_leaf` may crash or + # wrongly collapse them. Flatten fully to count how many replicas of `x` are needed. subtreespec = tree_structure( subtree, - is_leaf=is_leaf, + is_leaf=None, none_is_leaf=none_is_leaf, namespace=namespace, ) @@ -1821,6 +1832,13 @@ def broadcast_common( ``other_tree``. The number of replicas is determined by the corresponding subtree in the suffix structure. + .. note:: + If ``is_leaf`` classifies nodes by their leaf **value** rather than by type or structure + (e.g., treating an integer tuple as a leaf), broadcasting may replicate a value into a slot + whose filled form the predicate then re-classifies, yielding two trees that re-flatten to + different structures under the same ``is_leaf``. Prefer type- or structure-based predicates + when broadcasting. + >>> broadcast_common(1, [2, 3, 4]) ([1, 1, 1], [2, 3, 4]) >>> broadcast_common([1, 2, 3], [4, 5, 6]) @@ -2974,10 +2992,68 @@ def treespec_is_prefix( *, strict: bool = False, ) -> bool: - """Return whether ``treespec`` is a prefix of ``other_treespec``. + r"""Return whether ``treespec`` is a prefix of ``other_treespec``. + + A treespec is a *prefix* of another when the latter can be built by replacing some of the + former's leaves with subtrees. This is the broadcasting-compatibility relation behind + multi-input :func:`tree_map` and :func:`tree_broadcast_prefix`: a prefix treespec broadcasts + over any tree it is a prefix of, pairing each of its leaves with the corresponding subtree of + the fuller structure. See also :func:`treespec_is_suffix` and :meth:`PyTreeSpec.is_prefix`. + The prefix relation, together with :func:`treespec_is_suffix` and the ``<``, ``<=``, ``>``, + ``>=`` operators (``<`` and ``>`` are the ``strict=True`` variants), is a **preorder**: + reflexive (every treespec is a non-strict prefix of itself) and transitive, but **neither a + partial order nor a total order**. + + - Not antisymmetric (so not a partial order): ``a <= b and b <= a`` does **not** imply + ``a == b``. Metadata that does not change how children are partitioned, such as a + :class:`collections.deque` ``maxlen`` or a :class:`dict` key order, is transparent to the + prefix relation but significant to equality, so they are mutual prefixes yet unequal. + - Not total: two treespecs can be incomparable, with neither a prefix of the other. Examples + are a :class:`tuple` and a same-arity :class:`list`, or :class:`dict`\s with different key + sets. + + >>> treespec_is_prefix(tree_structure(1), tree_structure([1, 2, 3])) + True + >>> treespec_is_prefix(tree_structure([1, 2, 3]), tree_structure(1)) + False + + Mutual prefixes need not be equal, since the relation is a preorder, not a partial order. A + :class:`collections.deque` ``maxlen`` is transparent to the prefix relation but significant to + equality: + + >>> from collections import deque + >>> a = tree_structure(deque([1, 2], maxlen=2)) + >>> b = tree_structure(deque([1, 2], maxlen=5)) + >>> treespec_is_prefix(a, b) and treespec_is_prefix(b, a) + True + >>> a <= b and b <= a + True + >>> a == b + False + + Two treespecs can be incomparable, since the relation is not a total order. A :class:`tuple` + and a same-arity :class:`list` are neither a prefix of the other: + + >>> a = tree_structure((1, 2)) + >>> b = tree_structure([1, 2]) + >>> treespec_is_prefix(a, b) + False + >>> treespec_is_prefix(b, a) + False + >>> a <= b + False + >>> b <= a + False + >>> a < b + False + >>> b < a + False + >>> a == b + False + Args: treespec (PyTreeSpec): A treespec. other_treespec (PyTreeSpec): Another treespec to compare against. @@ -2999,8 +3075,19 @@ def treespec_is_suffix( ) -> bool: """Return whether ``treespec`` is a suffix of ``other_treespec``. + This is the reverse of :func:`treespec_is_prefix`: ``treespec_is_suffix(a, b)`` is equivalent to + ``treespec_is_prefix(b, a)``. Like the prefix relation, it is a **preorder**: reflexive and + transitive, but neither a partial order (mutual suffixes need not be equal) nor a total order + (two treespecs can be incomparable). See :func:`treespec_is_prefix` for the full discussion and + examples. + See also :func:`treespec_is_prefix` and :meth:`PyTreeSpec.is_suffix`. + >>> treespec_is_suffix(tree_structure([1, 2, 3]), tree_structure(1)) + True + >>> treespec_is_suffix(tree_structure(1), tree_structure([1, 2, 3])) + False + Args: treespec (PyTreeSpec): A treespec. other_treespec (PyTreeSpec): Another treespec to compare against. diff --git a/optree/typing.py b/optree/typing.py index 77dc4ce4..fd6c6d04 100644 --- a/optree/typing.py +++ b/optree/typing.py @@ -550,6 +550,11 @@ def is_structseq_class(cls: type, /) -> bool: # pylint: disable-next=line-too-long StructSequenceFieldType: type[types.MemberDescriptorType] = type(type(sys.version_info).major) # type: ignore[assignment] +# The name reported for an unnamed PyStructSequence slot; CPython's C-level marker (not a valid +# identifier, so accessors fall back to index access for such slots). +# pylint: disable-next=invalid-name +PyStructSequence_UnnamedField: str = _C.PyStructSequence_UnnamedField # 'unnamed field' + @_override_with_(_C.structseq_fields) def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: @@ -564,19 +569,33 @@ def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: raise TypeError(f'Expected an instance of PyStructSequence type, got {obj!r}.') if platform.python_implementation() == 'PyPy': # pragma: pypy cover - indices_by_name = { - name: member.index # type: ignore[attr-defined] + # PyPy has no unnamed sequence fields: a field descriptor exposes `.index` as its sequence + # position, and hidden fields have an index >= n_sequence_fields. (`n_unnamed_fields == 0`, + # see PyPy's `lib_pypy/_structseq.py`) Map index -> name and defensively fill any missing + # (i.e. unnamed) sequence slot with the marker, should that invariant ever change. + names_by_index = { + member.index: name # type: ignore[attr-defined] for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) } - fields = sorted(indices_by_name, key=indices_by_name.get) # type: ignore[arg-type] - else: # pragma: pypy no cover - fields = [ - name - for name, member in vars(cls).items() - if isinstance(member, StructSequenceFieldType) - ] - return tuple(fields[: cls.n_sequence_fields]) # type: ignore[attr-defined] + return tuple( + names_by_index.get(index, PyStructSequence_UnnamedField) + for index in range(cls.n_sequence_fields) # type: ignore[attr-defined] + ) + + # pragma: pypy no cover + # CPython's `member_descriptor` does not expose the field offset, so the exact position of an + # unnamed slot is not recoverable in pure Python (the C++ implementation maps by offset). Assume + # unnamed sequence fields trail the named ones (as in `os.stat_result`): keep the named sequence + # fields, then fill the remaining slots with the marker. + named = [ + name for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) + ] + n_sequence_fields: int = cls.n_sequence_fields # type: ignore[attr-defined] + n_unnamed_fields: int = cls.n_unnamed_fields # type: ignore[attr-defined] + return tuple(named[: n_sequence_fields - n_unnamed_fields]) + ( + (PyStructSequence_UnnamedField,) * n_unnamed_fields + ) del _tp_cache diff --git a/src/optree.cpp b/src/optree.cpp index f1dd9fd0..a5e51bdd 100644 --- a/src/optree.cpp +++ b/src/optree.cpp @@ -51,6 +51,7 @@ void BuildModule(py::module_ &mod) { // NOLINT[runtime/references] mod.doc() = "Optimized PyTree Utilities. (C extension module built from " + std::string(__FILE_RELPATH_FROM_PROJECT_ROOT__) + ")"; mod.attr("Py_TPFLAGS_BASETYPE") = py::int_(Py_TPFLAGS_BASETYPE); + mod.attr("PyStructSequence_UnnamedField") = py::str(PyStructSequenceUnnamedField()); // Meta information during build py::dict BUILDTIME_METADATA{}; @@ -510,7 +511,22 @@ void BuildModule(py::module_ &mod) { // NOLINT[runtime/references] }), "Serialization support for PyTreeSpec.", py::arg("state"), - py::pos_only()); + py::pos_only()) + .def( + "__reduce__", + [](const py::handle &self) -> py::object { + // pybind11's pickle support (`__getstate__`/`__setstate__`) reconstructs via the + // protocol >= 2 `copyreg.__newobj__` reduction. At protocol 0/1 the default + // reduction goes through `object.__new__`, which pybind11 rejects with an + // untranslated C++ exception that aborts the interpreter. Return the `__newobj__` + // reduction explicitly so every protocol reconstructs via `cls.__new__(cls)`. + const py::object newobj = py::module_::import("copyreg").attr("__newobj__"); + return py::make_tuple(newobj, + py::make_tuple(py::type::handle_of(self)), + self.attr("__getstate__")()); + }, + "Reduce the treespec to a picklable form supporting all pickle protocols.", + py::pos_only()); auto PyTreeIterTypeObject = #if defined(PYBIND11_HAS_INTERNALS_WITH_SMART_HOLDER_SUPPORT) diff --git a/src/registry.cpp b/src/registry.cpp index 5c12296c..932171de 100644 --- a/src/registry.cpp +++ b/src/registry.cpp @@ -69,9 +69,8 @@ template template PyTreeTypeRegistry &PyTreeTypeRegistry::GetSingleton(); template PyTreeTypeRegistry &PyTreeTypeRegistry::GetSingleton(); -ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_namespace) const { - const scoped_read_lock lock{sm_mutex}; - +ssize_t PyTreeTypeRegistry::SizeImpl(const std::optional ®istry_namespace) const { + // The caller must hold `sm_mutex`. ssize_t count = py::ssize_t_cast(m_registrations.size()); for (const auto &[named_type, _] : m_named_registrations) { if (!registry_namespace || named_type.first == *registry_namespace) [[likely]] { @@ -81,15 +80,17 @@ ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_name return count; } -template -/*static*/ void PyTreeTypeRegistry::RegisterImpl(const py::object &cls, - const py::function &flatten_func, - const py::function &unflatten_func, - const py::object &path_entry_type, - const std::string ®istry_namespace) { - auto ®istry = GetSingleton(); - - if (registry.m_builtins_types.find(cls) != registry.m_builtins_types.end()) [[unlikely]] { +ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_namespace) const { + const scoped_read_lock lock{sm_mutex}; + return SizeImpl(registry_namespace); +} + +void PyTreeTypeRegistry::RegisterImpl(const py::object &cls, + const py::function &flatten_func, + const py::function &unflatten_func, + const py::object &path_entry_type, + const std::string ®istry_namespace) { + if (m_builtins_types.find(cls) != m_builtins_types.end()) [[unlikely]] { throw py::value_error("PyTree type " + PyRepr(cls) + " is a built-in type and cannot be re-registered."); } @@ -101,29 +102,12 @@ template registration->unflatten_func = py::reinterpret_borrow(unflatten_func); registration->path_entry_type = py::reinterpret_borrow(path_entry_type); if (registry_namespace.empty()) [[unlikely]] { - if (!registry.m_registrations.emplace(cls, std::move(registration)).second) [[unlikely]] { + if (!m_registrations.emplace(cls, std::move(registration)).second) [[unlikely]] { throw py::value_error("PyTree type " + PyRepr(cls) + " is already registered in the global namespace."); } - if (IsStructSequenceClass(cls)) [[unlikely]] { - PyErr_WarnEx(PyExc_UserWarning, - ("PyTree type " + PyRepr(cls) + - " is a class of `PyStructSequence`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions.") - .c_str(), - /*stack_level=*/2); - } else if (IsNamedTupleClass(cls)) [[unlikely]] { - PyErr_WarnEx(PyExc_UserWarning, - ("PyTree type " + PyRepr(cls) + - " is a subclass of `collections.namedtuple`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions.") - .c_str(), - /*stack_level=*/2); - } } else [[likely]] { - if (!registry.m_named_registrations + if (!m_named_registrations .emplace(std::make_pair(registry_namespace, cls), std::move(registration)) .second) [[unlikely]] { std::ostringstream oss{}; @@ -131,27 +115,6 @@ template << PyRepr(registry_namespace) << "."; throw py::value_error(oss.str()); } - if (IsStructSequenceClass(cls)) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) - << " is a class of `PyStructSequence`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions in namespace " - << PyRepr(registry_namespace) << "."; - PyErr_WarnEx(PyExc_UserWarning, - oss.str().c_str(), - /*stack_level=*/2); - } else if (IsNamedTupleClass(cls)) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) - << " is a subclass of `collections.namedtuple`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions in namespace " - << PyRepr(registry_namespace) << "."; - PyErr_WarnEx(PyExc_UserWarning, - oss.str().c_str(), - /*stack_level=*/2); - } } } @@ -160,44 +123,86 @@ template const py::function &unflatten_func, const py::object &path_entry_type, const std::string ®istry_namespace) { - const scoped_write_lock lock{sm_mutex}; + // Emit the override warning for namedtuple / PyStructSequence subclasses BEFORE taking + // `sm_mutex`: the classification and `PyErr_WarnEx` run Python and release the GIL, and doing + // that under the write lock inverts the GIL to `sm_mutex` order against a concurrent flatten. + // Emitting before any insert also keeps registration atomic under warnings-as-errors. + if (IsStructSequenceClass(cls)) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTree type " << PyRepr(cls) + << " is a class of `PyStructSequence`, " + "which is already registered in the global namespace. " + "Override it with custom flatten/unflatten functions"; + if (!registry_namespace.empty()) [[likely]] { + oss << " in namespace " << PyRepr(registry_namespace); + } + oss << "."; + if (PyErr_WarnEx(PyExc_UserWarning, oss.str().c_str(), /*stack_level=*/2) < 0) + [[unlikely]] { + throw py::error_already_set(); + } + } else if (IsNamedTupleClass(cls)) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTree type " << PyRepr(cls) + << " is a subclass of `collections.namedtuple`, " + "which is already registered in the global namespace. " + "Override it with custom flatten/unflatten functions"; + if (!registry_namespace.empty()) [[likely]] { + oss << " in namespace " << PyRepr(registry_namespace); + } + oss << "."; + if (PyErr_WarnEx(PyExc_UserWarning, oss.str().c_str(), /*stack_level=*/2) < 0) + [[unlikely]] { + throw py::error_already_set(); + } + } + + // Acquire both singletons BEFORE `sm_mutex`, mirroring `Init`/`Clear`. Under + // `per_interpreter_gil`, `GetSingleton()` releases the GIL on every call once a subinterpreter + // has existed; doing that while holding `sm_mutex` inverts the GIL <-> `sm_mutex` lock order + // against a concurrent flatten (read lock) and deadlocks. + auto ®istry1 = GetSingleton(); + auto ®istry2 = GetSingleton(); - RegisterImpl(cls, + { + const scoped_write_lock lock{sm_mutex}; + + registry1.RegisterImpl(cls, flatten_func, unflatten_func, path_entry_type, registry_namespace); - RegisterImpl(cls, + registry2.RegisterImpl(cls, flatten_func, unflatten_func, path_entry_type, registry_namespace); - cls.inc_ref(); - flatten_func.inc_ref(); - unflatten_func.inc_ref(); - path_entry_type.inc_ref(); + cls.inc_ref(); + flatten_func.inc_ref(); + unflatten_func.inc_ref(); + path_entry_type.inc_ref(); + } } -template -/*static*/ PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::UnregisterImpl( +PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::UnregisterImpl( const py::object &cls, - const std::string ®istry_namespace) { - auto ®istry = GetSingleton(); - - if (registry.m_builtins_types.find(cls) != registry.m_builtins_types.end()) [[unlikely]] { + const std::string ®istry_namespace, + const bool &is_structsequence_class, + const bool &is_namedtuple_class) { + if (m_builtins_types.find(cls) != m_builtins_types.end()) [[unlikely]] { throw py::value_error("PyTree type " + PyRepr(cls) + " is a built-in type and cannot be unregistered."); } if (registry_namespace.empty()) [[unlikely]] { - const auto it = registry.m_registrations.find(cls); - if (it == registry.m_registrations.end()) [[unlikely]] { + const auto it = m_registrations.find(cls); + if (it == m_registrations.end()) [[unlikely]] { std::ostringstream oss{}; oss << "PyTree type " << PyRepr(cls) << " "; - if (IsStructSequenceClass(cls)) [[unlikely]] { + if (is_structsequence_class) [[unlikely]] { oss << "is a class of `PyStructSequence`, " << "which is not explicitly registered in the global namespace."; - } else if (IsNamedTupleClass(cls)) [[unlikely]] { + } else if (is_namedtuple_class) [[unlikely]] { oss << "is a subclass of `collections.namedtuple`, " << "which is not explicitly registered in the global namespace."; } else [[likely]] { @@ -206,18 +211,17 @@ template throw py::value_error(oss.str()); } RegistrationPtr registration = it->second; - registry.m_registrations.erase(it); + m_registrations.erase(it); return registration; } else [[likely]] { - const auto named_it = - registry.m_named_registrations.find(std::make_pair(registry_namespace, cls)); - if (named_it == registry.m_named_registrations.end()) [[unlikely]] { + const auto named_it = m_named_registrations.find(std::make_pair(registry_namespace, cls)); + if (named_it == m_named_registrations.end()) [[unlikely]] { std::ostringstream oss{}; oss << "PyTree type " << PyRepr(cls) << " "; - if (IsStructSequenceClass(cls)) [[unlikely]] { + if (is_structsequence_class) [[unlikely]] { oss << "is a class of `PyStructSequence`, " << "which is not explicitly registered "; - } else if (IsNamedTupleClass(cls)) [[unlikely]] { + } else if (is_namedtuple_class) [[unlikely]] { oss << "is a subclass of `collections.namedtuple`, " << "which is not explicitly registered "; } else [[likely]] { @@ -227,43 +231,70 @@ template throw py::value_error(oss.str()); } RegistrationPtr registration = named_it->second; - registry.m_named_registrations.erase(named_it); + m_named_registrations.erase(named_it); return registration; } } /*static*/ void PyTreeTypeRegistry::Unregister(const py::object &cls, const std::string ®istry_namespace) { - const scoped_write_lock lock{sm_mutex}; - - const auto registration1 = UnregisterImpl(cls, registry_namespace); - const auto registration2 = UnregisterImpl(cls, registry_namespace); - EXPECT_TRUE(registration1->type.is(registration2->type)); - EXPECT_TRUE(registration1->flatten_func.is(registration2->flatten_func)); - EXPECT_TRUE(registration1->unflatten_func.is(registration2->unflatten_func)); - EXPECT_TRUE(registration1->path_entry_type.is(registration2->path_entry_type)); - registration1->type.dec_ref(); - registration1->flatten_func.dec_ref(); - registration1->unflatten_func.dec_ref(); - registration1->path_entry_type.dec_ref(); + // Classify the type BEFORE taking `sm_mutex`. On the not-found path `UnregisterImpl` builds its + // error message from `IsStructSequenceClass` / `IsNamedTupleClass`, which run Python and + // release the GIL; calling them while holding `sm_mutex` in write mode inverts the GIL <-> + // `sm_mutex` lock order and deadlocks a concurrent flatten that holds the GIL while waiting on + // `sm_mutex` in read mode (mirrors `Register`). + const bool is_structsequence_class = IsStructSequenceClass(cls); + const bool is_namedtuple_class = IsNamedTupleClass(cls); + + // Acquire both singletons BEFORE `sm_mutex`, mirroring `Init`/`Clear` (see `Lookup`/`Register` + // for the lock-order rationale). + auto ®istry1 = GetSingleton(); + auto ®istry2 = GetSingleton(); + + { + const scoped_write_lock lock{sm_mutex}; + + const auto registration1 = registry1.UnregisterImpl(cls, + registry_namespace, + is_structsequence_class, + is_namedtuple_class); + const auto registration2 = registry2.UnregisterImpl(cls, + registry_namespace, + is_structsequence_class, + is_namedtuple_class); + EXPECT_TRUE(registration1->type.is(registration2->type)); + EXPECT_TRUE(registration1->flatten_func.is(registration2->flatten_func)); + EXPECT_TRUE(registration1->unflatten_func.is(registration2->unflatten_func)); + EXPECT_TRUE(registration1->path_entry_type.is(registration2->path_entry_type)); + registration1->type.dec_ref(); + registration1->flatten_func.dec_ref(); + registration1->unflatten_func.dec_ref(); + registration1->path_entry_type.dec_ref(); + } } template /*static*/ PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::Lookup( const py::object &cls, const std::string ®istry_namespace) { - const scoped_read_lock lock{sm_mutex}; - + // Acquire the singleton BEFORE `sm_mutex`, mirroring `Init`/`Clear`. Under + // `per_interpreter_gil`, `GetSingleton()` releases the GIL on every call once a subinterpreter + // has existed; doing that while holding `sm_mutex` inverts the GIL <-> `sm_mutex` lock order + // against a concurrent registration (write lock) and deadlocks. const auto ®istry = GetSingleton(); - if (!registry_namespace.empty()) [[unlikely]] { - const auto named_it = - registry.m_named_registrations.find(std::make_pair(registry_namespace, cls)); - if (named_it != registry.m_named_registrations.end()) [[likely]] { - return named_it->second; + + { + const scoped_read_lock lock{sm_mutex}; + if (!registry_namespace.empty()) [[unlikely]] { + const auto named_it = + registry.m_named_registrations.find(std::make_pair(registry_namespace, cls)); + if (named_it != registry.m_named_registrations.end()) [[likely]] { + return named_it->second; + } } + const auto it = registry.m_registrations.find(cls); + return it != registry.m_registrations.end() ? it->second : nullptr; } - const auto it = registry.m_registrations.find(cls); - return it != registry.m_registrations.end() ? it->second : nullptr; } template PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::Lookup( diff --git a/src/treespec/constructors.cpp b/src/treespec/constructors.cpp index 380f6b84..1664435f 100644 --- a/src/treespec/constructors.cpp +++ b/src/treespec/constructors.cpp @@ -72,9 +72,17 @@ template Node node; node.kind = PyTreeTypeRegistry::GetKind(handle, node.custom, registry_namespace); - const auto verify_children = [&handle, &node, ®istry_namespace]( - const std::vector &children, - std::vector &treespecs) -> void { + const auto dict_order_flags = + PyTreeTypeRegistry::GetDictInsertionOrderedFlags(registry_namespace); + const bool is_dict_insertion_ordered = dict_order_flags.with_inherited_global_namespace; + const bool is_dict_insertion_ordered_in_current_namespace = + dict_order_flags.in_current_namespace; + + const auto verify_children = + // NOLINTNEXTLINE[readability-function-cognitive-complexity] + [&handle, &node, ®istry_namespace, &is_dict_insertion_ordered_in_current_namespace]( + const std::vector &children, + std::vector &treespecs) -> void { for (const py::object &child : children) { if (!py::isinstance(child)) [[unlikely]] { std::ostringstream oss{}; @@ -114,7 +122,13 @@ template << ", got " << PyRepr(common_registry_namespace) << "."; throw py::value_error(oss.str()); } - } else if (node.kind != PyTreeKind::Custom) [[likely]] { + } else if (node.kind != PyTreeKind::Custom && + ((node.kind != PyTreeKind::Dict && node.kind != PyTreeKind::DefaultDict) || + !is_dict_insertion_ordered_in_current_namespace)) [[likely]] { + // Drop the namespace for namespace-independent nodes. A Dict/DefaultDict whose keys are + // kept in insertion order does depend on the namespace (see the sort at the Dict case), + // so keep it there, mirroring `Flatten` + // (`is_dict_insertion_ordered_in_current_namespace`). registry_namespace = ""; } }; @@ -122,9 +136,11 @@ template switch (node.kind) { case PyTreeKind::Leaf: { node.arity = 0; - PyErr_WarnEx(PyExc_UserWarning, - "PyTreeSpec::MakeFromCollection() is called on a leaf.", - /*stack_level=*/2); + if (PyErr_WarnEx(PyExc_UserWarning, + "PyTreeSpec::MakeFromCollection() is called on a leaf.", + /*stack_level=*/2) < 0) [[unlikely]] { + throw py::error_already_set(); + } break; } @@ -170,8 +186,7 @@ template keys = DictKeys(dict); if (node.kind != PyTreeKind::OrderedDict) [[likely]] { node.original_keys = DictFromKeys(dict); - if (!PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace)) - [[likely]] { + if (!is_dict_insertion_ordered) [[likely]] { TotalOrderSort(keys); } } @@ -272,6 +287,23 @@ template out->m_traversal.emplace_back(std::move(node)); out->m_none_is_leaf = NoneIsLeaf; out->m_namespace = registry_namespace; + // Reject a namespace promotion (an empty caller namespace adopting a child spec's namespace) + // that would rebind a custom node to a different registration than the one it holds, e.g. the + // root node, or a globally-resolved child, resolved in the global registry while the promoted + // namespace registers the type differently. The result keeps each node's original registration, + // mirroring the compose / transform / broadcast merge guards. Skipped for an empty namespace, + // which resolves every custom node globally. + if (!registry_namespace.empty()) [[unlikely]] { + if (const auto reregistered_type = out->FindReregisteredCustomType(registry_namespace)) + [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be composed into a collection: custom PyTree type " + << PyRepr(*reregistered_type) + << " resolves to a different registration in namespace " + << PyRepr(registry_namespace) << "."; + throw py::value_error(oss.str()); + } + } out->m_traversal.shrink_to_fit(); PYTREESPEC_SANITY_CHECK(*out); return out; diff --git a/src/treespec/flatten.cpp b/src/treespec/flatten.cpp index 0894824d..d608ab47 100644 --- a/src/treespec/flatten.cpp +++ b/src/treespec/flatten.cpp @@ -204,17 +204,11 @@ bool PyTreeSpec::FlattenInto(const py::handle &handle, const bool &none_is_leaf, const std::string ®istry_namespace) { bool found_custom = false; - bool is_dict_insertion_ordered = false; - bool is_dict_insertion_ordered_in_current_namespace = false; - { -#if defined(OPTREE_HAS_READ_WRITE_LOCK) - const scoped_read_lock lock{PyTreeTypeRegistry::sm_dict_order_mutex}; -#endif - is_dict_insertion_ordered = PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace); - is_dict_insertion_ordered_in_current_namespace = - PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace, - /*inherit_global_namespace=*/false); - } + const auto dict_order_flags = + PyTreeTypeRegistry::GetDictInsertionOrderedFlags(registry_namespace); + const bool is_dict_insertion_ordered = dict_order_flags.with_inherited_global_namespace; + const bool is_dict_insertion_ordered_in_current_namespace = + dict_order_flags.in_current_namespace; if (none_is_leaf) [[unlikely]] { if (!is_dict_insertion_ordered) [[likely]] { @@ -481,17 +475,11 @@ bool PyTreeSpec::FlattenIntoWithPath(const py::handle &handle, const bool &none_is_leaf, const std::string ®istry_namespace) { bool found_custom = false; - bool is_dict_insertion_ordered = false; - bool is_dict_insertion_ordered_in_current_namespace = false; - { -#if defined(OPTREE_HAS_READ_WRITE_LOCK) - const scoped_read_lock lock{PyTreeTypeRegistry::sm_dict_order_mutex}; -#endif - is_dict_insertion_ordered = PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace); - is_dict_insertion_ordered_in_current_namespace = - PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace, - /*inherit_global_namespace=*/false); - } + const auto dict_order_flags = + PyTreeTypeRegistry::GetDictInsertionOrderedFlags(registry_namespace); + const bool is_dict_insertion_ordered = dict_order_flags.with_inherited_global_namespace; + const bool is_dict_insertion_ordered_in_current_namespace = + dict_order_flags.in_current_namespace; auto stack = reserved_vector(4); if (none_is_leaf) [[unlikely]] { diff --git a/src/treespec/gc.cpp b/src/treespec/gc.cpp index f6c3700d..a030865f 100644 --- a/src/treespec/gc.cpp +++ b/src/treespec/gc.cpp @@ -43,6 +43,17 @@ namespace optree { Py_VISIT(node.node_data.ptr()); Py_VISIT(node.node_entries.ptr()); Py_VISIT(node.original_keys.ptr()); + // Report the registration's members only when this node is their sole owner. + // The registration holds one reference to each member however many nodes point at it, so + // reporting while it is shared would decrement the same object once per node and underflow + // its shadow refcount. While the registry holds the registration it also keeps the members + // alive, so skipping then leaks nothing. + if (node.custom != nullptr && node.custom.use_count() == 1) [[unlikely]] { + Py_VISIT(node.custom->type.ptr()); + Py_VISIT(node.custom->flatten_func.ptr()); + Py_VISIT(node.custom->unflatten_func.ptr()); + Py_VISIT(node.custom->path_entry_type.ptr()); + } } return 0; } @@ -58,11 +69,13 @@ namespace optree { Py_CLEAR(node.node_data.ptr()); Py_CLEAR(node.node_entries.ptr()); Py_CLEAR(node.original_keys.ptr()); + node.custom.reset(); } self.m_traversal.clear(); return 0; } +// NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ int PyTreeIter::PyTpTraverse(PyObject *self_base, visitproc visit, void *arg) { Py_VISIT(Py_TYPE(self_base)); if (!::is_holder_constructed(self_base)) [[unlikely]] { @@ -74,6 +87,11 @@ namespace optree { Py_VISIT(obj.ptr()); } Py_VISIT(self.m_root.ptr()); + if (self.m_leaf_predicate) [[likely]] { + // The leaf predicate is an owned Python callback; it must be visited so the cyclic GC can + // see reference cycles that pass through it (otherwise such cycles leak). + Py_VISIT(self.m_leaf_predicate->ptr()); + } return 0; } @@ -88,6 +106,10 @@ namespace optree { } self.m_agenda.clear(); Py_CLEAR(self.m_root.ptr()); + if (self.m_leaf_predicate) [[likely]] { + // Drop the owned leaf predicate reference to break cycles that pass through it. + Py_CLEAR(self.m_leaf_predicate->ptr()); + } return 0; } diff --git a/src/treespec/richcomparison.cpp b/src/treespec/richcomparison.cpp index 96b61798..42b60d2d 100644 --- a/src/treespec/richcomparison.cpp +++ b/src/treespec/richcomparison.cpp @@ -16,6 +16,7 @@ limitations under the License. */ #include // std::copy, std::reverse +#include // std::make_reverse_iterator #include // std::unordered_map #include // std::vector @@ -129,7 +130,13 @@ bool PyTreeSpec::IsPrefix(const PyTreeSpec &other, const bool &strict) const { EXPECT_EQ(reordered_other_offsets.front(), b->num_nodes, "PyTreeSpec traversal out of range."); - auto original_b = other.m_traversal.crbegin() + (b - other_traversal.crbegin()); + // Snapshot `b`'s subtree from the working copy before permuting its children + // below. The permutation copies overlapping child ranges, so the source must be + // a stable snapshot. It must be taken from the working copy (not the pristine + // `other.m_traversal`): a previously-processed ancestor dict may have relocated + // this subtree, so `b`'s offset no longer matches the pristine traversal. + const std::vector b_subtree(b.base() - b->num_nodes, b.base()); + const auto original_b = std::make_reverse_iterator(b_subtree.cend()); for (const auto &[i, j] : reordered_index_to_index) { std::copy(original_b + other_offsets[j + 1], original_b + other_offsets[j], diff --git a/src/treespec/serialization.cpp b/src/treespec/serialization.cpp index 04a36052..36df36c9 100644 --- a/src/treespec/serialization.cpp +++ b/src/treespec/serialization.cpp @@ -22,6 +22,7 @@ limitations under the License. #include // std::string #include // std::this_thread::get_id #include // std::unordered_set +#include // std::pair, std::move #include "optree/optree.h" @@ -137,9 +138,17 @@ std::string PyTreeSpec::ToStringImpl() const { case PyTreeKind::NamedTuple: { const py::object type = node.node_data; const auto fields = NamedTupleGetFields(type); - EXPECT_EQ(TupleGetSize(fields), - node.arity, - "Number of fields and entries does not match."); + // The field names are read from the (mutable) `_fields` attribute at repr time, so + // a caller may have changed them after the treespec was built. Report the mismatch + // as a `ValueError`, not an internal error, since the cause is external. + if (TupleGetSize(fields) != node.arity) [[unlikely]] { + std::ostringstream oss{}; + oss << "Number of fields (" << TupleGetSize(fields) << ") of namedtuple type " + << PyRepr(type) << " does not match the arity (" << node.arity + << ") of the treespec node. The `_fields` attribute may have been modified " + "after the treespec was created."; + throw py::value_error(oss.str()); + } const std::string kind = PyStr(EVALUATE_WITH_LOCK_HELD(py::getattr(type, "__name__"), type)); sstream << kind << "("; @@ -213,15 +222,24 @@ std::string PyTreeSpec::ToStringImpl() const { const py::object qualname = EVALUATE_WITH_LOCK_HELD(py::getattr(type, "__qualname__"), type); sstream << PyStr(qualname) << "("; - bool first = true; + ssize_t index = 0; auto child_it = agenda.cend() - node.arity; for (const py::handle &field : fields) { - if (!first) [[likely]] { + if (index > 0) [[likely]] { sstream << ", "; } - sstream << PyStr(field) << "=" << *child_it; + const std::string name = PyStr(field); + // An unnamed slot has no valid identifier, so render it angle-bracketed like + // CPython's `` rather than as the bare marker, which would read as an + // invalid keyword argument. The index disambiguates multiple unnamed slots. + if (name == PyStructSequenceUnnamedField()) [[unlikely]] { + sstream << ""; + } else [[likely]] { + sstream << name; + } + sstream << "=" << *child_it; ++child_it; - first = false; + ++index; } sstream << ")"; break; @@ -304,16 +322,31 @@ py::object PyTreeSpec::ToPicklable() const { const scoped_critical_section2 cs{ node.custom != nullptr ? py::handle{node.custom->type} : py::handle{}, node.node_data}; - TupleSetItem(node_states, - i++, - py::make_tuple(py::int_(static_cast(node.kind)), - py::int_(node.arity), - node.node_data ? node.node_data : py::none(), - node.node_entries ? node.node_entries : py::none(), - node.custom != nullptr ? node.custom->type : py::none(), - py::int_(node.num_leaves), - py::int_(node.num_nodes), - node.original_keys ? node.original_keys : py::none())); + + // Copy the node's mutable containers so the pickled state cannot alias (and, if the caller + // mutates it, corrupt) the immutable spec. Only dict-like keys are mutable and internal; a + // namedtuple/PyStructSequence type or custom metadata is left as-is. + py::object node_data = + node.node_data ? py::reinterpret_borrow(node.node_data) : py::none(); + if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::OrderedDict) [[unlikely]] { + node_data = node.node_data.attr("copy")(); + } else if (node.kind == PyTreeKind::DefaultDict) [[unlikely]] { + const auto metadata = py::reinterpret_borrow(node.node_data); + node_data = + py::make_tuple(TupleGetItem(metadata, 0), TupleGetItem(metadata, 1).attr("copy")()); + } + + TupleSetItem( + node_states, + i++, + py::make_tuple(py::int_(static_cast(node.kind)), + py::int_(node.arity), + std::move(node_data), + node.node_entries ? node.node_entries : py::none(), + node.custom != nullptr ? node.custom->type : py::none(), + py::int_(node.num_leaves), + py::int_(node.num_nodes), + node.original_keys ? node.original_keys.attr("copy")() : py::none())); } return py::make_tuple(node_states, py::bool_(m_none_is_leaf), py::str(m_namespace)); } @@ -321,9 +354,13 @@ py::object PyTreeSpec::ToPicklable() const { // NOLINTBEGIN[cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers] // NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ std::unique_ptr PyTreeSpec::FromPicklable(const py::object &picklable) { + const auto malformed = [](const std::string &reason) -> std::runtime_error { + return std::runtime_error("Malformed pickled PyTreeSpec: " + reason + "."); + }; + const auto state = thread_safe_cast(picklable); if (state.size() != 3) [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + throw malformed("the state is not a 3-tuple"); } bool none_is_leaf = false; std::string registry_namespace{}; @@ -332,56 +369,154 @@ py::object PyTreeSpec::ToPicklable() const { out->m_namespace = registry_namespace = thread_safe_cast(state[2]); const auto node_states = thread_safe_cast(state[0]); for (const auto &item : node_states) { - const auto t = thread_safe_cast(item); + const auto node_state = thread_safe_cast(item); + const auto node_state_size = node_state.size(); + if (node_state_size != 7 && node_state_size != 8) [[unlikely]] { + throw malformed("a node state is not a 7- or 8-tuple"); + } + Node &node = out->m_traversal.emplace_back(); - node.kind = static_cast(thread_safe_cast(t[0])); - node.arity = thread_safe_cast(t[1]); - if (t.size() != 7) [[unlikely]] { - if (t.size() == 8) [[likely]] { - if (t[7].is_none()) [[likely]] { - if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) - [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); - } - } else [[unlikely]] { - if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) - [[likely]] { - node.original_keys = DictFromKeys(t[7]); - } else [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); - } + const auto kind_value = thread_safe_cast(node_state[0]); + if (kind_value < 0 || kind_value >= static_cast(PyTreeKind::NumKinds)) + [[unlikely]] { + throw malformed("the node kind is out of range"); + } + node.kind = static_cast(kind_value); + node.arity = thread_safe_cast(node_state[1]); + if (node.arity < 0) [[unlikely]] { + throw malformed("the node arity is negative"); + } + if (node_state_size == 8) [[likely]] { + const auto &original_keys = node_state[7]; + if (original_keys.is_none()) [[likely]] { + if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) + [[unlikely]] { + throw malformed("a dict node is missing its original keys"); } } else [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) + [[likely]] { + node.original_keys = DictFromKeys(original_keys); + } else [[unlikely]] { + throw malformed("a non-dict node must not have original keys"); + } } } + + node.num_leaves = thread_safe_cast(node_state[5]); + node.num_nodes = thread_safe_cast(node_state[6]); + if (node.num_leaves < 0 || node.num_nodes < 1) [[unlikely]] { + throw malformed("a node has a negative or invalid size"); + } + + const auto &node_data = node_state[2]; + const auto &node_entries = node_state[3]; + const auto &custom_type = node_state[4]; switch (node.kind) { case PyTreeKind::Leaf: - case PyTreeKind::None: + case PyTreeKind::None: { + if (!node_data.is_none()) [[unlikely]] { + throw malformed("a leaf or none node must not have node data"); + } + // A leaf or none node is childless; a nonzero arity would let it absorb preceding + // subtrees (folding consistently) and silently drop leaves on unflatten. + if (node.arity != 0) [[unlikely]] { + throw malformed("a leaf or none node must have arity 0"); + } + // With `none_is_leaf`, None is flattened as a leaf, so a flattened tree never + // contains a None-kind node; a reconstructed one would later trip an InternalError + // in `FlattenUpTo` (`GetKind` never returns `None` under `NoneIsLeaf`). + if (node.kind == PyTreeKind::None && none_is_leaf) [[unlikely]] { + throw malformed("a none node cannot appear when none_is_leaf is set"); + } + break; + } + case PyTreeKind::Tuple: case PyTreeKind::List: { - if (!t[2].is_none()) [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + if (!node_data.is_none()) [[unlikely]] { + throw malformed("a tuple or list node must not have node data"); } break; } case PyTreeKind::Dict: case PyTreeKind::OrderedDict: { - node.node_data = thread_safe_cast(t[2]); + node.node_data = thread_safe_cast(node_data); + if (ListGetSize(node.node_data) != node.arity) [[unlikely]] { + throw malformed("the number of keys does not match the arity"); + } + // The keys must be hashable and distinct; a duplicate or unhashable key would + // collapse or fail when the dict is rebuilt, desyncing the keys from the children. + if (DistinctCount(node.node_data) != node.arity) [[unlikely]] { + throw malformed("the keys are not distinct"); + } + break; + } + + case PyTreeKind::NamedTuple: { + node.node_data = thread_safe_cast(node_data); + if (!IsNamedTupleClass(node.node_data)) [[unlikely]] { + throw malformed("the node data is not a namedtuple type"); + } + if (TupleGetSize(NamedTupleGetFields(node.node_data)) != node.arity) [[unlikely]] { + throw malformed("the number of fields does not match the arity"); + } break; } - case PyTreeKind::NamedTuple: case PyTreeKind::StructSequence: { - node.node_data = thread_safe_cast(t[2]); + node.node_data = thread_safe_cast(node_data); + if (!IsStructSequenceClass(node.node_data)) [[unlikely]] { + throw malformed("the node data is not a PyStructSequence type"); + } + if (TupleGetSize(StructSequenceGetFields(node.node_data)) != node.arity) + [[unlikely]] { + throw malformed("the number of fields does not match the arity"); + } + break; + } + + case PyTreeKind::DefaultDict: { + // A default dict stores its metadata as a 2-tuple `(default_factory, sorted_keys)`. + // `MakeNode` reads it with raw tuple/list accessors, so validate the shape here to + // avoid type-confusion on malformed input. + const auto metadata = thread_safe_cast(node_data); + if (metadata.size() != 2) [[unlikely]] { + throw malformed("the defaultdict metadata is not a 2-tuple"); + } + // `default_factory` is passed to `defaultdict(...)`, which requires None or + // callable. + if (!(metadata[0].is_none() || + static_cast(PyCallable_Check(metadata[0].ptr())))) [[unlikely]] { + throw malformed("the `default_factory` is not callable"); + } + const auto keys = thread_safe_cast(metadata[1]); + if (ListGetSize(keys) != node.arity) [[unlikely]] { + throw malformed("the number of keys does not match the arity"); + } + if (DistinctCount(keys) != node.arity) [[unlikely]] { + throw malformed("the keys are not distinct"); + } + node.node_data = metadata; + break; + } + + case PyTreeKind::Deque: { + // A deque's `maxlen` is None (unbounded) or a non-negative int bounding its length, + // so it must be at least the node's arity. + if (!node_data.is_none()) [[likely]] { + if (PyLong_Check(node_data.ptr()) == 0 || + thread_safe_cast(node_data) < node.arity) [[unlikely]] { + throw malformed("the deque maxlen is invalid"); + } + } + node.node_data = node_data; break; } - case PyTreeKind::DefaultDict: - case PyTreeKind::Deque: case PyTreeKind::Custom: { - node.node_data = t[2]; + node.node_data = node_data; break; } @@ -390,23 +525,23 @@ py::object PyTreeSpec::ToPicklable() const { INTERNAL_ERROR(); } if (node.kind == PyTreeKind::Custom) [[unlikely]] { // NOLINT - if (!t[3].is_none()) [[unlikely]] { - node.node_entries = thread_safe_cast(t[3]); + if (!node_entries.is_none()) [[unlikely]] { + node.node_entries = thread_safe_cast(node_entries); } - if (t[4].is_none()) [[unlikely]] { + if (custom_type.is_none()) [[unlikely]] { node.custom = nullptr; } else [[likely]] { if (none_is_leaf) [[unlikely]] { node.custom = - PyTreeTypeRegistry::Lookup(t[4], registry_namespace); + PyTreeTypeRegistry::Lookup(custom_type, registry_namespace); } else [[likely]] { node.custom = - PyTreeTypeRegistry::Lookup(t[4], registry_namespace); + PyTreeTypeRegistry::Lookup(custom_type, registry_namespace); } } if (node.custom == nullptr) [[unlikely]] { std::ostringstream oss{}; - oss << "Unknown custom type in pickled PyTreeSpec: " << PyRepr(t[4]); + oss << "Unknown custom type in pickled PyTreeSpec: " << PyRepr(custom_type); if (!registry_namespace.empty()) [[likely]] { oss << " in namespace " << PyRepr(registry_namespace); } else [[unlikely]] { @@ -415,21 +550,69 @@ py::object PyTreeSpec::ToPicklable() const { oss << "."; throw std::runtime_error(oss.str()); } - } else if (!t[3].is_none() || !t[4].is_none()) [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + } else if (!node_entries.is_none() || !custom_type.is_none()) [[unlikely]] { + throw malformed("a non-custom node must not have node entries or a custom type"); } - if (node.original_keys && DictGetSize(node.original_keys) != node.arity) [[unlikely]] { - throw std::runtime_error("Number of keys does not match arity in pickled PyTreeSpec."); + if (node.original_keys) [[unlikely]] { + if (DictGetSize(node.original_keys) != node.arity) [[unlikely]] { + throw malformed("the number of original keys does not match the arity"); + } + // `original_keys` records the insertion order of the same keys stored (sorted) in + // node_data; its key set must match, or unflatten would map children onto keys the dict + // never had. + const auto keys = (node.kind == PyTreeKind::DefaultDict + ? TupleGetItemAs(node.node_data, 1) + : py::reinterpret_borrow(node.node_data)); + for (const py::handle &key : keys) { + const int contains = PyDict_Contains(node.original_keys.ptr(), key.ptr()); + if (contains < 0) [[unlikely]] { + throw py::error_already_set(); + } + if (contains == 0) [[unlikely]] { + throw malformed("the keys do not match the original keys"); + } + } } if (node.node_entries && !node.node_entries.is_none() && TupleGetSize(node.node_entries) != node.arity) [[unlikely]] { - throw std::runtime_error( - "Number of node entries does not match arity in pickled PyTreeSpec."); + throw malformed("the number of node entries does not match the arity"); } + } - node.num_leaves = thread_safe_cast(t[5]); - node.num_nodes = thread_safe_cast(t[6]); + // Validate that the reconstructed traversal is structurally consistent. + // `PYTREESPEC_SANITY_CHECK` only checks the final node, so a malformed pickle could otherwise + // smuggle in inconsistent arity / num_nodes / num_leaves that cause out-of-bounds access when + // the spec is later used. Walk the post-order traversal, folding each node's children off a + // stack of subtree sizes. + { + auto subtree_sizes = + reserved_vector>( + out->m_traversal.size()); + for (const Node &node : out->m_traversal) { + if (static_cast(subtree_sizes.size()) < node.arity) [[unlikely]] { + throw malformed("a node has more children than available subtrees"); + } + ssize_t children_num_nodes = 0; + ssize_t children_num_leaves = 0; + for (ssize_t i = 0; i < node.arity; ++i) { + children_num_nodes += subtree_sizes.back().first; + children_num_leaves += subtree_sizes.back().second; + subtree_sizes.pop_back(); + } + const ssize_t expected_num_nodes = children_num_nodes + 1; + const ssize_t expected_num_leaves = + (node.kind == PyTreeKind::Leaf ? ssize_t{1} : children_num_leaves); + if (node.num_nodes != expected_num_nodes || node.num_leaves != expected_num_leaves) + [[unlikely]] { + throw malformed("a node's size is inconsistent with its children"); + } + subtree_sizes.emplace_back(node.num_nodes, node.num_leaves); + } + if (subtree_sizes.size() != 1) [[unlikely]] { + throw malformed("the traversal does not yield a single tree"); + } } + out->m_traversal.shrink_to_fit(); PYTREESPEC_SANITY_CHECK(*out); return out; diff --git a/src/treespec/treespec.cpp b/src/treespec/treespec.cpp index 44305dd0..0d0f827e 100644 --- a/src/treespec/treespec.cpp +++ b/src/treespec/treespec.cpp @@ -185,13 +185,36 @@ namespace optree { } } +std::optional PyTreeSpec::FindReregisteredCustomType( + const std::string &target_namespace) const { + for (const Node &node : m_traversal) { + if (node.kind == PyTreeKind::Custom) [[unlikely]] { + const auto registration = + (m_none_is_leaf + ? PyTreeTypeRegistry::Lookup(node.custom->type, target_namespace) + : PyTreeTypeRegistry::Lookup(node.custom->type, + target_namespace)); + if (registration != node.custom) [[unlikely]] { + return node.custom->type; + } + } + } + return {}; +} + // NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ std::tuple PyTreeSpec::BroadcastToCommonSuffixImpl( std::vector &nodes, const std::vector &traversal, const ssize_t &pos, const std::vector &other_traversal, - const ssize_t &other_pos) { + const ssize_t &other_pos, + const ssize_t &depth) { + if (depth > MAX_RECURSION_DEPTH) [[unlikely]] { + PyErr_SetString(PyExc_RecursionError, + "Maximum recursion depth exceeded during broadcasting the treespecs."); + throw py::error_already_set(); + } const Node &root = traversal.at(pos); const Node &other_root = other_traversal.at(other_pos); EXPECT_GE(pos + 1, @@ -276,15 +299,19 @@ namespace optree { const auto expected_keys = (root.kind != PyTreeKind::DefaultDict ? py::reinterpret_borrow(root.node_data) : TupleGetItemAs(root.node_data, 1)); - auto other_keys = (other_root.kind != PyTreeKind::DefaultDict - ? py::reinterpret_borrow(other_root.node_data) - : TupleGetItemAs(other_root.node_data, 1)); + const auto other_keys = (other_root.kind != PyTreeKind::DefaultDict + ? py::reinterpret_borrow(other_root.node_data) + : TupleGetItemAs(other_root.node_data, 1)); const py::dict dict{}; for (ssize_t i = 0; i < other_root.arity; ++i) { DictSetItem(dict, ListGetItem(other_keys, i), py::int_(i)); } if (!DictKeysEqual(expected_keys, dict)) [[unlikely]] { - TotalOrderSort(other_keys); + // Build the message from a sorted COPY of the keys. `other_keys` is a borrow of the + // argument spec's live `node_data`; sorting it in place would permute the keys + // while the child subtrees stay put, silently corrupting a spec the caller still + // holds. + const py::list sorted_other_keys = SortedDictKeys(dict); const auto [missing_keys, extra_keys] = DictKeysDifference(expected_keys, dict); std::ostringstream key_difference_sstream{}; if (ListGetSize(missing_keys) != 0) [[likely]] { @@ -295,7 +322,7 @@ namespace optree { } std::ostringstream oss{}; oss << "dictionary key mismatch; expected key(s): " << PyRepr(expected_keys) - << ", got key(s): " << PyRepr(other_keys) << key_difference_sstream.str() + << ", got key(s): " << PyRepr(sorted_other_keys) << key_difference_sstream.str() << "."; throw py::value_error(oss.str()); } @@ -314,7 +341,12 @@ namespace optree { other_cur = other_curs[py::cast(DictGetItem(dict, key))]; const auto [num_nodes, other_num_nodes, new_num_nodes, new_num_leaves] = // NOLINTNEXTLINE[misc-no-recursion] - BroadcastToCommonSuffixImpl(nodes, traversal, cur, other_traversal, other_cur); + BroadcastToCommonSuffixImpl(nodes, + traversal, + cur, + other_traversal, + other_cur, + depth + 1); cur -= num_nodes; nodes[start_num_nodes].num_nodes += new_num_nodes; nodes[start_num_nodes].num_leaves += new_num_leaves; @@ -393,7 +425,12 @@ namespace optree { for (ssize_t i = root.arity - 1; i >= 0; --i) { const auto [num_nodes, other_num_nodes, new_num_nodes, new_num_leaves] = // NOLINTNEXTLINE[misc-no-recursion] - BroadcastToCommonSuffixImpl(nodes, traversal, cur, other_traversal, other_cur); + BroadcastToCommonSuffixImpl(nodes, + traversal, + cur, + other_traversal, + other_cur, + depth + 1); cur -= num_nodes; other_cur -= other_num_nodes; nodes[start_num_nodes].num_nodes += new_num_nodes; @@ -405,6 +442,7 @@ namespace optree { nodes[start_num_nodes].num_leaves}; } +// NOLINTNEXTLINE[readability-function-cognitive-complexity] std::unique_ptr PyTreeSpec::BroadcastToCommonSuffix(const PyTreeSpec &other) const { PYTREESPEC_SANITY_CHECK(*this); PYTREESPEC_SANITY_CHECK(other); @@ -420,12 +458,36 @@ std::unique_ptr PyTreeSpec::BroadcastToCommonSuffix(const PyTreeSpec throw py::value_error(oss.str()); } + const std::string &target_namespace = m_namespace.empty() ? other.m_namespace : m_namespace; auto treespec = std::make_unique(); treespec->m_none_is_leaf = m_none_is_leaf; - if (other.m_namespace.empty()) [[likely]] { - treespec->m_namespace = m_namespace; + treespec->m_namespace = target_namespace; + + if (!target_namespace.empty()) [[likely]] { + // The compatibility check above rejects two distinct non-empty namespaces, so the adopted + // namespace equals one side's namespace and at most the other (empty) side differs from it. + // Re-check only that side: adopting a namespace under which its custom nodes resolve to a + // different registration would silently rebind them (the result keeps the original ones). + std::optional reregistered_type{}; + if (target_namespace != m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, other.m_namespace, "Namespace mismatch."); + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + reregistered_type = FindReregisteredCustomType(target_namespace); + } else if (target_namespace != other.m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, m_namespace, "Namespace mismatch."); + EXPECT_TRUE(other.m_namespace.empty(), "Namespace mismatch."); + reregistered_type = other.FindReregisteredCustomType(target_namespace); + } + if (reregistered_type) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be merged: custom PyTree type " << PyRepr(*reregistered_type) + << " resolves to a different registration in namespace " << PyRepr(target_namespace) + << "."; + throw py::value_error(oss.str()); + } } else [[unlikely]] { - treespec->m_namespace = other.m_namespace; + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + EXPECT_TRUE(other.m_namespace.empty(), "Namespace mismatch."); } const ssize_t num_nodes = GetNumNodes(); @@ -436,7 +498,8 @@ std::unique_ptr PyTreeSpec::BroadcastToCommonSuffix(const PyTreeSpec m_traversal, num_nodes - 1, other.m_traversal, - other_num_nodes - 1); + other_num_nodes - 1, + 0); std::reverse(treespec->m_traversal.begin(), treespec->m_traversal.end()); EXPECT_EQ(num_nodes_walked, num_nodes, @@ -573,11 +636,29 @@ std::unique_ptr PyTreeSpec::Transform(const std::optionalm_none_is_leaf = m_none_is_leaf; treespec->m_namespace = common_registry_namespace; + + // Reject a transform whose unified namespace would rebind a custom node to a different + // registration than the one it holds (the result keeps each node's original registration; e.g. + // a globally-resolved custom node from the input under a non-empty unified namespace). Only + // relevant for a non-empty namespace: an empty one resolves every custom node globally. + if (!common_registry_namespace.empty()) [[unlikely]] { + if (const auto &reregistered_type = + treespec->FindReregisteredCustomType(common_registry_namespace)) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be transformed: custom PyTree type " + << PyRepr(*reregistered_type) + << " resolves to a different registration in namespace " + << PyRepr(common_registry_namespace) << "."; + throw py::value_error(oss.str()); + } + } + treespec->m_traversal.shrink_to_fit(); PYTREESPEC_SANITY_CHECK(*treespec); return treespec; } +// NOLINTNEXTLINE[readability-function-cognitive-complexity] std::unique_ptr PyTreeSpec::Compose(const PyTreeSpec &inner) const { PYTREESPEC_SANITY_CHECK(*this); PYTREESPEC_SANITY_CHECK(inner); @@ -593,12 +674,36 @@ std::unique_ptr PyTreeSpec::Compose(const PyTreeSpec &inner) const { throw py::value_error(oss.str()); } + const std::string &target_namespace = m_namespace.empty() ? inner.m_namespace : m_namespace; auto treespec = std::make_unique(); treespec->m_none_is_leaf = m_none_is_leaf; - if (inner.m_namespace.empty()) [[likely]] { - treespec->m_namespace = m_namespace; + treespec->m_namespace = target_namespace; + + if (!target_namespace.empty()) [[likely]] { + // The compatibility check above rejects two distinct non-empty namespaces, so the adopted + // namespace equals one side's namespace and at most the other (empty) side differs from it. + // Re-check only that side: adopting a namespace under which its custom nodes resolve to a + // different registration would silently rebind them (the result keeps the original ones). + std::optional reregistered_type{}; + if (target_namespace != m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, inner.m_namespace, "Namespace mismatch."); + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + reregistered_type = FindReregisteredCustomType(target_namespace); + } else if (target_namespace != inner.m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, m_namespace, "Namespace mismatch."); + EXPECT_TRUE(inner.m_namespace.empty(), "Namespace mismatch."); + reregistered_type = inner.FindReregisteredCustomType(target_namespace); + } + if (reregistered_type) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be merged: custom PyTree type " << PyRepr(*reregistered_type) + << " resolves to a different registration in namespace " << PyRepr(target_namespace) + << "."; + throw py::value_error(oss.str()); + } } else [[unlikely]] { - treespec->m_namespace = inner.m_namespace; + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + EXPECT_TRUE(inner.m_namespace.empty(), "Namespace mismatch."); } const ssize_t num_outer_leaves = GetNumLeaves(); @@ -638,6 +743,11 @@ ssize_t PyTreeSpec::PathsImpl(PathVector &paths, // NOLINT[misc-no-recursion] const ssize_t &depth) const { const Node &root = m_traversal.at(pos); EXPECT_GE(pos + 1, root.num_nodes, "PyTreeSpec::Paths() walked off start of array."); + if (depth > MAX_RECURSION_DEPTH) [[unlikely]] { + PyErr_SetString(PyExc_RecursionError, + "Maximum recursion depth exceeded during walking the tree."); + throw py::error_already_set(); + } ssize_t cur = pos - 1; // NOLINTNEXTLINE[misc-no-recursion] @@ -736,6 +846,11 @@ ssize_t PyTreeSpec::AccessorsImpl(Span &accessors, // NOLINT[misc-no-recursion] const Node &root = m_traversal.at(pos); EXPECT_GE(pos + 1, root.num_nodes, "PyTreeSpec::TypedPaths() walked off start of array."); + if (depth > MAX_RECURSION_DEPTH) [[unlikely]] { + PyErr_SetString(PyExc_RecursionError, + "Maximum recursion depth exceeded during walking the tree."); + throw py::error_already_set(); + } ssize_t cur = pos - 1; const py::object node_type = GetType(root); diff --git a/tests/concurrent/test_subinterpreters.py b/tests/concurrent/test_subinterpreters.py index aec86e9a..ad50ee63 100644 --- a/tests/concurrent/test_subinterpreters.py +++ b/tests/concurrent/test_subinterpreters.py @@ -343,3 +343,125 @@ def check_import(): output='', rerun=NUM_FLAKY_RERUNS, ) + + +def test_type_cache_cleanup_across_subinterpreters(): + # Exercise the process-global type caches across subinterpreter churn. The caches key on the + # type's memory ADDRESS, which the allocator may recycle after a type is freed. Each interpreter + # flattens a RANDOMLY chosen type and reads its repr, firing both the classification and the + # field-name cache; distinct field names make a stale entry on a recycled address observable. + # The subprocess makes finalization real. + check_script_in_subprocess( + f""" + import contextlib + import textwrap + from concurrent import interpreters + + resolve = textwrap.dedent( + ''' + import keyword + import os + import random + import string + import time + from collections import namedtuple + + import optree + + # Distinct PyStructSequence types, each with a distinct first sequence field. + cases = [ + (os.stat_result, 'st_mode'), + (time.struct_time, 'tm_year'), + (os.times_result, 'user'), + ] + case = random.choice([*cases, None]) + if case is None: + # Exercise the classification cache with a namedtuple type. + field_names = [] + while not field_names: + field_names = [ + ''.join(random.choices(string.ascii_lowercase, k=random.randint(1, 8))) + for _ in range(random.randint(2, 5)) + ] + field_names = [n for n in field_names if not keyword.iskeyword(n)] # avoid reserved names + field_names = list(dict.fromkeys(field_names)) # deduplicate + NamedTupleType = namedtuple('NamedTupleType', field_names) + leaves, treespec = optree.tree_flatten(NamedTupleType(*range(len(field_names)))) + assert len(leaves) == len(field_names), (leaves, treespec) + assert (field_names[0] + '=*') in repr(treespec), repr(treespec) + else: + # Exercise the classification cache and the field-name cache with a PyStructSequence type. + PyStructSequenceType, first_field = case + obj = PyStructSequenceType(range(PyStructSequenceType.n_sequence_fields)) + assert (first_field + '=*') in repr(optree.tree_structure(obj)), first_field + ''' + ).strip() + + exec(resolve) # the main interpreter resolves a random type + for _ in range({NUM_FUTURES}): + with contextlib.closing(interpreters.create()) as subinterpreter: + subinterpreter.exec(resolve) + exec(resolve) # the main interpreter must still resolve correctly after the churn + """, + output='', + rerun=NUM_FLAKY_RERUNS, + ) + + +def test_registry_lookup_no_gil_lock_order_deadlock(): + # Regression: `Lookup`/`Register`/`Unregister` called `GetSingleton()` while holding `sm_mutex`. + # Under `per_interpreter_gil` that call releases the GIL once any subinterpreter has imported + # `optree._C`, so a flatten thread drops the GIL holding the read lock while a registration + # thread holds the GIL waiting on the write lock: a lock-order inversion that hangs the process. + # The watchdog turns the hang into a non-zero exit; the fix acquires the singleton before + # locking. + check_script_in_subprocess( + f""" + import contextlib + import faulthandler + import threading + import time + from concurrent import interpreters + + import optree + + # Latch pybind11's `has_seen_non_main_interpreter` so `GetSingleton()` releases the GIL on + # every subsequent call (the single-interpreter fast path is disabled process-wide). + with contextlib.closing(interpreters.create()) as subinterpreter: + subinterpreter.exec('import optree._C') + + faulthandler.dump_traceback_later(20, exit=True) # a deadlock becomes a non-zero exit + + stop = threading.Event() + tree = {{'a': 1, 'b': [2, 3], 'c': (4, 5)}} + + def flatten_worker(): + while not stop.is_set(): + optree.tree_flatten(tree) + + def register_worker(index): + cls = type(f'DeadlockNode_{{index}}', (), {{}}) + while not stop.is_set(): + optree.register_pytree_node( + cls, + lambda x: ((), None), + lambda metadata, children: cls(), + namespace='deadlock' + ) + optree.unregister_pytree_node(cls, namespace='deadlock') + + workers = [threading.Thread(target=flatten_worker) for _ in range({NUM_WORKERS})] + workers += [threading.Thread(target=register_worker, args=(index,)) for index in range(2)] + for worker in workers: + worker.start() + time.sleep(0.5) # let readers and writers contend on `sm_mutex` + stop.set() + for worker in workers: + worker.join(timeout=5) + assert not worker.is_alive(), 'worker thread did not terminate (deadlock)' + + faulthandler.cancel_dump_traceback_later() + """, + output=None, + rerun=NUM_FLAKY_RERUNS, + ) diff --git a/tests/helpers.py b/tests/helpers.py index 5440bb11..dde27fa0 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -75,6 +75,14 @@ reason='Py_GIL_DISABLED is set', ) +# Free-threaded builds before 3.14 hold deferred references to type objects in per-thread caches, so +# `gc_collect()` cannot force a heap type to be reclaimed there. +HAS_DEFERRED_TYPE_REFS = Py_GIL_DISABLED and sys.version_info < (3, 14) +skipif_deferred_type_refs = pytest.mark.skipif( + HAS_DEFERRED_TYPE_REFS, + reason='free-threaded builds before 3.14 keep deferred references to type objects', +) + PYPY = platform.python_implementation() == 'PyPy' skipif_pypy = pytest.mark.skipif( PYPY, diff --git a/tests/integrations/test_attrs.py b/tests/integrations/test_attrs.py index d8b7ac5c..44ea115b 100644 --- a/tests/integrations/test_attrs.py +++ b/tests/integrations/test_attrs.py @@ -97,6 +97,33 @@ def __attrs_post_init__(self): assert foo == optree.tree_unflatten(treespec, leaves) +def test_define_reconstruction_reapplies_field_converters(): + # Characterization: the generated unflatten rebuilds via `cls(**fields)`, which re-runs attrs + # field converters. Since `flatten` reads the already-converted value, an idempotent converter + # round-trips exactly while a non-idempotent one does not. A class needing exact reconstruction + # should register explicitly with custom flatten/unflatten (see `register_node`). + + # Idempotent converter: `int(int(x)) == int(x)`, so the round-trip is exact. + @optree.integrations.attrs.define(namespace='test-attrs-converter-idempotent') + class Cast: + x: int = optree.integrations.attrs.field(converter=int) + + cast = Cast(5.0) # __init__ converter: x = int(5.0) = 5 + leaves, treespec = optree.tree_flatten(cast, namespace='test-attrs-converter-idempotent') + assert leaves == [5] + assert optree.tree_unflatten(treespec, leaves).x == 5 + + # Non-idempotent converter: unflatten re-applies it, so the stored `6` becomes `7`. + @optree.integrations.attrs.define(namespace='test-attrs-converter-shift') + class Shift: + x: int = optree.integrations.attrs.field(converter=lambda v: v + 1) + + shift = Shift(5) # __init__ converter: x = 5 + 1 = 6 + leaves, treespec = optree.tree_flatten(shift, namespace='test-attrs-converter-shift') + assert leaves == [6] + assert optree.tree_unflatten(treespec, leaves).x == 7 # re-applied: 6 + 1 = 7, not 6 + + def test_define_frozen(): @optree.integrations.attrs.frozen(namespace='test-attrs-frozen') class FrozenPoint: @@ -143,7 +170,12 @@ def foo(): def test_define_with_duplicate_registrations(): with pytest.raises( TypeError, - match=r'Cannot register .* as a pytree node more than once\.', + match=( + r'Cannot register .* as a pytree node more than once with ' + r'`optree\.integrations\.attrs\.register_node\(\)`\. ' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)` ' + r'with explicit flatten/unflatten functions' + ), ): @optree.integrations.attrs.define(namespace='error') @@ -292,10 +324,52 @@ class Double: with pytest.raises( TypeError, - match=r'Cannot register .* as a pytree node more than once\.', + match=( + r'Cannot register .* as a pytree node more than once with ' + r'`optree\.integrations\.attrs\.register_node\(\)`\. ' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)` ' + r'with explicit flatten/unflatten functions' + ), ): optree.integrations.attrs.register_node(Double, namespace='test-attrs-double-2') + # As the error suggests, the class can still be registered in another namespace via the generic + # API with explicit flatten/unflatten functions. + optree.register_pytree_node( + Double, + lambda d: ((d.x,), None, None), + lambda _, children: Double(*children), + namespace='test-attrs-double-2', + ) + optree.unregister_pytree_node(Double, namespace='test-attrs-double-2') + + +def test_register_node_failure_does_not_leak_fields_guard(): + @attrs.define + class Leak: + x: int + + namespace = 'test-attrs-register-leak' + # Occupy the (class, namespace) slot directly so the attrs `register_node()`'s internal + # `register_pytree_node()` call fails, after `register_node()` would set its `_FIELDS` guard. + optree.register_pytree_node( + Leak, + lambda leak: ((leak.x,), None, None), + lambda _, children: Leak(*children), + namespace=namespace, + ) + try: + with pytest.raises(ValueError, match='already registered'): + optree.integrations.attrs.register_node(Leak, namespace=namespace) + finally: + optree.unregister_pytree_node(Leak, namespace=namespace) + + # A failed registration must not leave the `_FIELDS` guard behind, or the class becomes + # impossible to register ever again: every retry would raise "... more than once". + # Re-registration after clearing the conflicting entry must succeed. + optree.integrations.attrs.register_node(Leak, namespace=namespace) + optree.unregister_pytree_node(Leak, namespace=namespace) + def test_register_init_false_class_warns(): @attrs.define(init=False) @@ -356,9 +430,29 @@ class EntryTest: assert 'AttrsEntry' in repr(entry_str) assert "'x'" in repr(entry_str) - entry_int = optree.integrations.attrs.AttrsEntry(1, EntryTest, optree.PyTreeKind.CUSTOM) - assert entry_int.field == 'y' - assert entry_int.name == 'y' + +def test_attrs_entry_integer_indexes_children(): + # An integer entry indexes the tree CHILDREN (the fields that are BOTH `pytree_node=True` and + # `init`), not all init fields and not all `pytree_node` fields. A non-child field interleaved + # between children (a `pytree_node=False` field OR a non-`init` field) must not shift the mapping. + @attrs.define + class Foo: + a: int + b: int = optree.integrations.attrs.field(default=0, pytree_node=False) # not a child + d: int = attrs.field(init=False, default=0) # non-init -> not a tree child + c: int = 0 + + foo = Foo(1, 2, 3) # a=1, b=2, c=3 (d defaults to 0, not an init parameter) + assert tuple(a.name for a in attrs.fields(Foo)) == ('a', 'b', 'd', 'c') + + entry_int = optree.integrations.attrs.AttrsEntry(1, Foo, optree.PyTreeKind.CUSTOM) + assert entry_int.init_fields == ('a', 'b', 'c') # `d` is not an init field + assert entry_int.children_fields == ('a', 'c') # `b` is metadata and `d` is non-init + # The 2nd child is `c` (not the metadata field `b` nor the non-init field `d`). + assert entry_int.field == 'c' + assert entry_int.name == 'c' + assert entry_int.codify('x') == 'x.c' + assert entry_int(foo) == foo.c == 3 def test_accessor_codify(): diff --git a/tests/test_dataclasses.py b/tests/test_dataclasses.py index 8a4bae98..ca8c8595 100644 --- a/tests/test_dataclasses.py +++ b/tests/test_dataclasses.py @@ -932,6 +932,25 @@ def __post_init__(self): assert obj == optree.tree_unflatten(treespec, leaves) +def test_dataclass_post_init_rewriting_a_field_does_not_round_trip(): + # Characterization: the generated unflatten rebuilds via `cls(**fields)`, re-running `__init__` + # and `__post_init__`. A `__post_init__` that only derives non-field state round-trips (see + # `test_register_existing_class_with_metadata` above), but one that rewrites a stored field does + # not: flatten reads the rewritten value and unflatten rewrites it again. Such a class should + # register explicitly with custom flatten/unflatten (see `register_node`). + @optree.dataclasses.dataclass(namespace='test-dc-post-init-rewrite') + class Shift: + x: int + + def __post_init__(self): + self.x = self.x + 1 + + obj = Shift(5) # __post_init__: x = 5 + 1 = 6 + leaves, treespec = optree.tree_flatten(obj, namespace='test-dc-post-init-rewrite') + assert leaves == [6] + assert optree.tree_unflatten(treespec, leaves).x == 7 # re-applied: 6 + 1 = 7, not 6 + + def test_register_non_class(): with pytest.raises(TypeError, match='Expected a class'): optree.dataclasses.register_node(42, namespace='error') @@ -954,10 +973,94 @@ class Double: with pytest.raises( TypeError, - match=r'Cannot register .* as a pytree node more than once\.', + match=( + r'Cannot register .* as a pytree node more than once with ' + r'`optree\.dataclasses\.register_node\(\)`\. ' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)` ' + r'with explicit flatten/unflatten functions' + ), ): optree.dataclasses.register_node(Double, namespace='test-dc-double-2') + # As the error suggests, the class can still be registered in another namespace via the generic + # API with explicit flatten/unflatten functions. + optree.register_pytree_node( + Double, + lambda d: ((d.x,), None, None), + lambda _, children: Double(*children), + namespace='test-dc-double-2', + ) + optree.unregister_pytree_node(Double, namespace='test-dc-double-2') + + +def test_register_dataclass_with_initvar_rejected(): + # A dataclass with an `InitVar` cannot round-trip via the auto-generated flatten/unflatten: + # `InitVar`s are excluded from `dataclasses.fields()` (so they are neither flattened nor stored) + # yet remain required `__init__` parameters that `cls(**kwargs)` cannot restore. `register_node()` + # rejects it and points to the generic API with explicit flatten/unflatten functions. + @dataclasses.dataclass + class WithInitVar: + x: float + scale: dataclasses.InitVar[float] + scaled: float = optree.dataclasses.field(init=False, pytree_node=False, default=0.0) + + def __post_init__(self, scale): + self.scaled = self.x * scale + + with pytest.raises( + TypeError, + match=( + r'has `InitVar` field\(s\) .*cannot round-trip.*' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)`' + ), + ): + optree.dataclasses.register_node(WithInitVar, namespace='test-dc-initvar') + + # The generic API with explicit flatten/unflatten works: keep the derived value as metadata and + # rebuild the instance without re-running `__init__`/`__post_init__`. + def flatten(obj): + return (obj.x,), obj.scaled + + def unflatten(scaled, children): + rebuilt = WithInitVar.__new__(WithInitVar) + rebuilt.x = children[0] + rebuilt.scaled = scaled + return rebuilt + + optree.register_pytree_node(WithInitVar, flatten, unflatten, namespace='test-dc-initvar') + obj = WithInitVar(3.0, 10.0) + leaves, treespec = optree.tree_flatten(obj, namespace='test-dc-initvar') + assert leaves == [3.0] + assert optree.tree_unflatten(treespec, leaves) == obj + optree.unregister_pytree_node(WithInitVar, namespace='test-dc-initvar') + + +def test_register_node_failure_does_not_leak_fields_guard(): + @dataclasses.dataclass + class Leak: + x: int + + namespace = 'test-dc-register-leak' + # Occupy the (class, namespace) slot directly so the dataclass `register_node()`'s internal + # `register_pytree_node()` call fails, after `register_node()` would set its `_FIELDS` guard. + optree.register_pytree_node( + Leak, + lambda leak: ((leak.x,), None, None), + lambda _, children: Leak(*children), + namespace=namespace, + ) + try: + with pytest.raises(ValueError, match='already registered'): + optree.dataclasses.register_node(Leak, namespace=namespace) + finally: + optree.unregister_pytree_node(Leak, namespace=namespace) + + # A failed registration must not leave the `_FIELDS` guard behind, or the class becomes + # impossible to register ever again: every retry would raise "... more than once". + # Re-registration after clearing the conflicting entry must succeed. + optree.dataclasses.register_node(Leak, namespace=namespace) + optree.unregister_pytree_node(Leak, namespace=namespace) + def test_register_init_false_class_warns(): @dataclasses.dataclass(init=False) @@ -1131,6 +1234,28 @@ class EntryTest: assert 'DataclassEntry' in repr(entry_str) assert "'x'" in repr(entry_str) - entry_int = optree.DataclassEntry(1, EntryTest, optree.PyTreeKind.CUSTOM) - assert entry_int.field == 'y' - assert entry_int.name == 'y' + +def test_dataclass_entry_integer_indexes_children(): + # An integer entry indexes the tree CHILDREN (the fields that are BOTH `pytree_node=True` and + # `init`), not all init fields and not all `pytree_node` fields. A non-child field interleaved + # between children (a `pytree_node=False` field OR a non-`init` field) must not shift the mapping. + @dataclasses.dataclass + class Foo: + a: int + # metadata (init, not a child) + b: int = optree.dataclasses.field(default=0, pytree_node=False) + # non-init -> not a tree child + d: int = dataclasses.field(init=False, default=0) + c: int = 0 + + foo = Foo(1, 2, 3) # a=1, b=2, c=3 (d defaults to 0, not an init parameter) + assert tuple(f.name for f in dataclasses.fields(Foo)) == ('a', 'b', 'd', 'c') + + entry_int = optree.DataclassEntry(1, Foo, optree.PyTreeKind.CUSTOM) + assert entry_int.init_fields == ('a', 'b', 'c') # `d` is not an init field + assert entry_int.children_fields == ('a', 'c') # `b` is metadata and `d` is non-init + # The 2nd child is `c` (not the metadata field `b` nor the non-init field `d`). + assert entry_int.field == 'c' + assert entry_int.name == 'c' + assert entry_int.codify('x') == 'x.c' + assert entry_int(foo) == foo.c == 3 diff --git a/tests/test_ops.py b/tests/test_ops.py index 04fd8e89..bf0525c7 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -3120,6 +3120,99 @@ def test_tree_broadcast_common(): ) +def test_tree_broadcast_common_does_not_apply_is_leaf_to_internal_sentinel(): + # `tree_broadcast_common` fills an internal `common_suffix_tree` with a private `object()` + # sentinel, then replicates each leaf across the matching subtree. The user's `is_leaf` must run + # only on values from the input trees, never the sentinel: a value-dependent predicate would + # otherwise mis-structure the placeholder subtree or crash while inspecting it. + + # Collapses a list unless every element is an int. When the predicate keeps the other's list + # (all ints), the fix must not re-apply the predicate to the sentinel subtree, which would + # collapse `[, ]` and under-replicate. When the predicate collapses the real other + # tree to a leaf (a list holding a tuple), there is simply nothing to broadcast and the inputs + # pass through unchanged. + def collapse_non_int_lists(x): + return isinstance(x, list) and not all(isinstance(e, int) for e in x) + + assert optree.tree_broadcast_common( + 1, + [2, 3], + is_leaf=collapse_non_int_lists, + ) == ([1, 1], [2, 3]) + assert optree.tree_broadcast_common( + 1, + [2, (3, 4)], + is_leaf=collapse_non_int_lists, + ) == (1, [2, (3, 4)]) + assert optree.tree_broadcast_common( + {'a': 1}, + {'a': [2, 3]}, + is_leaf=collapse_non_int_lists, + ) == ({'a': [1, 1]}, {'a': [2, 3]}) + assert optree.tree_broadcast_common( + {'a': 1}, + {'a': [2, (3, 4)]}, + is_leaf=collapse_non_int_lists, + ) == ({'a': 1}, {'a': [2, (3, 4)]}) + + # Inspects the value; must not crash on the sentinel (`object() < 0` raises `TypeError`), + # regardless of the container (list, tuple, nested) the sentinel subtree takes. + def negative_scalar(x): + return not isinstance(x, (list, tuple, dict)) and x < 0 + + assert optree.tree_broadcast_common( + 1, + [2, 3], + is_leaf=negative_scalar, + ) == ([1, 1], [2, 3]) + assert optree.tree_broadcast_common( + [1], + [(2, 3)], + is_leaf=negative_scalar, + ) == ([(1, 1)], [(2, 3)]) + assert optree.tree_broadcast_common( + [1, 1], + [[2], [3, 4]], + is_leaf=negative_scalar, + ) == ([[1], [1, 1]], [[2], [3, 4]]) + # `broadcast_common` delegates to `tree_broadcast_common`, so it must not crash either. + assert optree.broadcast_common(1, [2, 3], is_leaf=negative_scalar) == ([1, 1], [2, 3]) + + +def test_tree_broadcast_common_value_dependent_is_leaf_is_lossy(): + # Characterization, separate from the R18 sentinel fix: a `is_leaf` that classifies by leaf + # *value* rather than type/structure can be lossy under broadcasting. `is_shape_like` treats a + # tuple of ints as one leaf, so broadcasting the scalar `5` into a two-slot yields `(5, 5)`, + # which the predicate re-reads as a single shape. Such a result still shares the literal + # common-suffix structure, but re-flattens to a different leaf count under the predicate. The + # cases below show it is lossy for a number yet faithful for a shape or a list slot. This pins + # the known limitation; prefer type/structure predicates for broadcasting. + def is_shape_like(x): + return type(x) is tuple and all(isinstance(v, int) for v in x) + + # A tuple slot is lossy: `(5, 5)` re-reads as a shape leaf. + assert optree.tree_broadcast_common( + [5], + [(1, 'x')], + is_leaf=is_shape_like, + ) == ([(5, 5)], [(1, 'x')]) + # A list slot is not: `[5, 5]` is not a shape tuple, so it stays consistent. + assert optree.tree_broadcast_common( + [5], + [[1, 2]], + is_leaf=is_shape_like, + ) == ([[5, 5]], [[1, 2]]) + + # A mixed tree shows both outcomes at once. The shape `(5, 5)` broadcast into a two-slot becomes + # `((5, 5), (5, 5))`, two shape leaves, faithful; the number `3` broadcast the same way becomes + # `(3, 3)`, which the predicate re-reads as one shape leaf, lossy. + assert optree.tree_broadcast_common( + [(5, 5), (1.0, 2)], + [(2, 'x'), 3], + is_leaf=is_shape_like, + ) == ([((5, 5), (5, 5)), (1.0, 2)], [(2, 'x'), (3, 3)]) + + def test_broadcast_common(): assert optree.broadcast_common(1, [2, 3, 4]) == ([1, 1, 1], [2, 3, 4]) assert optree.broadcast_common([1, 2, 3], [4, 5, 6]) == ([1, 2, 3], [4, 5, 6]) diff --git a/tests/test_registry.py b/tests/test_registry.py index 3900b3e5..c08b0718 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -19,6 +19,7 @@ import pickle import re import sys +import warnings import weakref from collections import UserDict, UserList, namedtuple from dataclasses import dataclass @@ -32,10 +33,14 @@ NODETYPE_REGISTRY, PYPY, Py_GIL_DISABLED, + check_script_in_subprocess, disable_systrace, gc_collect, parametrize, + skipif_android, + skipif_ios, skipif_pypy, + skipif_wasm, ) from optree.registry import DictMetaData @@ -401,6 +406,226 @@ def test_register_pytree_node_namedtuple(): assert treespec1 != treespec3 +def test_register_pytree_node_warning_as_error_does_not_corrupt_registry(): + # Registering a `namedtuple` / `PyStructSequence` subclass emits a `UserWarning`. Under + # warnings-as-errors that warning is raised as an exception; previously the C++ registry had + # already committed the registration and ignored the escalation, leaving corrupt state and a + # `SystemError`. The registration must instead be rolled back and the escalated warning raised. + mytuple = namedtuple('mytuple_warn_error', ['a', 'b']) # noqa: PYI024 + tree = mytuple(1, 2) + + with warnings.catch_warnings(): + warnings.simplefilter('error') + with pytest.raises(UserWarning): + optree.register_pytree_node( + mytuple, + lambda t: (tuple(t), None, None), + lambda _, t: mytuple(*t), + namespace='mytuple_warn_error', + ) + + # The escalated registration must not have stuck (no custom node registered). + assert 'CustomTreeNode' not in str( + optree.tree_structure(tree, namespace='mytuple_warn_error'), + ) + + # A subsequent non-error registration must succeed cleanly, proving no half-committed state. + with pytest.warns(UserWarning, match=re.escape('is a subclass of `collections.namedtuple`')): + optree.register_pytree_node( + mytuple, + lambda t: (tuple(t), None, None), + lambda _, t: mytuple(*t), + namespace='mytuple_warn_error', + ) + assert 'CustomTreeNode' in str( + optree.tree_structure(tree, namespace='mytuple_warn_error'), + ) + optree.unregister_pytree_node(mytuple, namespace='mytuple_warn_error') + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_register_pytree_node_no_deadlock_with_concurrent_flatten(): + # Regression: `register_pytree_node` takes the registry write lock, and the namedtuple / + # PyStructSequence detection it runs releases the GIL. Holding the registry lock across that GIL + # release inverts the GIL <-> lock order and deadlocks a concurrent `tree_flatten` that holds + # the GIL while waiting on the registry read lock, wedging the whole interpreter. Run register / + # unregister concurrently with flatten in a subprocess under a watchdog; the process must finish. + check_script_in_subprocess( + """ + import faulthandler + import threading + + import optree + + class MyType: + def __init__(self, children): + self.children = children + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + def register_loop(): + for _ in range(5000): + try: + optree.register_pytree_node( + MyType, + lambda o: (tuple(o.children), None, None), + lambda _, c: MyType(list(c)), + namespace='deadlock', + ) + except ValueError: + pass + try: + optree.unregister_pytree_node(MyType, namespace='deadlock') + except ValueError: + pass + + def flatten_loop(): + tree = {'a': 1, 'b': 2, 'c': 3} + for _ in range(5000): + optree.tree_flatten(tree) + + threads = [ + threading.Thread(target=register_loop), + threading.Thread(target=flatten_loop), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_unregister_pytree_node_not_found_no_deadlock_with_concurrent_flatten(): + # Regression: `unregister_pytree_node` takes the registry write lock, and on the NOT-FOUND path + # it builds its error message using the namedtuple / PyStructSequence detection, which releases + # the GIL. Holding the registry lock across that GIL release inverts the GIL <-> lock order and + # deadlocks a concurrent `tree_flatten` that holds the GIL while waiting on the registry read + # lock (the symmetric twin of the `register_pytree_node` deadlock). Unregister a never-registered + # type (always not-found) concurrently with flatten in a subprocess under a watchdog; the process + # must finish. + check_script_in_subprocess( + """ + import faulthandler + import threading + + import optree + + class NeverRegistered: # never registered -> unregister always hits the not-found path + pass + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + def unregister_loop(): + for _ in range(5000): + try: + optree.unregister_pytree_node(NeverRegistered, namespace='deadlock') + except ValueError: + pass + + def flatten_loop(): + tree = {'a': 1, 'b': 2, 'c': 3} + for _ in range(5000): + optree.tree_flatten(tree) + + threads = [threading.Thread(target=unregister_loop) for _ in range(2)] + threads += [threading.Thread(target=flatten_loop) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_get_registry_size_concurrent_no_spurious_error(): + # Regression: `GetRegistrySize()` reads the two internal registries (NoneIsNode / NoneIsLeaf) + # and checks they differ by exactly one (the extra `None` node). Reading them in two separate + # locked `Size()` calls dropped the lock between the reads, so a concurrent (un)registration + # could slip in and spuriously trip that check, surfacing as a `SystemError`. This manifests on + # free-threading builds (the GIL otherwise serializes the two reads). Hammer `get_registry_size()` + # while a type churns in and out of the registry, in a subprocess under a watchdog; the process + # must finish with no error recorded. + check_script_in_subprocess( + """ + import faulthandler + import threading + import time + + import optree + import optree._C as _C + + class MyType: + def __init__(self, children): + self.children = children + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + errors = [] + stop = threading.Event() + + def register_loop(): + for _ in range(5000): + try: + optree.register_pytree_node( + MyType, + lambda o: (tuple(o.children), None, None), + lambda _, c: MyType(list(c)), + namespace='size-race', + ) + except ValueError: + pass + try: + optree.unregister_pytree_node(MyType, namespace='size-race') + except ValueError: + pass + stop.set() + + def size_loop(): + while not stop.is_set(): + try: + _C.get_registry_size() + _C.get_registry_size('size-race') + except Exception as exc: # noqa: BLE001 + errors.append(repr(exc)) + return + # `get_registry_size()` takes the registry lock in read mode, so spinning here keeps + # overlapping shared holds on it and starves `register_loop`: `stop` is never set + # and the watchdog fires. The sleep must be non-zero, since `time.sleep(0)` maps to + # `Sleep(0)` on Windows and can return without ever dropping the lock. + time.sleep(0.001) + + threads = [ + threading.Thread(target=register_loop), + threading.Thread(target=size_loop), + threading.Thread(target=size_loop), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not errors, errors + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + def test_flatten_with_wrong_number_of_returns(): @optree.register_pytree_node_class(namespace='error') class MyList1(UserList): @@ -579,6 +804,36 @@ def __tree_unflatten__(cls, metadata, children): assert handler is registry3[cls] +def test_pytree_node_registry_get_named_namespace_takes_precedence_over_global(): + # When a type is registered in BOTH the global namespace and a named namespace, the dict form of + # `.get(namespace=...)` must return the NAMED handler for that type, matching the single-class + # lookup, which checks the named namespace before falling back to the global one. Registering + # the global entry LAST would otherwise let it shadow the named one in the returned dict. + class Both: + def __init__(self, a): + self.a = a + + def flatten(both): + return (both.a,), None, None + + def unflatten(metadata, children): + return Both(*children) + + optree.register_pytree_node(Both, flatten, unflatten, namespace='both-ns') + optree.register_pytree_node(Both, flatten, unflatten, namespace=GLOBAL_NAMESPACE) + try: + # Single-class lookup already prefers the named namespace. + assert optree.register_pytree_node.get(Both, namespace='both-ns').namespace == 'both-ns' + # The dict form must agree: the named handler wins over the global one. + registry = optree.register_pytree_node.get(namespace='both-ns') + assert registry[Both].namespace == 'both-ns' + # The global-only view still shows the global handler. + assert optree.register_pytree_node.get(namespace=GLOBAL_NAMESPACE)[Both].namespace == '' + finally: + optree.unregister_pytree_node(Both, namespace='both-ns') + optree.unregister_pytree_node(Both, namespace=GLOBAL_NAMESPACE) + + def test_pytree_node_registry_get_with_invalid_arguments(): registry = optree.register_pytree_node.get() assert optree.register_pytree_node.get(None) == registry diff --git a/tests/test_treespec.py b/tests/test_treespec.py index 643176a0..942f14e7 100644 --- a/tests/test_treespec.py +++ b/tests/test_treespec.py @@ -16,7 +16,10 @@ # pylint: disable=missing-function-docstring,invalid-name import contextlib +import gc import itertools +import math +import os import pickle import platform import re @@ -24,8 +27,10 @@ import subprocess import sys import tempfile +import time +import warnings import weakref -from collections import OrderedDict, UserList, defaultdict, deque +from collections import OrderedDict, UserList, defaultdict, deque, namedtuple import pytest @@ -33,6 +38,7 @@ import optree from helpers import ( GLOBAL_NAMESPACE, + HAS_DEFERRED_TYPE_REFS, NAMESPACED_TREE, PYPY, STANDARD_DICT_TYPES, @@ -48,6 +54,7 @@ parametrize, recursionlimit, skipif_android, + skipif_deferred_type_refs, skipif_ios, skipif_pypy, skipif_wasm, @@ -260,6 +267,25 @@ def test_treespec_string_representation(data): assert new_tree == reconstructed_tree +@skipif_pypy # CPython-only: `os.stat_result` slots 7, 8, 9 are unnamed; PyPy names them +def test_treespec_structseq_unnamed_field_string_representation(): + # `os.stat_result` renders its UNNAMED sequence slots (7, 8, 9) with the synthetic `` + # placeholder, following CPython's `` convention for names that are not identifiers, + # rather than the bare `unnamed field` marker which reads as an invalid keyword. CPython's + # `stat_result_desc` has pinned the 7 named + 3 unnamed sequence fields for 16 years, so the + # repr is asserted exactly; the hidden float `st_atime` fields (indices >= 10) are not part of + # the sequence and must not leak into it. + assert os.stat_result.n_sequence_fields == 10 + assert os.stat_result.n_unnamed_fields == 3 + st = os.stat_result(range(os.stat_result.n_fields)) + representation = str(optree.tree_structure(st)) + assert representation == ( + 'PyTreeSpec(os.stat_result(' + 'st_mode=*, st_ino=*, st_dev=*, st_nlink=*, st_uid=*, st_gid=*, st_size=*, ' + '=*, =*, =*))' + ) + + def test_treespec_with_empty_tuple_string_representation(): assert str(optree.tree_structure(())) == r'PyTreeSpec(())' @@ -276,6 +302,50 @@ def test_treespec_with_empty_dict_string_representation(): assert str(optree.tree_structure({})) == r'PyTreeSpec({})' +def test_treespec_namedtuple_repr_with_divergent_fields_raises_value_error(): + # If a namedtuple's `_fields` is mutated after the treespec is built, the recorded arity and the + # now-divergent field count disagree. The repr must raise a clear `ValueError` attributing the + # cause, not an `InternalError` telling the user to file a bug report. + Point = namedtuple('Point', ('x', 'y')) # noqa: PYI024 + treespec = optree.tree_structure(Point(1, 2)) + assert str(treespec) == 'PyTreeSpec(Point(x=*, y=*))' + + Point._fields = ('x', 'y', 'z') # diverge: 3 fields vs the treespec's arity of 2 + with pytest.raises(ValueError, match=r'does not match the arity'): + repr(treespec) + + +def test_treespec_setstate_rejects_structseq_field_arity_mismatch(): + # A PyStructSequence type's sequence-field count is fixed in C, so a node's arity must equal it + # (unlike a namedtuple, whose `_fields` can be mutated after the fact). `FromPicklable` (via + # `__setstate__`/`pickle`) must reject a crafted state pairing a PyStructSequence type with a + # mismatched arity at load time, rather than build a corrupt treespec that later aborts (e.g. in + # repr with an `InternalError`). + spec = optree.tree_structure(time.gmtime()) # struct_time: 9 sequence fields + node_states, none_is_leaf, namespace = spec.__getstate__() + # Swap the type to os.stat_result (10 sequence fields) while keeping the arity of 9. + crafted = tuple( + (kind, arity, os.stat_result if data is time.struct_time else data, *remaining) + for (kind, arity, data, *remaining) in node_states + ) + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + with pytest.raises(RuntimeError, match=r'does not match the arity'): + obj.__setstate__((crafted, none_is_leaf, namespace)) + + +def test_treespec_setstate_rejects_namedtuple_field_arity_mismatch(): + # A namedtuple's `_fields` can be mutated, so a crafted state can pair the type with an arity + # that no longer matches its field count. `FromPicklable` must reject it at load, rather than + # build a corrupt spec (the repr guards the post-load mutation case separately). + Point = namedtuple('Point', ('x', 'y')) # noqa: PYI024 + state = optree.tree_structure(Point(1, 2)).__getstate__() # arity 2 + Point._fields = ('x', 'y', 'z') # diverge: 3 fields vs the pickled arity of 2 + + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + with pytest.raises(RuntimeError, match=r'does not match the arity'): + obj.__setstate__(state) + + @disable_systrace def test_treespec_self_referential(): class Holder: @@ -344,6 +414,72 @@ def __repr__(self): assert wr() is None +@skipif_pypy # relies on CPython's reference-cycle collector +@skipif_deferred_type_refs +def test_treespec_custom_node_reference_cycle_is_collectable(): + # A treespec reaches a registered custom type through the shared registration held by its custom + # node. Once the registry no longer pins the registration, the node is its sole owner, so + # `PyTreeSpec::PyTpTraverse` reports those objects and the cyclic GC can collect a cycle through + # them. The cycle keeps the heap type alive, and free-threaded builds before 3.14 hold deferred + # references to type objects in per-thread caches, so `gc_collect()` cannot reclaim it there. + # 3.14t does, so this keeps free-threaded coverage. + class Cyclic: + pass + + optree.register_pytree_node( + Cyclic, + lambda cyclic: ((), None), + lambda metadata, children: None, # never called; the test only needs the registration + namespace='cycle_gc', + ) + try: + treespec = optree.tree_structure(Cyclic(), namespace='cycle_gc') + # Cycle: Cyclic -> __dict__ -> treespec -> registration -> Cyclic + Cyclic.self_spec = treespec + finally: + optree.unregister_pytree_node(Cyclic, namespace='cycle_gc') + + wr = weakref.ref(Cyclic) + del Cyclic, treespec + gc_collect() + assert wr() is None + + +@skipif_pypy # relies on CPython's reference-cycle collector +def test_treespec_shared_registration_refs_are_not_reported(): + # `PyTreeSpec::PyTpTraverse` must report only references the treespec owns. + # A shared registration holds one reference to each member however many nodes point at it, so + # reporting per node would underflow the object's shadow refcount and abort on debug builds. + # `gc.get_referrers()` walks `tp_traverse`, so it shows what the traversal reports. + class Shared: + pass + + optree.register_pytree_node( + Shared, + lambda shared: ((), None), + lambda metadata, children: None, + namespace='shared_gc', + ) + try: + # The registry holds the registration, so no treespec is its sole owner. + treespecs = [optree.tree_structure(Shared(), namespace='shared_gc') for _ in range(4)] + gc_collect() + assert not any(treespec in gc.get_referrers(Shared) for treespec in treespecs) + # The registration is also shared between treespecs, so dropping the registry's hold while + # more than one treespec remains must not make them report it either. + optree.unregister_pytree_node(Shared, namespace='shared_gc') + gc_collect() + assert not any(treespec in gc.get_referrers(Shared) for treespec in treespecs) + finally: + del treespecs + + wr = weakref.ref(Shared) + del Shared + gc_collect() + if not HAS_DEFERRED_TYPE_REFS: + assert wr() is None + + @disable_systrace def test_treeiter_self_referential(): sentinel = object() @@ -375,6 +511,24 @@ def test_treeiter_self_referential(): assert wr() is None +def test_treeiter_leaf_predicate_no_reference_leak(): + # A reference cycle that runs through the `leaf_predicate` callback must be collectable. + # Regression: `PyTreeIter` tp_traverse / tp_clear previously ignored `m_leaf_predicate`, so a + # cycle through the predicate was invisible to the cyclic garbage collector and leaked. + def is_leaf(x): + return False + + it = optree.tree_iter({'a': 1, 'b': {'c': 2}}, is_leaf) + wr = weakref.ref(it) + assert next(it) == 1 + is_leaf.self_ref = it # cycle: it -> m_leaf_predicate (is_leaf) -> is_leaf.self_ref -> it + + del it, is_leaf + gc_collect() + if not PYPY: + assert wr() is None + + def test_treespec_with_namespace(): tree = NAMESPACED_TREE @@ -541,6 +695,29 @@ def test_treespec_pickle_roundtrip( ) +@skipif_wasm +@skipif_android +@skipif_ios +def test_treespec_pickle_all_protocols_roundtrip(): + # pybind11's pickle support reconstructs cleanly only at protocol >= 2. Protocols 0 and 1 used + # to reconstruct via `object.__new__`, which pybind11 rejects with an untranslated C++ exception + # that aborts the interpreter (SIGABRT). Run in a subprocess so a regression fails this test + # rather than killing the whole suite. + check_script_in_subprocess( + r""" + import pickle + + import optree + + spec = optree.tree_structure({'a': [1, 2], 'b': (3, 4)}) + for protocol in range(pickle.HIGHEST_PROTOCOL + 1): + restored = pickle.loads(pickle.dumps(spec, protocol=protocol)) + assert restored == spec, (protocol, restored, spec) + """, + output=None, + ) + + class Foo: def __init__(self, x, y): self.x = x @@ -592,6 +769,330 @@ def test_treespec_pickle_missing_registration(): treespec = pickle.loads(serialized) +def test_treespec_getstate_does_not_alias_internal_node_data(): + # `__getstate__` (used by `pickle`) must return a snapshot, not aliases of the immutable spec's + # internal mutable containers: the keys of a dict/OrderedDict/defaultdict node and its + # insertion-order keys dict. Mutating the returned state otherwise reaches back into the spec, + # desyncing the keys from the arity (repr raises an InternalError) or adding a spurious + # original key (unflatten returns an extra entry). A custom node's entries are immutable. + class Custom: + def __init__(self, *values): + self.values = values + + optree.register_pytree_node( + Custom, + lambda custom: (custom.values, None, tuple(range(len(custom.values)))), + lambda metadata, children: Custom(*children), + namespace='getstate_snapshot', + ) + try: + tree = { + 'b': Custom(1, 2), + 'a': 3, + 'od': OrderedDict([('y', 4), ('x', 5)]), + 'dd': defaultdict(int, {'q': 6, 'p': 7}), + } + spec = optree.tree_structure(tree, namespace='getstate_snapshot') + node_states, _, _ = state = spec.__getstate__() + before = repr(state) + + for node in node_states: + kind, node_data, node_entries, original_keys = node[0], node[2], node[3], node[7] + if kind in {optree.PyTreeKind.DICT, optree.PyTreeKind.ORDEREDDICT}: + node_data.append('injected') # a dict/OrderedDict node's keys list + elif kind == optree.PyTreeKind.DEFAULTDICT: + node_data[1].append('injected') # a defaultdict's (default_factory, keys) tuple + if isinstance(original_keys, dict): + original_keys['injected'] = None + assert node_entries is None or isinstance(node_entries, tuple) # entries are immutable + + assert repr(spec.__getstate__()) == before, 'mutating the pickled state corrupted the spec' + finally: + optree.unregister_pytree_node(Custom, namespace='getstate_snapshot') + + +def test_treespec_getstate_aliases_custom_node_data(): + # Limitation (characterization test): a custom node's `node_data` is the user-provided metadata, + # which `__getstate__` passes through by reference. optree copies its own dict/defaultdict keys + # (see `test_treespec_getstate_does_not_alias_internal_node_data`) but cannot generically + # deep-copy arbitrary metadata, so mutating it via the pickled state reaches back into the spec. + # Protecting custom metadata is the caller's responsibility; this pins the behavior. + class Custom: + def __init__(self, *children, alpha, beta=None): + self.children = children + self.metadata = {'alpha': alpha, 'beta': beta} + + def __eq__(self, other): + return ( + isinstance(other, Custom) + and self.children == other.children + and self.metadata == other.metadata + ) + + __hash__ = None + + optree.register_pytree_node( + Custom, + lambda custom: (custom.children, custom.metadata), # mutable dict metadata + lambda metadata, children: Custom(*children, **metadata), + namespace='getstate_alias_custom', + ) + try: + leaves, treespec = optree.tree_flatten( + Custom(1, 2, alpha=3, beta=4), + namespace='getstate_alias_custom', + ) + before = repr(treespec) + custom_state = next( + node for node in treespec.__getstate__()[0] if node[0] == optree.PyTreeKind.CUSTOM + ) + assert custom_state[2] == {'alpha': 3, 'beta': 4} + + # Mutate the aliased metadata in place via the pickled state. + custom_state[2]['gamma'] = 5 # mutate the aliased metadata in place + + aliased = next( + node for node in treespec.__getstate__()[0] if node[0] == optree.PyTreeKind.CUSTOM + ) + assert aliased[2] is custom_state[2] + assert aliased[2] == {'alpha': 3, 'beta': 4, 'gamma': 5} # the mutation reached the spec + assert repr(treespec) == before.replace( + repr({'alpha': 3, 'beta': 4}), + repr({'alpha': 3, 'beta': 4, 'gamma': 5}), + ) + # The corruption even reaches what `tree_unflatten` rebuilds, not just repr/getstate. + with pytest.raises(TypeError, match=r'unexpected keyword argument'): + optree.tree_unflatten(treespec, leaves) + + # Replacing the metadata wholesale reaches the spec the same way, but here unflatten + # succeeds and rebuilds a different object: the corruption is silent, not an error. + custom_state[2].clear() + custom_state[2]['alpha'] = 42 + aliased = next( + node for node in treespec.__getstate__()[0] if node[0] == optree.PyTreeKind.CUSTOM + ) + assert aliased[2] is custom_state[2] + assert aliased[2] == {'alpha': 42} + assert repr(treespec) == before.replace( + repr({'alpha': 3, 'beta': 4}), + repr({'alpha': 42}), + ) + reconstructed = optree.tree_unflatten(treespec, leaves) + reconstructed_treespec = optree.tree_structure( + reconstructed, + namespace='getstate_alias_custom', + ) + assert reconstructed == Custom(1, 2, alpha=42) + assert reconstructed.metadata == {'alpha': 42, 'beta': None} + assert reconstructed_treespec != treespec + assert repr(reconstructed_treespec) == before.replace( + repr({'alpha': 3, 'beta': 4}), + repr({'alpha': 42, 'beta': None}), + ) + finally: + optree.unregister_pytree_node(Custom, namespace='getstate_alias_custom') + + +def test_treespec_setstate_rejects_malformed_state(): + # `PyTreeSpec.__setstate__` (used by `pickle`) must reject structurally malformed state rather + # than build a corrupt spec that triggers out-of-bounds reads / crashes when later used. The + # per-node tuple layout is (kind, arity, node_data, node_entries, custom, num_leaves, num_nodes, + # original_keys); see `PyTreeSpec::FromPicklable`. + def setstate(state): + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + obj.__setstate__(state) + return obj + + CUSTOM = int(optree.PyTreeKind.CUSTOM) # noqa: N806 + LEAF = int(optree.PyTreeKind.LEAF) # noqa: N806 + NONE = int(optree.PyTreeKind.NONE) # noqa: N806 + TUPLE = int(optree.PyTreeKind.TUPLE) # noqa: N806 + DICT = int(optree.PyTreeKind.DICT) # noqa: N806 + NAMEDTUPLE = int(optree.PyTreeKind.NAMEDTUPLE) # noqa: N806 + DEFAULTDICT = int(optree.PyTreeKind.DEFAULTDICT) # noqa: N806 + DEQUE = int(optree.PyTreeKind.DEQUE) # noqa: N806 + STRUCTSEQUENCE = int(optree.PyTreeKind.STRUCTSEQUENCE) # noqa: N806 + NUM_KINDS = int(optree.PyTreeKind.NUM_KINDS) # noqa: N806 + leaf_node = (LEAF, 0, None, None, None, 1, 1, None) # arity 0, 1 leaf, 1 node + keys_ab = {'a': None, 'b': None} # original_keys for a 2-key ('a', 'b') dict node + + # Sanity: well-formed states still round-trip. + for spec in [ + optree.tree_structure((0, 0)), + optree.tree_structure({'a': 0, 'b': 0}), + optree.tree_structure(defaultdict(int, {'a': 0, 'b': 0})), + ]: + assert setstate(spec.__getstate__()) == spec + + malformed_exceptions = (RuntimeError, ValueError, TypeError) + + # The rejection cases below follow the order of the checks in `PyTreeSpec::FromPicklable`. + + # A state that is not a 3-tuple. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node,), False)) + + # A node state that is not a 7- or 8-tuple. + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, 1),), False, '')) + + # Kind out of range: the raw integer is validated before the narrowing `uint8_t` enum cast, + # which would otherwise wrap a bogus value to a valid-looking kind. + with pytest.raises(malformed_exceptions): + setstate((((NUM_KINDS, 0, None, None, None, 0, 1, None),), False, '')) + + # Negative arity. + with pytest.raises(malformed_exceptions): + setstate((((TUPLE, -1, None, None, None, 0, 1, None),), False, '')) + + # A dict node missing its original keys, and a non-dict node carrying them. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, (DICT, 2, ['a', 'b'], None, None, 2, 3, None)), False, '')) + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, 1, 1, keys_ab),), False, '')) + + # A negative leaf count, or a non-positive node count. + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, -1, 1, None),), False, '')) + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, 1, 0, None),), False, '')) + + # Node data on a leaf or none node (childless kinds that must not carry any). + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, 'data', None, None, 1, 1, None),), False, '')) + + # Leaf or none nodes are childless; a nonzero arity absorbs the preceding subtrees while still + # folding consistently, so the reconstructed spec reports a leaf/None while its num_leaves counts + # the absorbed children and unflatten silently drops them. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, (NONE, 1, None, None, None, 1, 2, None)), False, '')) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, (LEAF, 1, None, None, None, 1, 2, None)), False, '')) + + # A None-kind node cannot appear when none_is_leaf is set (None is flattened as a leaf then, so + # a flattened tree never contains a None node); accepting one later raises an InternalError. + with pytest.raises(malformed_exceptions): + setstate((((NONE, 0, None, None, None, 0, 1, None),), True, '')) + + # Node data on a tuple or list node. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, (TUPLE, 2, 'data', None, None, 2, 3, None)), False, '')) + + # Dict key list shorter than arity (MakeNode would index past the list end). + short_keys = (DICT, 2, ['a'], None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, short_keys), False, '')) + + # Dict with duplicate keys (would collapse the rebuilt dict), and with an unhashable key. + dup_keys = (DICT, 2, ['a', 'a'], None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, dup_keys), False, '')) + unhashable_key = (DICT, 2, [[], []], None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, unhashable_key), False, '')) + + # NamedTuple / StructSequence node_data that is not the expected kind of type. + with pytest.raises(malformed_exceptions): + setstate((((NAMEDTUPLE, 0, int, None, None, 0, 1, None),), False, '')) + with pytest.raises(malformed_exceptions): + setstate((((STRUCTSEQUENCE, 0, int, None, None, 0, 1, None),), False, '')) + + # DefaultDict metadata as a list where a 2-tuple is expected previously caused a raw tuple-item + # read to segfault; it is now coerced to a tuple and used safely. + restored = setstate( + ( + ( + leaf_node, + leaf_node, + (DEFAULTDICT, 2, [int, ['a', 'b']], None, None, 2, 3, keys_ab), + ), + False, + '', + ), + ) + assert optree.tree_unflatten(restored, [10, 20]) == defaultdict(int, {'a': 10, 'b': 20}) + + # DefaultDict metadata with the wrong tuple size is rejected. + wrong_metadata = (DEFAULTDICT, 2, (int, ['a', 'b'], 'extra'), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, wrong_metadata), False, '')) + + # DefaultDict default_factory that is neither None nor callable. + bad_factory = (DEFAULTDICT, 2, (42, ['a', 'b']), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, bad_factory), False, '')) + + # DefaultDict keys too few, and DefaultDict keys not distinct (the Dict variants are above). + defaultdict_short = (DEFAULTDICT, 2, (int, ['a']), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, defaultdict_short), False, '')) + defaultdict_dup = (DEFAULTDICT, 2, (int, ['a', 'a']), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, defaultdict_dup), False, '')) + + # Deque maxlen that is neither None nor an int, and maxlen smaller than the arity (a deque holds + # at most maxlen items, so arity <= maxlen). + with pytest.raises(malformed_exceptions): + setstate((((DEQUE, 0, 'x', None, None, 0, 1, None),), False, '')) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, (DEQUE, 2, 1, None, None, 2, 3, None)), False, '')) + + # A non-custom node carrying node entries or a custom type. + with pytest.raises(malformed_exceptions): + setstate( + ((leaf_node, leaf_node, (TUPLE, 2, None, ('a', 'b'), None, 2, 3, None)), False, ''), + ) + + # Original keys whose count (not just key set) disagrees with the arity. + short_original = (DICT, 2, ['a', 'b'], None, None, 2, 3, {'a': None}) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, short_original), False, '')) + + # Dict original_keys whose key set differs from the sorted key list. + mismatched_original = (DICT, 2, ['a', 'b'], None, None, 2, 3, {'a': None, 'c': None}) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, mismatched_original), False, '')) + + # A custom node whose node-entries count disagrees with the arity (needs a registered type). + class MalformedCustomNode: + pass + + optree.register_pytree_node( + MalformedCustomNode, + lambda obj: ((), None), + lambda metadata, children: MalformedCustomNode(), + namespace='malformed', + ) + try: + custom_node = (CUSTOM, 2, None, ('one-entry',), MalformedCustomNode, 2, 3, None) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, custom_node), False, 'malformed')) + finally: + optree.unregister_pytree_node(MalformedCustomNode, namespace='malformed') + + # A node claiming more children than the traversal provides. + with pytest.raises(malformed_exceptions): + setstate((((TUPLE, 2, None, None, None, 2, 3, None),), False, '')) + + # Inconsistent intermediate num_nodes (previously only the last node was checked). + with pytest.raises(malformed_exceptions): + setstate( + ( + ( + (LEAF, 0, None, None, None, 1, 5, None), # leaf claims num_nodes == 5 + leaf_node, + (TUPLE, 2, None, None, None, 2, 3, None), + ), + False, + '', + ), + ) + + # A traversal that yields more than one tree. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node), False, '')) + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -744,6 +1245,541 @@ def test_treespec_compose_children( assert optree.treespec_is_suffix(expected_treespec, treespec, strict=False) +def test_treespec_compose_rejects_incompatible_namespace_merge(): + # Regression: composing an empty-namespace spec (whose custom nodes are resolved globally) with + # a spec in another namespace adopted that namespace but kept the global registrations. When the + # same type is registered differently in the two namespaces, the composed spec silently used the + # wrong flatten/unflatten (spurious flatten_up_to errors; corrupt pickle). Reject the merge. + class Pair: + def __init__(self, a, b): + self.a, self.b = a, b + + class Single: + def __init__(self, x): + self.x = x + + optree.register_pytree_node( + Pair, + lambda t: ((t.a, t.b), None, None), + lambda m, c: Pair(c[0], c[1]), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( # behavior differs from the global registration + Pair, + lambda t: ((t.b, t.a), None, None), + lambda m, c: Pair(c[1], c[0]), + namespace='behavior_change', + ) + optree.register_pytree_node( + Single, + lambda t: ((t.x,), None, None), + lambda m, c: Single(c[0]), + namespace='behavior_change', + ) + try: + outer = optree.tree_structure(Pair(0, 0)) + inner = optree.tree_structure(Single(0), namespace='behavior_change') + assert outer.namespace == '' + assert inner.namespace == 'behavior_change' + with pytest.raises(ValueError, match='different registration'): + outer.compose(inner) + + # `tree_transpose` builds its expected structure with `compose`, so the rejection surfaces + # through the public API too (here via its structure-mismatch diagnostic path). + with pytest.raises(ValueError, match='different registration'): + optree.tree_transpose(outer, inner, [1, 2, 3]) + + # `broadcast_to_common_suffix` adopts the namespace the same way. + self_spec = optree.tree_structure({'k': Pair(0, 0)}) + other_spec = optree.tree_structure({'k': Single(0)}, namespace='behavior_change') + with pytest.raises(ValueError, match='different registration'): + self_spec.broadcast_to_common_suffix(other_spec) + finally: + optree.unregister_pytree_node(Pair, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(Pair, namespace='behavior_change') + optree.unregister_pytree_node(Single, namespace='behavior_change') + + +def test_treespec_compose_allows_compatible_namespace_merge(): + # The namespace-merge rejection must not over-reject: a custom type registered only globally + # resolves identically under any namespace (via global fallback), so merging an empty-namespace + # spec that uses it into another namespace is allowed and the result stays consistent. + class GlobalOnly: + def __init__(self, a, b): + self.a, self.b = a, b + + optree.register_pytree_node( + GlobalOnly, + lambda t: ((t.a, t.b), None, None), + lambda m, c: GlobalOnly(c[0], c[1]), + namespace=GLOBAL_NAMESPACE, + ) + try: + outer = optree.tree_structure(GlobalOnly(0, 0)) + assert outer.namespace == '' + + # Both empty -> the merge stays in the global namespace. + assert outer.compose(optree.tree_structure(0)).namespace == '' + + # Empty side (global-only custom) merged into a namespace: allowed, adopts the namespace, + # and unflattens consistently (the global registration is used throughout). + with optree.dict_insertion_ordered(True, namespace='no_override'): + inner = optree.tree_structure({'x': 0}, namespace='no_override') + assert inner.namespace == 'no_override' + composed = outer.compose(inner) + assert composed.namespace == 'no_override' + result = optree.tree_unflatten(composed, [1, 2]) + assert isinstance(result, GlobalOnly) + assert result.a == {'x': 1} + assert result.b == {'x': 2} + + # The cross-namespace merge equals building the composed structure directly with `tree_map` + # in the adopted namespace, compose's defining identity. + expected = optree.tree_structure( + optree.tree_map(lambda _: {'x': 0}, GlobalOnly(0, 0), namespace='no_override'), + namespace='no_override', + ) + assert composed == expected + + # broadcast_to_common_suffix likewise allows the compatible merge. + broadcasted = outer.broadcast_to_common_suffix( + optree.tree_structure(GlobalOnly(0, 0), namespace='no_override'), + ) + assert broadcasted.namespace == 'no_override' + finally: + optree.unregister_pytree_node(GlobalOnly, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_broadcast_to_common_suffix_does_not_mutate_argument_on_key_mismatch(): + # Regression: BroadcastToCommonSuffixImpl built the "got key(s)" part of its key-mismatch error + # message by sorting the ARGUMENT spec's live dict-node key list IN PLACE: `other_keys` was a + # borrow of `node_data`, not a copy. For an OrderedDict the child subtrees stay in insertion + # order while the keys get permuted, silently corrupting a spec the caller still holds: repr, + # equality, hash, and unflatten all go wrong. The message must be built from a sorted COPY. + other = optree.tree_structure(OrderedDict([('c', 1), ('b', 2)])) + before_repr = str(other) + before_hash = hash(other) + this = optree.tree_structure({'a': 1}) + with pytest.raises(ValueError, match='dictionary key mismatch'): + this.broadcast_to_common_suffix(other) + # The argument spec must be byte-for-byte unchanged by the failed call. + assert str(other) == before_repr + assert hash(other) == before_hash + # And it must still unflatten in its ORIGINAL insertion order (c, b), not a sorted (b, c) order. + assert other.unflatten([10, 20]) == OrderedDict([('c', 10), ('b', 20)]) + + +def test_treespec_broadcast_to_common_suffix_preserves_custom_node_entries(): + # Regression: BroadcastToCommonSuffixImpl rebuilds each non-leaf node with a designated + # initializer that lists `.node_data` then jumps to `.custom`, silently skipping `.node_entries` + # (which is declared between them). C++ value-initializes the omitted member to a null handle, + # so a custom node registered with explicit path entries (a 3-tuple flatten `(children, + # metadata, entries)`) lost them in the broadcasted spec. `entries()`/`paths()`/`accessors()` + # then fell back to the `range(arity)` integer indices, producing not just different but BROKEN + # accessors (`GetAttrEntry(entry=0)` calls `getattr(obj, 0)`, a `TypeError`, instead of + # `getattr(obj, 'a')`). + class Vector: + def __init__(self, a, c): + self.a, self.c = a, c + + optree.register_pytree_node( + Vector, + lambda o: ((o.a, o.c), None, ('a', 'c')), # 3-tuple flatten -> node_entries = ('a', 'c') + lambda metadata, children: Vector(*children), + path_entry_type=optree.GetAttrEntry, + namespace=GLOBAL_NAMESPACE, + ) + try: + spec = optree.tree_structure(Vector(1, 2)) + other = optree.tree_structure(Vector(3, 4)) + assert spec.entries() == ['a', 'c'] + + # Both specs share the same custom structure, so the common suffix is that structure and the + # explicit string entries must survive unchanged, not degrade to the fallback [0, 1]. + broadcasted = spec.broadcast_to_common_suffix(other) + assert broadcasted.entries() == ['a', 'c'] + assert broadcasted.paths() == spec.paths() + assert broadcasted.accessors() == spec.accessors() + finally: + optree.unregister_pytree_node(Vector, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_deep_walk_raises_recursion_error_not_segfault(): + # Regression: `PathsImpl`, `AccessorsImpl`, and `BroadcastToCommonSuffixImpl` recurse once per + # tree level. Without a depth guard, a deeply-nested spec (trivially built via doubling + # `compose`) overflowed the native C++ stack and crashed the interpreter with a SIGSEGV instead + # of raising a catchable `RecursionError`. + # Each `compose` doubles the depth, so ceil(log2(limit)) + 1 composes push it above the limit. + num_composes = math.ceil(math.log2(optree.MAX_RECURSION_DEPTH)) + 1 + deep = optree.tree_structure([0]) + for _ in range(num_composes): + deep = deep.compose(deep) + assert 2**num_composes > optree.MAX_RECURSION_DEPTH + with pytest.raises(RecursionError): + deep.paths() + with pytest.raises(RecursionError): + deep.accessors() + with pytest.raises(RecursionError): + deep.broadcast_to_common_suffix(deep) + + # Broadcasting the deep spec against a shallower spec whose depth is still below the limit + # recurses only as far as the common suffix, so it must succeed (not raise RecursionError or + # crash) in either direction, returning the deeper spec. + shallower = optree.tree_structure([0]) + for _ in range(num_composes - 2): # depth 2 ** (num_composes - 2), safely below the limit + shallower = shallower.compose(shallower) + assert 2 ** (num_composes - 2) < optree.MAX_RECURSION_DEPTH + assert deep.broadcast_to_common_suffix(shallower) == deep + assert shallower.broadcast_to_common_suffix(deep) == deep + + +def test_treespec_compose_rejects_namespace_override_with_different_arity(): + # A type registered globally flattens both members as children (arity 2); a namespace override + # flattens one member as a child and stores the other as node metadata (arity 1). Both + # registrations round-trip, but merging an empty-namespace spec (global, arity 2) into that + # namespace must be rejected: the composed spec would claim the namespace while carrying an + # arity-2 node that the namespace's registration cannot unflatten. + class TwoMember: + def __init__(self, a, b): + self.a, self.b = a, b + + def __eq__(self, other): + return isinstance(other, TwoMember) and (self.a, self.b) == (other.a, other.b) + + __hash__ = None + + optree.register_pytree_node( + TwoMember, + lambda t: ((t.a, t.b), None, None), # global: both members are children + lambda metadata, children: TwoMember(children[0], children[1]), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( + TwoMember, + lambda t: ((t.a,), t.b, None), # override: one child, the other is metadata + lambda metadata, children: TwoMember(children[0], metadata), + namespace='arity_change', + ) + try: + obj = TwoMember(1, 2) + + # Both registrations round-trip on their own. + global_leaves, global_spec = optree.tree_flatten(obj) + assert global_leaves == [1, 2] + assert optree.tree_unflatten(global_spec, global_leaves) == obj + custom_leaves, custom_spec = optree.tree_flatten(obj, namespace='arity_change') + assert custom_leaves == [1] + assert optree.tree_unflatten(custom_spec, custom_leaves) == obj + + assert global_spec.namespace == '' + assert global_spec.num_leaves == 2 + assert custom_spec.namespace == 'arity_change' + assert custom_spec.num_leaves == 1 + with pytest.raises(ValueError, match='different registration'): + global_spec.compose(custom_spec) + finally: + optree.unregister_pytree_node(TwoMember, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(TwoMember, namespace='arity_change') + + +def test_treespec_transform_rejects_incompatible_namespace_merge(): + # `transform` unifies the namespace across the input spec and the transform outputs. If that + # unified (non-empty) namespace rebinds a custom node (e.g. the input's globally-resolved + # custom node) to a different registration, the transform must be rejected (same class as the + # compose / broadcast merge rejection). A globally-only-registered type is still allowed via + # fallback. + class Diverge: # variable arity; registered differently in the global and named namespaces + def __init__(self, *children): + self.children = children + + class GlobalOnly: # variable arity; registered only globally -> resolves via fallback anywhere + def __init__(self, *children): + self.children = children + + optree.register_pytree_node( + Diverge, + lambda d: (d.children, None, None), + lambda metadata, children: Diverge(*children), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( + Diverge, + lambda d: (tuple(reversed(d.children)), None, None), # divergent from the global reg + lambda metadata, children: Diverge(*reversed(children)), + namespace='transform_change', + ) + optree.register_pytree_node( + GlobalOnly, + lambda g: (g.children, None, None), + lambda metadata, children: GlobalOnly(*children), + namespace=GLOBAL_NAMESPACE, + ) + + def to_namespaced_leaf(_): + # Replace a leaf with a namespaced Diverge to inject the namespace (leaves have no arity). + return optree.tree_structure(Diverge(0), namespace='transform_change') + + def to_global_node(spec): + # Replace a node with a same-arity globally-resolved Diverge; it rebinds under the promoted + # namespace (Diverge is registered differently there). Generic over the node's arity. + return optree.tree_structure(Diverge(*range(spec.num_children))) + + def to_namespaced_node(spec): + # Outer node -> global Diverge (rebinds); inner nodes -> namespaced Diverge (injects the + # namespace). Both same-arity, so `f_node` alone drives the (f_node, None) rejection. + if spec.type is tuple: + return to_global_node(spec) + return optree.tree_structure( + Diverge(*range(spec.num_children)), + namespace='transform_change', + ) + + try: + # The rejection must fire for every `(f_node, f_leaf)` combination that puts a + # globally-resolved custom node under the non-empty unified namespace. + + # (None, f_leaf): the input's global Diverge is kept, f_leaf injects the namespace. + outer = optree.tree_structure(Diverge(0, 0)) + assert outer.namespace == '' + with pytest.raises(ValueError, match='different registration'): + outer.transform(None, to_namespaced_leaf) + + # (f_node, None): f_node alone yields a global Diverge above namespaced children. + with pytest.raises(ValueError, match='different registration'): + optree.tree_structure(([0], [0])).transform(to_namespaced_node, None) + + # (f_node, f_leaf): f_node injects the global Diverge, f_leaf injects the namespace. + with pytest.raises(ValueError, match='different registration'): + optree.tree_structure([0, 0]).transform(to_global_node, to_namespaced_leaf) + + # Compatible: GlobalOnly resolves identically under any namespace via fallback. + global_outer = optree.tree_structure(GlobalOnly(0, 0)) + transformed = global_outer.transform(None, to_namespaced_leaf) + assert transformed.namespace == 'transform_change' + finally: + optree.unregister_pytree_node(Diverge, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(Diverge, namespace='transform_change') + optree.unregister_pytree_node(GlobalOnly, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_from_collection_rejects_incompatible_namespace_promotion(): + # `treespec_from_collection` promotes an empty caller namespace to a child spec's namespace. If + # that promoted namespace rebinds a custom node the collection resolved globally (the root node, + # or a globally-resolved child) to a different registration, the result would claim the + # namespace while carrying the wrong registration: it must be rejected, exactly like compose / + # transform / broadcast. A globally-only-registered type is still allowed via fallback. + class Diverge: # variable arity; registered differently in the global and named namespaces + def __init__(self, *children): + self.children = children + + class GlobalOnly: # variable arity; registered only globally -> resolves via fallback anywhere + def __init__(self, *children): + self.children = children + + optree.register_pytree_node( + Diverge, + lambda d: (d.children, None, None), + lambda metadata, children: Diverge(*children), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( + Diverge, + lambda d: (tuple(reversed(d.children)), None, None), # divergent from the global reg + lambda metadata, children: Diverge(*reversed(children)), + namespace='from_coll_change', + ) + optree.register_pytree_node( + GlobalOnly, + lambda g: (g.children, None, None), + lambda metadata, children: GlobalOnly(*children), + namespace=GLOBAL_NAMESPACE, + ) + try: + # Incompatible: the globally-resolved Diverge rebinds under the promoted namespace. + foo = optree.tree_structure(Diverge(0, 0)) + child = optree.tree_structure(Diverge(0), namespace='from_coll_change') + assert foo.namespace == '' + with pytest.raises(ValueError, match='different registration'): + optree.treespec_from_collection([foo, child], namespace='') + + # Compatible: GlobalOnly resolves identically under any namespace via fallback. + global_spec = optree.tree_structure(GlobalOnly(0, 0)) + promoted = optree.treespec_from_collection([global_spec, child], namespace='') + assert promoted.namespace == 'from_coll_change' + finally: + optree.unregister_pytree_node(Diverge, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(Diverge, namespace='from_coll_change') + optree.unregister_pytree_node(GlobalOnly, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_dict_key_order_survives_namespace_promotion(): + # A dict node's key order is fixed at BUILD time by the namespace passed then. Operations that + # merge/promote a spec's namespace (`treespec_from_collection`, `compose`, `transform`) only + # re-tag it for custom-node resolution; like `compose` they NEVER reorder an already-built dict. + # So a dict built under the global ('') namespace (sorted keys) keeps that order even after + # promotion to an insertion-ordered namespace, intentionally differing from the same dict built + # directly under that namespace, while a dict built directly under the namespace keeps its + # insertion order (matching). This test locks that behavior across all three operations. + class Wrap: # a variable-arity custom node, so `f_node` can build same-arity replacements + def __init__(self, *children): + self.children = children + + optree.register_pytree_node( + Wrap, + lambda w: (w.children, None, None), + lambda metadata, children: Wrap(*children), + namespace='promote_order', + ) + try: + with optree.dict_insertion_ordered(True, namespace='promote_order'): + child = optree.tree_structure(Wrap(0), namespace='promote_order') + # Built directly under the insertion-ordered namespace: keys in insertion order (b, a). + genuine = optree.tree_structure({'b': Wrap(0), 'a': Wrap(0)}, namespace='promote_order') + assert genuine.entries() == ['b', 'a'] + + def to_namespaced_node(spec): + # Rewrite every non-dict node into a same-arity `Wrap` in the namespace (generic + # over the node's arity rather than tied to this test's shapes) so `f_node` alone + # can promote the spec. `transform` promotes only when some output carries a + # namespace, and only a custom node can. The outer dict node is kept so its key + # order stays observable. + if spec.type is dict: + return spec + return optree.tree_structure( + Wrap(*range(spec.num_children)), + namespace='promote_order', + ) + + def transform_combos(outer): + return { + 'transform(None, f_leaf)': outer.transform(None, lambda _: child), + 'transform(f_node, None)': outer.transform(to_namespaced_node, None), + 'transform(f_node, f_leaf)': outer.transform( + to_namespaced_node, + lambda _: child, + ), + } + + # Dicts built under the GLOBAL ('') namespace, sorted keys (a, b), then promoted. + from_global = { + 'from_collection': optree.treespec_from_collection( + {'b': child, 'a': child}, + namespace='', + ), + 'compose': optree.tree_structure({'b': 0, 'a': 0}).compose(child), + **transform_combos(optree.tree_structure({'b': [0], 'a': [0]})), + } + # Dicts built directly under the namespace, insertion-order keys (b, a). + from_namespace = { + 'from_collection': optree.treespec_from_collection( + {'b': child, 'a': child}, + namespace='promote_order', + ), + 'compose': optree.tree_structure( + {'b': 0, 'a': 0}, + namespace='promote_order', + ).compose(child), + **transform_combos( + optree.tree_structure({'b': [0], 'a': [0]}, namespace='promote_order'), + ), + } + + for name, spec in from_global.items(): + assert spec.namespace == 'promote_order', name # promoted for custom resolution ... + assert spec.entries() == ['a', 'b'], name # ... but the dict keeps '' (sorted) order + + for name, spec in from_namespace.items(): + assert spec.namespace == 'promote_order', name + assert spec.entries() == ['b', 'a'], name # insertion order kept, matches direct build + + # from_collection / compose reproduce genuine's flat structure exactly, so the only + # difference is the dict key order: global-built differs, namespace-built matches. + assert from_global['from_collection'] != genuine + assert from_global['compose'] != genuine + assert from_namespace['from_collection'] == genuine + assert from_namespace['compose'] == genuine + finally: + optree.unregister_pytree_node(Wrap, namespace='promote_order') + + +def test_treespec_is_prefix_nested_dict_key_reorder(): + # Regression: `IsPrefix` reorders a dict node's children in a working copy of the traversal to + # make key order irrelevant. When a NESTED dict also needed reordering, it indexed the pristine + # traversal by an offset into the already-mutated working copy, corrupting it -> a spurious + # `optree._C.InternalError` or a wrong boolean. Two treespecs that describe the SAME tree + # (differing only in dict key insertion order, at nested levels) must be mutual non-strict + # prefixes / suffixes. + + # Top-level AND nested dict keys reordered; the top-level reorder relocates the nested dict. + tree_a = OrderedDict([('a', 0), ('b', OrderedDict([('e', 0), ('g', 0)])), ('d', 0)]) + tree_b = OrderedDict([('b', OrderedDict([('g', 0), ('e', 0)])), ('a', 0), ('d', 0)]) + a = optree.tree_structure(tree_a) + b = optree.tree_structure(tree_b) + assert optree.treespec_is_prefix(a, b, strict=False) + assert optree.treespec_is_prefix(b, a, strict=False) + assert optree.treespec_is_suffix(a, b, strict=False) + assert optree.treespec_is_suffix(b, a, strict=False) + assert a <= b + assert b <= a + assert a >= b + assert b >= a + + # A nested dict whose reorder relocates a subtree containing another out-of-order dict. + tree_a2 = OrderedDict([('a', 0), ('d', 0), ('b', OrderedDict([('e', 0), ('f', 0)]))]) + tree_b2 = OrderedDict([('b', OrderedDict([('f', 0), ('e', 0)])), ('d', 0), ('a', 0)]) + a2 = optree.tree_structure(tree_a2) + b2 = optree.tree_structure(tree_b2) + assert optree.treespec_is_prefix(a2, b2, strict=False) + assert optree.treespec_is_prefix(b2, a2, strict=False) + assert a2 <= b2 + assert b2 <= a2 + + +def test_treespec_is_prefix_deque_maxlen_agnostic(): + # A deque's treespec stores both its arity and its `maxlen`, but `is_prefix` is arity-based and + # deliberately `maxlen`-AGNOSTIC. A deque holds at most `maxlen` items, so `arity <= maxlen` + # always holds and two flatten-compatible deques necessarily share the same arity while carrying + # any `maxlen1`/`maxlen2`; `maxlen` does not affect how children are partitioned, so gating the + # prefix relation on it would wrongly reject valid `flatten_up_to`/`broadcast_prefix` operations. + # `EqualTo`, by contrast, IS `maxlen`-sensitive (`unflatten` restores the exact `maxlen`), so + # `a <= b and b <= a` does NOT imply `a == b`: `is_prefix` is a preorder, not a partial order. + a = optree.tree_structure(deque([1, 2, 3], maxlen=3)) + b = optree.tree_structure(deque([1, 2, 3], maxlen=5)) + unbounded = optree.tree_structure(deque([1, 2, 3])) # maxlen=None + # Equality distinguishes maxlen (bounded vs bounded, and bounded vs unbounded). + assert a != b + assert a != unbounded + assert b != unbounded + + # Same arity, any maxlen (bounded or unbounded): mutual non-strict prefixes and suffixes, + # even though the specs are unequal, so mutual prefixes do NOT imply equality (a preorder). + for x, y in itertools.permutations([a, b, unbounded], 2): + assert x != y + assert optree.treespec_is_prefix(x, y, strict=False) + assert optree.treespec_is_suffix(x, y, strict=False) + assert x <= y + assert x >= y + + # Practical consequence: a prefix deque flattens / broadcasts a full deque of a different maxlen. + prefix_spec = optree.tree_structure(deque([1, 2, 3], maxlen=None)) + assert prefix_spec.flatten_up_to(deque([[10], [20, 21], [30]], maxlen=5)) == [ + [10], + [20, 21], + [30], + ] + assert optree.broadcast_prefix( + deque([1, 2, 3], maxlen=3), + deque([[0], [0, 0], [0]], maxlen=7), + ) == [1, 2, 2, 3] + + # Arity still gates the relation: a different-arity deque is not a prefix. + assert not optree.treespec_is_prefix( + optree.tree_structure(deque([1, 2], maxlen=9)), + a, + strict=False, + ) + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -1420,6 +2456,21 @@ def test_treespec_leaf_none(namespace): ) +def test_treespec_from_collection_on_leaf_propagates_escalated_warning(): + # `treespec_from_collection()` on a leaf issues a UserWarning via `PyErr_WarnEx()`. When + # warnings are escalated to errors (e.g. `-W error`), the escalation must propagate cleanly as + # that UserWarning: the C++ code must check `PyErr_WarnEx()`'s return value and raise, not + # ignore it and return a result with the exception left set (which pybind11 surfaces as a + # confusing `SystemError: ... returned a result with an exception set`). + with warnings.catch_warnings(): + warnings.simplefilter('error') + with pytest.raises( + UserWarning, + match=re.escape('PyTreeSpec::MakeFromCollection() is called on a leaf.'), + ): + optree.treespec_from_collection(1) + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -1784,6 +2835,30 @@ def __tree_unflatten__(cls, metadata, children): assert treespec1 == treespec2 +def test_treespec_dict_constructor_preserves_insertion_ordered_namespace(): + # Regression: under `dict_insertion_ordered` mode the key order of a dict spec depends on the + # namespace, so `treespec_dict(..., namespace=...)` must keep that namespace (like + # `tree_flatten`) instead of resetting it to '': an empty-namespace spec with unsorted keys is + # otherwise unreachable via `tree_flatten` and breaks equality/consistency. + leaf = optree.tree_structure(0) + + with optree.dict_insertion_ordered(True, namespace='namespace'): + constructed = optree.treespec_dict({'b': leaf, 'a': leaf}, namespace='namespace') + _, flattened = optree.tree_flatten({'b': 1, 'a': 2}, namespace='namespace') + + assert constructed.entries() == ['b', 'a'] # insertion order preserved + assert flattened.namespace == 'namespace' + assert constructed.namespace == 'namespace' # was '' before the fix + assert constructed == flattened + + # Without the mode, keys are sorted and the namespace is dropped, same as `tree_flatten`. + outside = optree.treespec_dict({'b': leaf, 'a': leaf}, namespace='namespace') + _, flattened_outside = optree.tree_flatten({'b': 1, 'a': 2}, namespace='namespace') + assert outside.entries() == ['a', 'b'] + assert outside.namespace == '' + assert outside == flattened_outside + + def test_treespec_constructor_none_treespec_inputs(): with pytest.raises(ValueError, match=r'Expected a\(n\) list of PyTreeSpec\(s\), got .*\.'): optree.treespec_list([optree.treespec_leaf(), 1]) diff --git a/tests/test_typing.py b/tests/test_typing.py index 51cb95cc..d831300b 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -16,6 +16,7 @@ # pylint: disable=missing-function-docstring import enum +import os import re import sys import time @@ -28,14 +29,19 @@ import optree from helpers import ( PYBIND11_HAS_NATIVE_ENUM, + PYPY, CustomNamedTupleSubclass, CustomTuple, Py_GIL_DISABLED, Vector2D, + check_script_in_subprocess, disable_systrace, gc_collect, getrefcount, + skipif_android, + skipif_ios, skipif_pypy, + skipif_wasm, ) @@ -547,6 +553,27 @@ def test_structseq_fields(): 'tm_yday', 'tm_isdst', ) + # On CPython, `os.stat_result` has UNNAMED sequence slots 7, 8, 9 (the integer + # atime/mtime/ctime); the st_atime/st_mtime/st_ctime attributes are hidden FLOAT fields at + # higher field indices. `tp_members` must be mapped by offset, not by position, or those + # slots get mislabeled with the trailing hidden float names. + stat_fields = structseq_fields(os.stat_result) + assert len(stat_fields) == os.stat_result.n_sequence_fields + assert stat_fields[:7] == ( + 'st_mode', + 'st_ino', + 'st_dev', + 'st_nlink', + 'st_uid', + 'st_gid', + 'st_size', + ) + if not PYPY: + # PyPy has no unnamed fields: it names slots 7-9 `_integer_atime`/etc. and puts the + # hidden float `st_atime` at a later index, so this CPython-only check does not apply. + for name in stat_fields[7:10]: + assert name not in {'st_atime', 'st_mtime', 'st_ctime'} + assert not name.isidentifier() # the PyStructSequence unnamed-field marker with pytest.raises( TypeError, @@ -597,6 +624,36 @@ def test_structseq_fields(): structseq_fields(FakeStructSequence) +def test_structseq_accessor_unnamed_fields_codify_by_index(): + # The accessor round-trip (the generated code evaluates to the accessed value) must hold for + # every slot on every implementation. It exercises both codify styles: CPython leaves + # `os.stat_result` slots 7, 8, 9 UNNAMED, so their accessors codify to index access (matching the + # index-based `__call__`); PyPy names those slots (`_integer_atime` etc.) and codifies them by + # attribute. Either way `accessor.codify(...)` and `accessor(...)` resolve to the same `st[i]`. + st = os.stat(os.curdir) # a real stat_result, valid on both CPython and PyPy + accessors = optree.tree_accessors(st) + assert len(accessors) == os.stat_result.n_sequence_fields + for i, accessor in enumerate(accessors): + assert eval(accessor.codify('__st'), {'__st': st}, {}) == accessor(st) == st[i] + assert accessors[6].codify('__st') == '__st.st_size' # a named slot -> attribute access + + # Repeat with DISTINCT per-field values (`st[i] == i`) so the round-trip reliably catches the + # unnamed-slot mislabel: a real stat's whole-second atime could coincide with integer slot 7. On + # CPython slots 7, 8, 9 are UNNAMED, so their accessors must codify to index access; codifying slot + # 7 as `.st_atime` (the hidden FLOAT field CPython's own repr mislabels it with) would eval to + # that wrong value. PyPy names those slots (`_integer_atime` etc.) and aliases `st_atime` back to + # `self[7]`, so the unnamed-slot specifics below are asserted CPython-only. + st = os.stat_result(range(os.stat_result.n_fields)) + accessors = optree.tree_accessors(st) + assert len(accessors) == os.stat_result.n_sequence_fields + for i, accessor in enumerate(accessors): + assert eval(accessor.codify('__st'), {'__st': st}, {}) == accessor(st) == i + if not PYPY: + assert st.st_atime != st[7] # a different (hidden) field, not sequence slot 7 + for i in (7, 8, 9): + assert accessors[i].codify('__st') == f'__st[{i}]' + + @skipif_pypy @disable_systrace def test_structseq_fields_cache(): @@ -660,3 +717,38 @@ class Foo(metaclass=FooMeta): if not Py_GIL_DISABLED: assert called_with == 'Foo' assert wr() is None + + +@skipif_wasm +@skipif_android +@skipif_ios +@skipif_pypy # CPython-only: uses `atexit._ncallbacks()` and CPython type caches +def test_type_caches_register_interpreter_cleanup(): + # optree keeps three process-global type caches: namedtuple classification, PyStructSequence + # classification, and PyStructSequence field names. Each registers one per-interpreter `atexit` + # cleanup on its first insert (the classification caches at import via the registry, the + # field-name cache on first use). Measuring in a clean subprocess before importing optree pins + # optree's whole footprint: one callback for the registry plus one per cache. + check_script_in_subprocess( + r""" + import atexit + import time + + n0 = atexit._ncallbacks() + import optree + n1 = atexit._ncallbacks() + optree.is_namedtuple(int) + n2 = atexit._ncallbacks() + optree.is_structseq(int) + n3 = atexit._ncallbacks() + optree.structseq_fields(time.struct_time) + n4 = atexit._ncallbacks() + + assert n0 < n1, (n0, n1) + assert n1 <= n2, (n1, n2) + assert n2 <= n3, (n2, n3) + assert n3 <= n4, (n3, n4) + assert n4 - n0 == 4, (n0, n1, n2, n3, n4) + """, + output=None, + )