This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 0f8e4a42 [API] Consolidate structural mutation entrypoints and
borrowed fallback (#786)
0f8e4a42 is described below
commit 0f8e4a429b3278dde0f41fe0ca088f9db418473a
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Sep 14 12:26:46 2026 -0400
[API] Consolidate structural mutation entrypoints and borrowed fallback
(#786)
Consolidate structural mutation into `MutateExpected(AnyView,
InplaceMode)` and `DefaultMutateExpected(AnyView, InplaceMode)`, with a
matching throwing `Mutate` entrypoint. All three default to
`InplaceMode::kDisallow`. `kAllow` permits in-place dispatch only along
an owned path with a unique current value; default descent trusts the
callback's validated mode without rechecking uniqueness. Remove the
superseded entrypoints and migrate recursive callers while preserving
the raw FFI hooks, vtable layout, and copy-on-write behavior.
C++ mutation callbacks receive the explicit mode. Python exposes
`InplaceMode.DISALLOW` and `InplaceMode.ALLOW`; optional callback modes
use their corresponding integer values through the existing FFI enum
conversion.
Add `ValueOrUnchanged(const T&) &&` to copy a borrowed fallback only
when unchanged and move replacements otherwise.
## Migration
This intentionally changes the C++ convenience API. Consumers, including
Apache TVM, must migrate these calls together with their tvm-ffi
dependency update.
| Removed call | Replacement |
| --- | --- |
| `Mutate<T>` / `MutateExpected<T>` | Use `MutateExpected(value,
InplaceMode::kDisallow)` with checked typed unwrapping below, or
`Mutate` at a throwing boundary. |
| `MaybeInplaceMutateIfUniqueExpected`, `MaybeInplaceMutate<T>` /
`MaybeInplaceMutateExpected<T>` | Use `MutateExpected(value,
inplace_mode)`, or `Mutate(value, inplace_mode)` at a throwing boundary.
|
| `DefaultMaybeInplaceMutateExpected` | Use
`DefaultMutateExpected(value, inplace_mode)` for the current node's
default descent with its already-validated mode. |
Change a typed callback's optional third argument to `InplaceMode` and
forward it explicitly during recursion. `kAllow` requires permission
along the entire ownership path, established by an owned root or a valid
in-place hook; a unique borrowed child alone does not grant that
permission. Otherwise use `kDisallow`. `MutateExpected` checks the
current value's uniqueness before callback arguments acquire ownership.
Default descent bypasses the current callback and trusts the established
mode without rechecking uniqueness.
For `Expected` helpers and raw hooks,
`TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN` preserves the fixed typed-mutation
`TypeError` diagnostic. For example, with `Expr` as the required result
type:
```cpp
Expected<Expr> MutateExpr(const Expr& original, StructuralMutatorObj*
mutator,
InplaceMode inplace_mode) {
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
UnchangedOr<Expr>, result, mutator->MutateExpected(original,
inplace_mode));
return std::move(result).ValueOrUnchanged(original);
}
```
Call `.value()` only at a throwing boundary. An explicit
`.ValueOrUnchanged(...).cast<T>()` there is also possible when its
separate runtime cast diagnostic is acceptable; it is not a replacement
for checked narrowing in an `Expected` helper or raw hook.
---
include/tvm/ffi/extra/structural_mutate.h | 353 +++++++++++++-----------------
include/tvm/ffi/reflection/accessor.h | 8 +-
python/tvm_ffi/__init__.py | 2 +
python/tvm_ffi/structural.py | 37 +++-
src/ffi/extra/structural_mutate.cc | 49 +++--
tests/cpp/extra/test_big_int.cc | 4 +-
tests/cpp/extra/test_structural_mutate.cc | 231 +++++++++++--------
tests/cpp/testing_object.h | 4 +-
tests/python/test_structural.py | 51 +++--
9 files changed, 398 insertions(+), 341 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index 06384b27..62a741d7 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -49,6 +49,19 @@
namespace tvm {
namespace ffi {
+/*!
+ * \brief Permission to mutate a value in place along its ownership path.
+ *
+ * \note The numeric values 0 and 1 are part of the Python integer binding
contract,
+ * matching ``InplaceMode.DISALLOW`` and ``InplaceMode.ALLOW``
respectively.
+ */
+enum class InplaceMode : int32_t {
+ /*! \brief Preserve the source through copy-on-write mutation. */
+ kDisallow = 0,
+ /*! \brief Permit in-place mutation when the current value is uniquely
owned. */
+ kAllow = 1,
+};
+
class StructuralMutatorObj;
template <typename T>
class UnchangedOr;
@@ -298,6 +311,16 @@ class UnchangedOr {
:
details::AnyUnsafe::MoveFromAnyAfterCheck<T>(std::move(data_));
}
+ /*!
+ * \brief Move the replacement, or copy \p original when unchanged.
+ * \param original The borrowed original value, which is left unmodified.
+ * \return The replacement or original value.
+ */
+ TVM_FFI_INLINE T ValueOrUnchanged(const T& original) && {
+ return IsUnchanged() ? original
+ :
details::AnyUnsafe::MoveFromAnyAfterCheck<T>(std::move(data_));
+ }
+
/*!
* \brief Move the replacement, or move \p original when unchanged.
* \param original The owned original value.
@@ -374,153 +397,74 @@ class StructuralMutatorObj : public Object {
using MutatorObjType = StructuralMutatorObj;
/*!
- * \brief Mutate a value through the mutator vtable.
- *
- * \param value The value to mutate.
- * \tparam T The declared replacement type.
+ * \brief Throwing form of \ref MutateExpected.
+ * \param value The borrowed value to mutate.
+ * \param inplace_mode Whether the caller permits mutation along this
ownership path.
* \return The replacement or unchanged marker.
* \throws Error if mutation fails.
*
- * This entry point never intentionally mutates \p value in place. Recursive
mutations
- * also use \ref Mutate.
+ * \note The default InplaceMode::kDisallow uses copy-on-write. In-place
mutation is permitted
+ * only when inplace_mode is InplaceMode::kAllow and the current value
is uniquely owned.
+ * See \ref MutateExpected for the full permission and error semantics.
*
+ * Use \ref TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN in Expected-returning helpers
or raw hooks to
+ * preserve checked typed mutation results and the fixed mismatch
diagnostic. At a throwing
+ * boundary, an explicit cast may instead report its own TypeError:
* \code{.cpp}
- * Expr new_node =
mutator->Mutate<Expr>(node).ValueOrUnchanged(std::move(node));
+ * Expr new_node =
mutator->Mutate(node).ValueOrUnchanged(AnyView(node)).cast<Expr>();
* \endcode
*/
- template <typename T = Any>
- TVM_FFI_INLINE UnchangedOr<T> Mutate(AnyView value) {
- return std::move(MutateExpected<T>(value)).value();
- }
-
- /*!
- * \brief Exception-free form of \ref Mutate.
- *
- * \param value The value to mutate.
- * \tparam T The declared replacement type.
- * \return The replacement or unchanged marker, or an Error if mutation
failed.
- */
- template <typename T = Any>
- TVM_FFI_INLINE Expected<UnchangedOr<T>> MutateExpected(AnyView value)
noexcept {
- if constexpr (std::is_same_v<T, Any>) {
- return details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<T>>(
- (*vtable_->mutate)(this, value));
- } else {
- TVMFFIAny result = (*vtable_->mutate)(this, value);
- if
(TVM_FFI_PREDICT_FALSE(!TypeTraits<Expected<UnchangedOr<T>>>::CheckAnyStrict(&result)))
{
- (void)details::AnyUnsafe::MoveTVMFFIAnyRawToAny(result);
- return details::SMutateDeclaredTypeError();
- }
- return
details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<T>>(result);
- }
+ TVM_FFI_INLINE UnchangedOr<Any> Mutate(AnyView value,
+ InplaceMode inplace_mode =
InplaceMode::kDisallow) {
+ return std::move(MutateExpected(value, inplace_mode)).value();
}
/*!
- * \brief Mutate a value, permitting an in-place implementation when it is
safe.
- *
+ * \brief Mutate a value, permitting in-place mutation only when uniquely
owned.
* \param value The borrowed value to mutate.
- * \tparam T The declared replacement type.
- * \return The replacement or unchanged marker.
- * \throws Error if mutation fails.
- *
- * The returned value may refer to the same object as \p value. Callers must
use the return value
- * as the result of the mutation rather than assuming that the input object
was reused.
- */
- template <typename T = Any>
- TVM_FFI_INLINE UnchangedOr<T> MaybeInplaceMutate(AnyView value) {
- return std::move(MaybeInplaceMutateExpected<T>(value)).value();
- }
-
- /*!
- * \brief Exception-free form of \ref MaybeInplaceMutate.
- *
- * \param value The borrowed value to mutate.
- * \tparam T The declared replacement type.
+ * \param inplace_mode Whether the caller permits mutation along this
ownership path.
* \return The replacement or unchanged marker, or an Error if mutation
failed.
*
- * \note Call only from a ``__s_maybe_inplace_mutate__`` hook, which is
dispatched
- * only for a value whose entire path from the root is uniquely owned.
+ * \note The default InplaceMode::kDisallow uses copy-on-write. In-place
mutation is permitted
+ * only when inplace_mode is InplaceMode::kAllow and the current value
is uniquely owned.
+ * InplaceMode::kAllow requires permission along the entire path from
the root.
+ * Recursive calls must forward their established mode explicitly.
Uniqueness is checked
+ * before callback arguments acquire ownership. This permits in-place
dispatch but does
+ * not guarantee reuse: a hook may return a replacement. In-place
changes completed before
+ * an Error are not rolled back.
*/
- template <typename T = Any>
- TVM_FFI_INLINE Expected<UnchangedOr<T>> MaybeInplaceMutateExpected(AnyView
value) noexcept {
- if constexpr (std::is_same_v<T, Any>) {
- return details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<T>>(
+ TVM_FFI_INLINE Expected<UnchangedOr<Any>> MutateExpected(
+ AnyView value, InplaceMode inplace_mode = InplaceMode::kDisallow)
noexcept {
+ const Object* object = value.as<Object>();
+ // Check uniqueness on the borrowed view before callbacks can acquire
owning references.
+ if (inplace_mode == InplaceMode::kAllow && object != nullptr &&
object->unique()) {
+ return details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<Any>>(
(*vtable_->maybe_inplace_mutate)(this, value));
- } else {
- TVMFFIAny result = (*vtable_->maybe_inplace_mutate)(this, value);
- if
(TVM_FFI_PREDICT_FALSE(!TypeTraits<Expected<UnchangedOr<T>>>::CheckAnyStrict(&result)))
{
- (void)details::AnyUnsafe::MoveTVMFFIAnyRawToAny(result);
- return details::SMutateDeclaredTypeError();
- }
- return
details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<T>>(result);
}
+ return details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<Any>>(
+ (*vtable_->mutate)(this, value));
}
/*!
- * \brief Mutate a value, using in-place mutation only for a uniquely owned
object
- * and \p allow_inplace set to true.
- * \tparam T The declared replacement type.
- * \param value The borrowed value to mutate.
- * \param allow_inplace Whether in-place mutation is permitted. Defaults to
true.
- * If false, use ordinary mutation without checking uniqueness.
+ * \brief Apply default structural mutation with validated in-place
permission.
+ * \param value The borrowed current value to mutate.
+ * \param inplace_mode The in-place mode already established by the caller
for value.
* \return The replacement or unchanged marker, or an Error if mutation
failed.
*
- * \note When \p allow_inplace is true, the caller must already know the
entire
- * path from the root is uniquely owned, either through an owning
moved-in
- * root or while handling a ``__s_maybe_inplace_mutate__`` hook. This
method
- * checks only \p value itself, not its ancestors.
+ * \note The default InplaceMode::kDisallow is a copy-on-write convenience
for one-off calls.
+ * Recursive code and hooks must explicitly forward the mode
established for the current
+ * value. This method uses that mode without checking uniqueness
again, even if a typed
+ * callback argument has acquired another reference. It bypasses the
current engine
+ * callback. Permission must cover the entire path from the root.
+ * Without an in-place hook, ordinary mutation runs. In-place changes
completed before an
+ * Error are not rolled back. Registered hooks own variable-remap
handling; the reflected
+ * fallback applies it automatically and always uses copy-on-write
mutation.
*/
- template <typename T = Any>
- TVM_FFI_INLINE Expected<UnchangedOr<T>> MaybeInplaceMutateIfUniqueExpected(
- AnyView value, bool allow_inplace = true) noexcept {
- const Object* obj = value.as<Object>();
- if (allow_inplace && obj != nullptr && obj->unique()) {
- return MaybeInplaceMutateExpected<T>(value);
- }
- return MutateExpected<T>(value);
- }
-
- /*!
- * \brief Apply the default structural mutation with copy-on-write behavior.
- *
- * \param value The value to mutate.
- * \return The replacement or unchanged marker, or an Error if hook
dispatch, copying, or field
- * mutation failed.
- *
- * \note A registered ``__s_mutate__`` hook is dispatched before the
reflected fallback. A
- * FreeVar hook owns the definition-only remap policy for that type;
the reflected fallback
- * applies the same policy automatically.
- */
-
- TVM_FFI_INLINE Expected<UnchangedOr<Any>> DefaultMutateExpected(AnyView
value) noexcept {
- return
details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<Any>>(DefaultMutateRaw(value));
- }
-
- /*!
- * \brief Apply default structural mutation with optional in-place
permission.
- *
- * \param value The borrowed value to mutate.
- * \param allow_inplace Whether in-place mutation is permitted. Defaults to
true.
- * If false, call \ref DefaultMutateExpected.
- * \return The replacement or unchanged marker, or an Error if mutation
failed. In-place
- * changes completed before an Error are not rolled back.
- *
- * \note In-place mutation is explicitly opt-in. A registered
- * ``__s_maybe_inplace_mutate__`` hook may rely on its input being
safe to mutate and owns
- * any variable-remap handling. When the hook is absent, this method
calls
- * \ref DefaultMutateExpected. This method does not check uniqueness.
When
- * \p allow_inplace is true, the caller must already know that mutating
- * \p value in place is safe, including ownership of the path from the
root.
- *
- * \code
- * return mutator->DefaultMaybeInplaceMutateExpected(value, allow_inplace);
- * \endcode
- */
- TVM_FFI_INLINE Expected<UnchangedOr<Any>> DefaultMaybeInplaceMutateExpected(
- AnyView value, bool allow_inplace = true) noexcept {
- if (!allow_inplace) return DefaultMutateExpected(value);
+ TVM_FFI_INLINE Expected<UnchangedOr<Any>> DefaultMutateExpected(
+ AnyView value, InplaceMode inplace_mode = InplaceMode::kDisallow)
noexcept {
return details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<Any>>(
- DefaultMaybeInplaceMutateRaw(value));
+ inplace_mode == InplaceMode::kAllow ?
DefaultMaybeInplaceMutateRaw(value)
+ : DefaultMutateRaw(value));
}
/*!
@@ -744,7 +688,8 @@ class StructuralMutatorObj : public Object {
}
return result;
}
- return
details::ExpectedUnsafe::MoveToTVMFFIAny(DefaultMutateExpected(value));
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ DefaultMutateExpected(value, InplaceMode::kDisallow));
}
protected:
@@ -840,14 +785,14 @@ TVM_FFI_INLINE static Expected<Any>
MutateReflectedFieldsExpected(StructuralMuta
Expected<UnchangedOr<Any>> mutated_field = [&]() ->
Expected<UnchangedOr<Any>> {
if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashDefSimple) {
return mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple,
[&]() {
- return mutator->MutateExpected(field_value);
+ return mutator->MutateExpected(field_value,
InplaceMode::kDisallow);
});
} else if (field_info->flags &
kTVMFFIFieldFlagBitMaskSEqHashDefPattern) {
return mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern,
[&]() {
- return mutator->MutateExpected(field_value);
+ return mutator->MutateExpected(field_value,
InplaceMode::kDisallow);
});
} else {
- return mutator->MutateExpected(field_value);
+ return mutator->MutateExpected(field_value,
InplaceMode::kDisallow);
}
}();
if (TVM_FFI_PREDICT_FALSE(mutated_field.is_err())) {
@@ -945,19 +890,20 @@ namespace details {
* returns ``TypeError`` through the surrounding raw or ``Expected`` function
without throwing,
* reported with a fixed string. The check is omitted when the declared type
subsumes the result's
* success type. Its early returns work from either a raw ``TVMFFIAny`` hook
or an
- * ``Expected<T>`` helper, including one with a different success type. This
macro declares ``Name``
- * into the enclosing scope and must be used in a braced block, never as an
unbraced control-flow
- * body.
+ * ``Expected<T>`` helper, including one with a different success type. This
macro declares
+ * ``Name`` into the enclosing scope and must be used in a braced block, never
as an unbraced
+ * control-flow body.
*
* Example:
* \code{.cpp}
* TVMFFIAny FooMutate(StructuralMutatorObj* mutator, AnyView value) noexcept {
* const FooNode* self =
* details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FooNode>(value);
+ * constexpr InplaceMode inplace_mode = InplaceMode::kDisallow;
* TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Expr>, a,
- * mutator->MutateExpected(self->a));
+ * mutator->MutateExpected(self->a,
inplace_mode));
* TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Expr>, b,
- * mutator->MutateExpected(self->b));
+ * mutator->MutateExpected(self->b,
inplace_mode));
* if (a.UnchangedOrSameAs(self->a) && b.UnchangedOrSameAs(self->b)) {
* return Unchanged().CopyToTVMFFIAny();
* }
@@ -1276,14 +1222,14 @@ class StructuralMapEngine : public Parent {
/*!
* \brief Test one link against \p value and, if it matches, mutate the node
through it.
*
- * \tparam kMaybeInplace Whether the caller may mutate a uniquely owned node
in place.
+ * \tparam kInplaceMode Whether the caller may mutate a uniquely owned node
in place.
* \tparam Callback The link's callback type.
* \param callback The link's callback.
* \param value The borrowed value to test and mutate.
* \param out Receives the mutated value or Error when the link matched.
* \return Whether the link matched, in which case \p out was written.
*/
- template <bool kMaybeInplace, typename Callback>
+ template <InplaceMode kInplaceMode, typename Callback>
TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value,
Expected<Any>* out) noexcept {
using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
static_assert(FuncInfo::num_args >= 1,
@@ -1320,17 +1266,18 @@ class StructuralMapEngine : public Parent {
mapped_value.type_index() == TypeIndex::kTVMFFIUnchanged ? value :
AnyView(mapped_value);
// Each descent names the node it actually ran on in the error context.
*out = [&]() -> Expected<Any> {
- if constexpr (kMaybeInplace) {
+ if constexpr (kInplaceMode == InplaceMode::kAllow) {
// A pre-order result can be mutated in place if unchanged or
uniquely owned.
if (descent_view.same_as(value)) {
- return this->DefaultMaybeInplaceMutateExpected(value);
+ return this->DefaultMutateExpected(value, InplaceMode::kAllow);
}
const Object* mapped_obj = descent_view.as<Object>();
- bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
- return can_inplace ?
this->DefaultMaybeInplaceMutateExpected(descent_view)
- : this->DefaultMutateExpected(descent_view);
+ InplaceMode inplace_mode = mapped_obj != nullptr &&
mapped_obj->unique()
+ ? InplaceMode::kAllow
+ : InplaceMode::kDisallow;
+ return this->DefaultMutateExpected(descent_view, inplace_mode);
} else {
- return this->DefaultMutateExpected(descent_view);
+ return this->DefaultMutateExpected(descent_view,
InplaceMode::kDisallow);
}
}();
if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
@@ -1371,10 +1318,10 @@ class StructuralMapEngine : public Parent {
* \brief Test every link in declaration order, stopping at the first that
matches.
* \return Whether some link matched, in which case \p out was written.
*/
- template <bool kMaybeInplace, size_t... Is>
+ template <InplaceMode kInplaceMode, size_t... Is>
TVM_FFI_INLINE bool TryLinks(AnyView value, Expected<Any>* out,
std::index_sequence<Is...>) noexcept {
- return (TryLink<kMaybeInplace>(std::get<Is>(callbacks_), value, out) ||
...);
+ return (TryLink<kInplaceMode>(std::get<Is>(callbacks_), value, out) ||
...);
}
/*!
@@ -1385,17 +1332,18 @@ class StructuralMapEngine : public Parent {
TVM_FFI_INLINE TVMFFIAny MutateImplRaw(AnyView value) noexcept {
Expected<Any> out{Any()};
if constexpr (order == WalkOrder::kPostOrder) {
- out = this->DefaultMutateExpected(value);
+ out = this->DefaultMutateExpected(value, InplaceMode::kDisallow);
if (TVM_FFI_PREDICT_FALSE(out.is_err())) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- TryLinks<false>(value, &out, std::index_sequence_for<Callbacks...>{});
+ TryLinks<InplaceMode::kDisallow>(value, &out,
std::index_sequence_for<Callbacks...>{});
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
} else {
- if (TryLinks<false>(value, &out,
std::index_sequence_for<Callbacks...>{})) {
+ if (TryLinks<InplaceMode::kDisallow>(value, &out,
std::index_sequence_for<Callbacks...>{})) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return
ExpectedUnsafe::MoveToTVMFFIAny(this->DefaultMutateExpected(value));
+ return ExpectedUnsafe::MoveToTVMFFIAny(
+ this->DefaultMutateExpected(value, InplaceMode::kDisallow));
}
}
@@ -1407,17 +1355,18 @@ class StructuralMapEngine : public Parent {
TVM_FFI_INLINE TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
Expected<Any> out{Any()};
if constexpr (order == WalkOrder::kPostOrder) {
- out = this->DefaultMaybeInplaceMutateExpected(value);
+ out = this->DefaultMutateExpected(value, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(out.is_err())) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- TryLinks<true>(value, &out, std::index_sequence_for<Callbacks...>{});
+ TryLinks<InplaceMode::kAllow>(value, &out,
std::index_sequence_for<Callbacks...>{});
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
} else {
- if (TryLinks<true>(value, &out,
std::index_sequence_for<Callbacks...>{})) {
+ if (TryLinks<InplaceMode::kAllow>(value, &out,
std::index_sequence_for<Callbacks...>{})) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return
ExpectedUnsafe::MoveToTVMFFIAny(this->DefaultMaybeInplaceMutateExpected(value));
+ return ExpectedUnsafe::MoveToTVMFFIAny(
+ this->DefaultMutateExpected(value, InplaceMode::kAllow));
}
}
@@ -1515,12 +1464,12 @@ class StructuralMapDynEngine : public Parent {
/*!
* \brief Test the runtime link table against \p value and mutate through
the first match.
- * \tparam kMaybeInplace Whether a uniquely owned node may be mutated in
place.
+ * \tparam kInplaceMode Whether a uniquely owned node may be mutated in
place.
* \param value The borrowed value to test and mutate.
* \param out Receives the mutated value or Error when a link matched.
* \return Whether a link matched, in which case \p out was written.
*/
- template <bool kMaybeInplace>
+ template <InplaceMode kInplaceMode>
TVM_FFI_INLINE bool TryLink(AnyView value, Expected<Any>* out) noexcept {
if constexpr (order == WalkOrder::kPreOrder) {
bool with_kind = false;
@@ -1538,16 +1487,17 @@ class StructuralMapDynEngine : public Parent {
const AnyView descent_view =
mapped_value.type_index() == TypeIndex::kTVMFFIUnchanged ? value :
AnyView(mapped_value);
*out = [&]() -> Expected<Any> {
- if constexpr (kMaybeInplace) {
+ if constexpr (kInplaceMode == InplaceMode::kAllow) {
if (descent_view.same_as(value)) {
- return this->DefaultMaybeInplaceMutateExpected(value);
+ return this->DefaultMutateExpected(value, InplaceMode::kAllow);
}
const Object* mapped_obj = descent_view.as<Object>();
- bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
- return can_inplace ?
this->DefaultMaybeInplaceMutateExpected(descent_view)
- : this->DefaultMutateExpected(descent_view);
+ InplaceMode inplace_mode = mapped_obj != nullptr &&
mapped_obj->unique()
+ ? InplaceMode::kAllow
+ : InplaceMode::kDisallow;
+ return this->DefaultMutateExpected(descent_view, inplace_mode);
} else {
- return this->DefaultMutateExpected(descent_view);
+ return this->DefaultMutateExpected(descent_view,
InplaceMode::kDisallow);
}
}();
if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
@@ -1580,17 +1530,18 @@ class StructuralMapDynEngine : public Parent {
TVM_FFI_INLINE TVMFFIAny MutateImplRaw(AnyView value) noexcept {
Expected<Any> out{Any()};
if constexpr (order == WalkOrder::kPostOrder) {
- out = this->DefaultMutateExpected(value);
+ out = this->DefaultMutateExpected(value, InplaceMode::kDisallow);
if (TVM_FFI_PREDICT_FALSE(out.is_err())) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- TryLink<false>(value, &out);
+ TryLink<InplaceMode::kDisallow>(value, &out);
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
} else {
- if (TryLink<false>(value, &out)) {
+ if (TryLink<InplaceMode::kDisallow>(value, &out)) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return
ExpectedUnsafe::MoveToTVMFFIAny(this->DefaultMutateExpected(value));
+ return ExpectedUnsafe::MoveToTVMFFIAny(
+ this->DefaultMutateExpected(value, InplaceMode::kDisallow));
}
}
@@ -1598,17 +1549,18 @@ class StructuralMapDynEngine : public Parent {
TVM_FFI_INLINE TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
Expected<Any> out{Any()};
if constexpr (order == WalkOrder::kPostOrder) {
- out = this->DefaultMaybeInplaceMutateExpected(value);
+ out = this->DefaultMutateExpected(value, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(out.is_err())) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- TryLink<true>(value, &out);
+ TryLink<InplaceMode::kAllow>(value, &out);
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
} else {
- if (TryLink<true>(value, &out)) {
+ if (TryLink<InplaceMode::kAllow>(value, &out)) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return
ExpectedUnsafe::MoveToTVMFFIAny(this->DefaultMaybeInplaceMutateExpected(value));
+ return ExpectedUnsafe::MoveToTVMFFIAny(
+ this->DefaultMutateExpected(value, InplaceMode::kAllow));
}
}
@@ -1655,7 +1607,7 @@ class StructuralMutateEngine : public Parent {
static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
auto* self = static_cast<StructuralMutateEngine*>(mutator);
if constexpr (sizeof...(Callbacks) == 1) {
- return self->template MutateSingleCallbackRaw<false>(value);
+ return self->template
MutateSingleCallbackRaw<InplaceMode::kDisallow>(value);
} else {
return self->MutateImplRaw(value);
}
@@ -1666,7 +1618,7 @@ class StructuralMutateEngine : public Parent {
AnyView value) noexcept {
auto* self = static_cast<StructuralMutateEngine*>(mutator);
if constexpr (sizeof...(Callbacks) == 1) {
- return self->template MutateSingleCallbackRaw<true>(value);
+ return self->template
MutateSingleCallbackRaw<InplaceMode::kAllow>(value);
} else {
return self->MaybeInplaceMutateImplRaw(value);
}
@@ -1675,7 +1627,7 @@ class StructuralMutateEngine : public Parent {
/*! \brief Mutate one value, handing a matched callback ownership of
descent. */
TVMFFIAny MutateImplRaw(AnyView value) noexcept {
Expected<Any> result{Any()};
- if (DispatchCallbacks(value, false, &result)) {
+ if (DispatchCallbacks(value, InplaceMode::kDisallow, &result)) {
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
// Keep callback-boundary context in addition to the default-descent
// context: a callback may return a rebuilt value, so the two nodes
can differ.
@@ -1683,13 +1635,14 @@ class StructuralMutateEngine : public Parent {
}
return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
- return
details::ExpectedUnsafe::MoveToTVMFFIAny(Parent::DefaultMutateExpected(value));
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ Parent::DefaultMutateExpected(value, InplaceMode::kDisallow));
}
/*! \brief Maybe mutate one value in place, with callback-owned descent. */
TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
Expected<Any> result{Any()};
- if (DispatchCallbacks(value, true, &result)) {
+ if (DispatchCallbacks(value, InplaceMode::kAllow, &result)) {
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
// Keep callback-boundary context in addition to the default-descent
// context: a callback may return a rebuilt value, so the two nodes
can differ.
@@ -1698,17 +1651,17 @@ class StructuralMutateEngine : public Parent {
return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
return details::ExpectedUnsafe::MoveToTVMFFIAny(
- Parent::DefaultMaybeInplaceMutateExpected(value));
+ Parent::DefaultMutateExpected(value, InplaceMode::kAllow));
}
/*! \brief Invoke a matched callback with the Parent view and preserve
returned/thrown Error. */
template <typename Callback, typename Matched>
TVM_FFI_INLINE Expected<Any> InvokeCallback(Callback& callback, Matched&&
matched,
- bool allow_inplace) noexcept {
+ InplaceMode inplace_mode)
noexcept {
using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
static_assert(FuncInfo::num_args == 2 || FuncInfo::num_args == 3,
"StructuralMutate callback must take (value, mutator) or "
- "(value, mutator, allow_inplace)");
+ "(value, mutator, inplace_mode)");
using SecondArg = std::decay_t<std::tuple_element_t<1, typename
FuncInfo::ArgType>>;
using Second = std::remove_pointer_t<SecondArg>;
static_assert(std::is_same_v<Second, typename Parent::MutatorObjType>,
@@ -1716,13 +1669,13 @@ class StructuralMutateEngine : public Parent {
"Parent::MutatorObjType*");
if constexpr (FuncInfo::num_args == 3) {
using ThirdArg = std::decay_t<std::tuple_element_t<2, typename
FuncInfo::ArgType>>;
- static_assert(std::is_same_v<ThirdArg, bool>,
- "third StructuralMutate callback argument must be bool");
+ static_assert(std::is_same_v<ThirdArg, InplaceMode>,
+ "third StructuralMutate callback argument must be
InplaceMode");
}
auto* mutator = static_cast<typename Parent::MutatorObjType*>(this);
try {
if constexpr (FuncInfo::num_args == 3) {
- return callback(std::forward<Matched>(matched), mutator,
allow_inplace);
+ return callback(std::forward<Matched>(matched), mutator, inplace_mode);
} else {
return callback(std::forward<Matched>(matched), mutator);
}
@@ -1740,7 +1693,7 @@ class StructuralMutateEngine : public Parent {
// One callback forwards its result directly, without a callback-chain
envelope.
// AnyView/Any always match; only a typed miss needs the Parent's default
descent.
- template <bool kMaybeInplace>
+ template <InplaceMode kInplaceMode>
TVM_FFI_INLINE TVMFFIAny MutateSingleCallbackRaw(AnyView value) noexcept {
auto& callback = std::get<0>(callbacks_);
using FuncInfo = details::FunctionInfo<std::decay_t<decltype(callback)>>;
@@ -1749,19 +1702,20 @@ class StructuralMutateEngine : public Parent {
TVMFFIAny result;
if constexpr (std::is_same_v<TSub, AnyView>) {
result =
- details::ExpectedUnsafe::MoveToTVMFFIAny(InvokeCallback(callback,
value, kMaybeInplace));
+ details::ExpectedUnsafe::MoveToTVMFFIAny(InvokeCallback(callback,
value, kInplaceMode));
} else if constexpr (std::is_same_v<TSub, Any>) {
result = details::ExpectedUnsafe::MoveToTVMFFIAny(
- InvokeCallback(callback, Any(value), kMaybeInplace));
+ InvokeCallback(callback, Any(value), kInplaceMode));
} else if (auto matched = value.template as<TSub>()) {
result = details::ExpectedUnsafe::MoveToTVMFFIAny(
- InvokeCallback(callback, *std::move(matched), kMaybeInplace));
+ InvokeCallback(callback, *std::move(matched), kInplaceMode));
} else {
- if constexpr (kMaybeInplace) {
+ if constexpr (kInplaceMode == InplaceMode::kAllow) {
return details::ExpectedUnsafe::MoveToTVMFFIAny(
- Parent::DefaultMaybeInplaceMutateExpected(value));
+ Parent::DefaultMutateExpected(value, InplaceMode::kAllow));
} else {
- return
details::ExpectedUnsafe::MoveToTVMFFIAny(Parent::DefaultMutateExpected(value));
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ Parent::DefaultMutateExpected(value, InplaceMode::kDisallow));
}
}
// Release any owning match before naming the callback boundary, as in the
general path.
@@ -1773,19 +1727,19 @@ class StructuralMutateEngine : public Parent {
/*! \brief Write a matched callback result, leaving out untouched on a type
miss. */
template <typename Callback>
- TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value, bool
allow_inplace,
+ TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value, InplaceMode
inplace_mode,
Expected<Any>* out) noexcept {
using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
if constexpr (std::is_same_v<TSub, AnyView>) {
- *out = InvokeCallback(callback, value, allow_inplace);
+ *out = InvokeCallback(callback, value, inplace_mode);
return true;
} else if constexpr (std::is_same_v<TSub, Any>) {
- *out = InvokeCallback(callback, Any(value), allow_inplace);
+ *out = InvokeCallback(callback, Any(value), inplace_mode);
return true;
} else if (auto matched = value.template as<TSub>()) {
- *out = InvokeCallback(callback, *std::move(matched), allow_inplace);
+ *out = InvokeCallback(callback, *std::move(matched), inplace_mode);
return true;
}
return false;
@@ -1793,15 +1747,15 @@ class StructuralMutateEngine : public Parent {
/*! \brief Fold callbacks in declaration order, stopping at the first match.
*/
template <size_t... Is>
- TVM_FFI_INLINE bool TryLinks(AnyView value, bool allow_inplace,
Expected<Any>* out,
+ TVM_FFI_INLINE bool TryLinks(AnyView value, InplaceMode inplace_mode,
Expected<Any>* out,
std::index_sequence<Is...>) noexcept {
- return (TryLink(std::get<Is>(callbacks_), value, allow_inplace, out) ||
...);
+ return (TryLink(std::get<Is>(callbacks_), value, inplace_mode, out) ||
...);
}
/*! \brief Run the callback chain, returning whether a callback matched. */
- TVM_FFI_INLINE bool DispatchCallbacks(AnyView value, bool allow_inplace,
+ TVM_FFI_INLINE bool DispatchCallbacks(AnyView value, InplaceMode
inplace_mode,
Expected<Any>* out) noexcept {
- return TryLinks(value, allow_inplace, out,
std::index_sequence_for<Callbacks...>{});
+ return TryLinks(value, inplace_mode, out,
std::index_sequence_for<Callbacks...>{});
}
/*! \brief Typed callbacks tested in declaration order, first match wins. */
@@ -1848,7 +1802,7 @@ Expected<Any> StructuralMapExpected(
static_assert(sizeof...(Callbacks) != 0, "StructuralMap requires at least
one callback");
using Mutator = StructuralMapEngine<StructuralMapEngineBase, order,
std::decay_t<Callbacks>...>;
StructuralMutator
mutator(make_object<Mutator>(std::forward<Callbacks>(callbacks)...));
- auto result = mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ auto result = mutator->MutateExpected(root, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) return
Unexpected(std::move(result).error());
UnchangedOr<Any> mapped =
details::AnyUnsafe::MoveFromAnyAfterCheck<UnchangedOr<Any>>(
std::move(details::ExpectedUnsafe::GetData(result)));
@@ -1888,7 +1842,7 @@ Any StructuralMap(Any root,
* A callback takes one of two forms, where ``R`` is a supported return type:
*
* - ``R(const T& value, StructuralMutatorObj* mutator)``
- * - ``R(const T& value, StructuralMutatorObj* mutator, bool allow_inplace)``
+ * - ``R(const T& value, StructuralMutatorObj* mutator, InplaceMode
inplace_mode)``
*
* A callback returns a replacement, ``Unchanged``, or ``Expected<Any>``. The
first argument
* selects by FFI type; callbacks are tried in declaration order and the first
match owns mutation
@@ -1899,9 +1853,10 @@ Any StructuralMap(Any root,
* \param callbacks Callbacks tested in declaration order.
* \return The mutated owning value, or an Error if mutation or a callback
fails.
*
- * \note A two-argument callback descends with ``MutateExpected`` and remains
copy-on-write.
- * A three-argument callback receives ``allow_inplace=true`` only when
its value is on a
- * uniquely owned path and may then explicitly use the maybe-in-place
mutator operation.
+ * \note A two-argument callback descends with ``MutateExpected(value)``,
using the default
+ * copy-on-write permission. A three-argument callback receives
InplaceMode::kAllow only
+ * when its value is on a uniquely owned path. Forward this mode to
``MutateExpected`` for
+ * child values or ``DefaultMutateExpected`` for the current value's
default descent.
* \note Pass an owned root with ``std::move(root)`` to permit root reuse. In a
* ``__s_maybe_inplace_mutate__`` hook, the corresponding nested idiom is
* ``self->field = StructuralMap(std::move(self->field), callback)``;
const-correctness
@@ -1914,7 +1869,7 @@ Expected<Any> StructuralMutateExpected(
static_assert(sizeof...(Callbacks) != 0, "StructuralMutate requires at least
one callback");
using Mutator = StructuralMutateEngine<StructuralMapEngineBase,
std::decay_t<Callbacks>...>;
StructuralMutator
mutator(make_object<Mutator>(std::forward<Callbacks>(callbacks)...));
- auto result = mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ auto result = mutator->MutateExpected(root, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) return
Unexpected(std::move(result).error());
UnchangedOr<Any> mapped =
details::AnyUnsafe::MoveFromAnyAfterCheck<UnchangedOr<Any>>(
std::move(details::ExpectedUnsafe::GetData(result)));
diff --git a/include/tvm/ffi/reflection/accessor.h
b/include/tvm/ffi/reflection/accessor.h
index c0e3d93f..94249176 100644
--- a/include/tvm/ffi/reflection/accessor.h
+++ b/include/tvm/ffi/reflection/accessor.h
@@ -534,10 +534,10 @@ inline constexpr const char* kStructuralMutate =
"__s_mutate__";
*
* ``(StructuralMutator mutator, Any value) -> Any``.
*
- * This hook is optional. When it is absent,
``DefaultMaybeInplaceMutateExpected`` falls back to
- * non-in-place mutation through ``kStructuralMutate`` or reflected structural
fields. In-place
- * mutation is therefore explicitly opt-in and is never inferred from
ownership by the reflected
- * fallback.
+ * This hook is optional. When it is absent, ``DefaultMutateExpected`` with
+ * ``inplace_mode=InplaceMode::kAllow`` falls back to non-in-place mutation
through
+ * ``kStructuralMutate`` or reflected structural fields. In-place mutation is
therefore explicitly
+ * opt-in and is never inferred from ownership by the reflected fallback.
*/
inline constexpr const char* kStructuralMaybeInplaceMutate =
"__s_maybe_inplace_mutate__";
diff --git a/python/tvm_ffi/__init__.py b/python/tvm_ffi/__init__.py
index f3ee8c86..7eafc132 100644
--- a/python/tvm_ffi/__init__.py
+++ b/python/tvm_ffi/__init__.py
@@ -77,6 +77,7 @@ if TYPE_CHECKING or not _is_config_mode():
from .stream import StreamContext, get_raw_stream, use_raw_stream,
use_torch_stream
from .structural import (
DefRegionKind,
+ InplaceMode,
StructuralKey,
StructuralMutator,
StructuralVisitor,
@@ -150,6 +151,7 @@ __all__ = [
"Device",
"Dict",
"Function",
+ "InplaceMode",
"List",
"Map",
"Module",
diff --git a/python/tvm_ffi/structural.py b/python/tvm_ffi/structural.py
index d1f369e8..fc10596a 100644
--- a/python/tvm_ffi/structural.py
+++ b/python/tvm_ffi/structural.py
@@ -33,6 +33,7 @@ from .registry import register_object
__all__ = [
"DefRegionKind",
+ "InplaceMode",
"StructuralKey",
"StructuralMutator",
"StructuralVisitor",
@@ -84,6 +85,25 @@ class WalkResult(IntEnum):
SKIP = 1
+class InplaceMode(IntEnum):
+ """Permission to mutate a value in place along its ownership path.
+
+ DISALLOW preserves the source through copy-on-write. ALLOW permits in-place
+ mutation only when the current value is uniquely owned; it does not require
+ a callback to reuse the value. Structural mutation callbacks receive the
+ corresponding integer, which compares directly with these enum members.
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.structural_mutate`
+ Mutate a value with callbacks that own recursive mutation.
+
+ """
+
+ DISALLOW = 0
+ ALLOW = 1
+
+
class DefRegionKind(IntEnum):
"""Def-region state active during structural visiting.
@@ -697,9 +717,10 @@ def structural_mutate(
"""Mutate a value with callbacks that own recursive mutation.
Each callback receives ``(value, mutator)`` and may optionally receive a
- third ``allow_inplace`` boolean, which is true only when the callback's
- value is on a uniquely owned path. The flag lets a callback choose an
- ownership-aware implementation; recursive descent still uses
+ third ``inplace_mode`` integer corresponding to :class:`InplaceMode`.
+ It equals ``InplaceMode.ALLOW`` only when the callback's value is on a
+ uniquely owned path, and ``InplaceMode.DISALLOW`` otherwise. The mode lets
+ a callback choose an ownership-aware implementation; recursive descent uses
:meth:`StructuralMutator.mutate` for selected children or
:meth:`StructuralMutator.default_mutate` for the matched value's default
mutation. Its returned value is final and is not traversed again. Entries
@@ -726,7 +747,7 @@ def structural_mutate(
(
_callback_type_to_type_index(t, api_name="structural_mutate"),
fn,
- _callback_accepts_allow_inplace(fn),
+ _callback_accepts_inplace_mode(fn),
)
for t, fn in callback_entries
]
@@ -860,8 +881,8 @@ def _normalize_callbacks(
return callback_entries
-def _callback_accepts_allow_inplace(callback: Callable[..., Any]) -> bool:
- """Return whether a StructuralMutate callback accepts its optional flag."""
+def _callback_accepts_inplace_mode(callback: Callable[..., Any]) -> bool:
+ """Return whether a StructuralMutate callback accepts its optional mode."""
try:
signature = inspect.signature(callback)
except (TypeError, ValueError):
@@ -870,14 +891,14 @@ def _callback_accepts_allow_inplace(callback:
Callable[..., Any]) -> bool:
return False
try:
- signature.bind(None, None, False)
+ signature.bind(None, None, InplaceMode.DISALLOW)
except TypeError:
try:
signature.bind(None, None)
except TypeError as err:
raise TypeError(
"structural_mutate callback must accept (value, mutator) or "
- "(value, mutator, allow_inplace)"
+ "(value, mutator, inplace_mode)"
) from err
return False
return True
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index f5f2fcff..3764322e 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -57,11 +57,11 @@ Expected<Any> StructuralMapExpected(
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
using Mutator = StructuralMapDynEngine<StructuralMapEngineBase,
WalkOrder::kPreOrder>;
StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
- return mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ return mutator->MutateExpected(root, InplaceMode::kAllow);
}
using Mutator = StructuralMapDynEngine<StructuralMapEngineBase,
WalkOrder::kPostOrder>;
StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
- return mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ return mutator->MutateExpected(root, InplaceMode::kAllow);
}();
if (TVM_FFI_PREDICT_FALSE(result.is_err())) return
Unexpected(std::move(result).error());
UnchangedOr<Any> mapped = AnyUnsafe::MoveFromAnyAfterCheck<UnchangedOr<Any>>(
@@ -96,12 +96,12 @@ class StructuralMutateDynEngine : public Parent {
return
static_cast<StructuralMutateDynEngine*>(mutator)->MaybeInplaceMutateImplRaw(value);
}
- std::optional<Expected<Any>> DispatchCallback(AnyView value, bool
allow_inplace) noexcept {
+ std::optional<Expected<Any>> DispatchCallback(AnyView value, InplaceMode
inplace_mode) noexcept {
for (const auto& entry : callbacks_) {
if (!RuntimeTypeIndexMatch(value.type_index(), entry.template get<0>()))
continue;
if (entry.template get<2>()) {
return entry.template get<1>().template CallExpected<Any>(
- value, GetRef<StructuralMutator>(this), allow_inplace);
+ value, GetRef<StructuralMutator>(this), inplace_mode);
}
return entry.template get<1>().template CallExpected<Any>(value,
GetRef<StructuralMutator>(this));
@@ -110,7 +110,7 @@ class StructuralMutateDynEngine : public Parent {
}
TVMFFIAny MutateImplRaw(AnyView value) noexcept {
- if (std::optional<Expected<Any>> matched = DispatchCallback(value, false))
{
+ if (std::optional<Expected<Any>> matched = DispatchCallback(value,
InplaceMode::kDisallow)) {
Expected<Any> result = *std::move(matched);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
Parent::UpdateVisitErrorContext(result, value);
@@ -121,7 +121,7 @@ class StructuralMutateDynEngine : public Parent {
}
TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
- if (std::optional<Expected<Any>> matched = DispatchCallback(value, true)) {
+ if (std::optional<Expected<Any>> matched = DispatchCallback(value,
InplaceMode::kAllow)) {
Expected<Any> result = *std::move(matched);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
Parent::UpdateVisitErrorContext(result, value);
@@ -137,9 +137,10 @@ class StructuralMutateDynEngine : public Parent {
/*!
* \brief Runtime callback-driven structural mutation.
* \param root The root value to mutate.
- * \param callbacks Runtime ``(type_index, callback, accepts_allow_inplace)``
entries. A callback
- * whose marker is true is invoked as ``callback(value,
mutator, allow_inplace)``;
- * otherwise it is invoked as ``callback(value, mutator)``.
+ * \param callbacks Runtime ``(type_index, callback, accepts_inplace_mode)``
entries. A callback
+ * whose marker is true is invoked as ``callback(value,
mutator, inplace_mode)``;
+ * the mode uses the FFI integer representation of
InplaceMode. Otherwise the
+ * callback is invoked as ``callback(value, mutator)``.
* \return The mutated owning value, or an Error.
*/
// The owning parameter makes caller ownership visible to the uniqueness check.
@@ -148,7 +149,7 @@ Expected<Any> StructuralMutateExpected(
const Array<Tuple<int32_t, Function, bool>>& callbacks) noexcept {
using Mutator = StructuralMutateDynEngine<StructuralMapEngineBase>;
StructuralMutator mutator(make_object<Mutator>(callbacks));
- auto result = mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ auto result = mutator->MutateExpected(root, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) return
Unexpected(std::move(result).error());
UnchangedOr<Any> mapped = AnyUnsafe::MoveFromAnyAfterCheck<UnchangedOr<Any>>(
std::move(ExpectedUnsafe::GetData(result)));
@@ -182,7 +183,7 @@ TVM_FFI_INLINE TVMFFIAny
MutateSeqContainerChanged(StructuralMutatorObj* mutator
for (int64_t i = index + 1; i < size; ++i) {
const Any& item = items[i];
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, mapped_value,
- mutator->MutateExpected(item));
+ mutator->MutateExpected(item,
InplaceMode::kDisallow));
output->SetItemAfterCheck(i,
std::move(mapped_value).ValueOrUnchanged(AnyView(item)));
}
return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
@@ -204,7 +205,7 @@ TVMFFIAny MutateSeqContainerRaw(StructuralMutatorObj*
mutator, const SeqObj* sel
for (int64_t i = 0; i < size; ++i) {
const Any& item = items[i];
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, mapped_value,
- mutator->MutateExpected(item));
+ mutator->MutateExpected(item,
InplaceMode::kDisallow));
if (!mapped_value.UnchangedOrSameAs(item)) {
return MutateSeqContainerChanged(mutator, self, i,
std::move(mapped_value).ValueUnchecked());
}
@@ -225,7 +226,7 @@ TVMFFIAny
MaybeInplaceMutateSeqContainerRaw(StructuralMutatorObj* mutator, SeqOb
for (int64_t i = 0; i < static_cast<int64_t>(self->size()); ++i) {
const Any& item = self->begin()[i];
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, mapped_value,
-
mutator->MaybeInplaceMutateIfUniqueExpected(item));
+ mutator->MutateExpected(item,
InplaceMode::kAllow));
if (!mapped_value.UnchangedOrSameAs(item)) {
self->SetItemAfterCheck(i, std::move(mapped_value).ValueUnchecked());
}
@@ -261,7 +262,7 @@ TVM_FFI_INLINE TVMFFIAny
MutateMapValuesChanged(StructuralMutatorObj* mutator,
for (; source_it != self->end(); ++source_it, ++output_it) {
const Any& old_value = source_it->second;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, new_value,
- mutator->MutateExpected(old_value));
+ mutator->MutateExpected(old_value,
InplaceMode::kDisallow));
if (!new_value.UnchangedOrSameAs(old_value)) {
output_it->second = std::move(new_value).ValueUnchecked();
}
@@ -283,7 +284,7 @@ TVMFFIAny MutateMapValuesRaw(StructuralMutatorObj* mutator,
const MapObjType* se
for (auto source_it = self->begin(); source_it != self->end(); ++source_it,
++index) {
const Any& old_value = source_it->second;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, new_value,
- mutator->MutateExpected(old_value));
+ mutator->MutateExpected(old_value,
InplaceMode::kDisallow));
if (!new_value.UnchangedOrSameAs(old_value)) {
return MutateMapValuesChanged(mutator, self, source_it, index,
std::move(new_value).ValueUnchecked());
@@ -305,7 +306,7 @@ TVMFFIAny
MaybeInplaceMutateMapValuesRaw(StructuralMutatorObj* mutator, MapObjTy
for (auto it = self->begin(); it != self->end(); ++it) {
const Any& old_value = it->second;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, new_value,
-
mutator->MaybeInplaceMutateIfUniqueExpected(old_value));
+ mutator->MutateExpected(old_value,
InplaceMode::kAllow));
if (!new_value.UnchangedOrSameAs(old_value)) {
it->second = std::move(new_value).ValueUnchecked();
}
@@ -377,14 +378,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def_method("ffi.StructuralMutatorMutate",
[](const StructuralMutator& mutator, AnyView value) {
- return
std::move(mutator->Mutate(value)).ValueOrUnchanged(value);
- })
- .def_method("ffi.StructuralMutatorDefaultMutate",
- [](const StructuralMutator& mutator, AnyView value) -> Any {
- UnchangedOr<Any> result =
-
std::move(mutator->DefaultMutateExpected(value)).value();
- return std::move(result).ValueOrUnchanged(value);
+ return std::move(mutator->Mutate(value,
InplaceMode::kDisallow))
+ .ValueOrUnchanged(value);
})
+ .def_method(
+ "ffi.StructuralMutatorDefaultMutate",
+ [](const StructuralMutator& mutator, AnyView value) -> Any {
+ UnchangedOr<Any> result =
+ std::move(mutator->DefaultMutateExpected(value,
InplaceMode::kDisallow)).value();
+ return std::move(result).ValueOrUnchanged(value);
+ })
.def_method("ffi.StructuralMutatorVarRemapGet",
[](const StructuralMutator& mutator, AnyView var) {
return mutator->VarRemapGetExpected(var).value();
diff --git a/tests/cpp/extra/test_big_int.cc b/tests/cpp/extra/test_big_int.cc
index 3d85e790..dcd952bf 100644
--- a/tests/cpp/extra/test_big_int.cc
+++ b/tests/cpp/extra/test_big_int.cc
@@ -141,8 +141,8 @@ TEST(BigIntExtra, StructuralMutation) {
auto never_matches = [](int64_t, StructuralMutatorObj*) -> Expected<Any> {
return Any(); };
using Mutator = StructuralMutateEngine<StructuralMapEngineBase,
decltype(never_matches)>;
StructuralMutator mutator(make_object<Mutator>(never_matches));
- EXPECT_TRUE(mutator->Mutate(AnyView(value)).IsUnchanged());
- EXPECT_TRUE(mutator->MaybeInplaceMutate(AnyView(value)).IsUnchanged());
+ EXPECT_TRUE(mutator->Mutate(AnyView(value),
InplaceMode::kDisallow).IsUnchanged());
+ EXPECT_TRUE(mutator->Mutate(AnyView(value),
InplaceMode::kAllow).IsUnchanged());
EXPECT_TRUE(StructuralMutateExpected(Any(value),
never_matches).value().same_as(Any(value)));
auto map_int = [](int64_t x) -> Expected<Any> { return Any(x + 1); };
EXPECT_TRUE(StructuralMapExpected<WalkOrder::kPostOrder>(Any(value), map_int)
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 5024613e..f60bdf3e 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -40,13 +40,12 @@ using namespace tvm::ffi::testing;
using AnyArray = Array<Any>;
using StringMap = Map<String, Any>;
+static_assert(std::is_same_v<decltype(std::declval<StructuralMutatorObj&>().MutateExpected(
+ std::declval<AnyView>(),
InplaceMode::kDisallow)),
+ Expected<UnchangedOr<Any>>>);
static_assert(std::is_same_v<decltype(std::declval<StructuralMutatorObj&>().DefaultMutateExpected(
- std::declval<AnyView>())),
+ std::declval<AnyView>(),
InplaceMode::kAllow)),
Expected<UnchangedOr<Any>>>);
-static_assert(
-
std::is_same_v<decltype(std::declval<StructuralMutatorObj&>().DefaultMaybeInplaceMutateExpected(
- std::declval<AnyView>())),
- Expected<UnchangedOr<Any>>>);
// ---------------------------------------------------------------------------
// Unchanged result protocol.
@@ -79,6 +78,21 @@ TEST(UnchangedOr, ConversionsAndAssignmentMacro) {
EXPECT_TRUE(std::move(moved).ValueUnchecked().same_as(original));
EXPECT_EQ(original.use_count(), 1);
+ const TInt borrowed(7);
+ TInt typed_original =
UnchangedOr<TInt>(Unchanged()).ValueOrUnchanged(borrowed);
+ EXPECT_TRUE(typed_original.same_as(borrowed));
+ EXPECT_EQ(borrowed.use_count(), 2);
+ TNumber base_original =
UnchangedOr<TNumber>(Unchanged()).ValueOrUnchanged(borrowed);
+ EXPECT_TRUE(base_original.same_as(borrowed));
+ EXPECT_EQ(borrowed.use_count(), 3);
+
+ UnchangedOr<TNumber> replacement = TInt(8);
+ TNumber replaced = std::move(replacement).ValueOrUnchanged(borrowed);
+ EXPECT_EQ(replaced.as_or_throw<TInt>()->value, 8);
+ EXPECT_TRUE(replaced.unique());
+ EXPECT_EQ(borrowed->value, 7);
+ EXPECT_EQ(borrowed.use_count(), 3);
+
UnchangedOr<double> numeric = UnchangedOr<int>(42);
EXPECT_EQ(AnyView(numeric).type_index(), TypeIndex::kTVMFFIFloat);
EXPECT_DOUBLE_EQ(std::move(numeric).ValueUnchecked(), 42.0);
@@ -165,32 +179,20 @@ TEST(StructuralMutate,
UnchangedProtocolResolvesAtThrowingEntryPoints) {
EXPECT_TRUE(ReturnTypedUnchangedExpected().value().IsUnchanged());
auto never_matches = [](int64_t, StructuralMutatorObj*) -> Expected<Any> {
return Any(); };
using Mutator = StructuralMutateEngine<StructuralMapEngineBase,
decltype(never_matches)>;
- StructuralMutator mutator(make_object<Mutator>(std::move(never_matches)));
+ StructuralMutator mutator(make_object<Mutator>(never_matches));
String value("value longer than small-string storage");
- auto untyped = mutator->MutateExpected(AnyView(value));
- ASSERT_TRUE(untyped.is_ok());
- EXPECT_TRUE(std::move(untyped).value().IsUnchanged());
- auto untyped_throwing = mutator->Mutate(AnyView(value));
- EXPECT_TRUE(untyped_throwing.IsUnchanged());
-
EXPECT_TRUE(std::move(untyped_throwing).ValueOrUnchanged(AnyView(value)).same_as(value));
-
- auto typed = mutator->MutateExpected<String>(value);
- ASSERT_TRUE(typed.is_ok());
- EXPECT_TRUE(std::move(typed).value().IsUnchanged());
- auto typed_throwing = mutator->Mutate<String>(value);
- EXPECT_TRUE(typed_throwing.IsUnchanged());
- String moved_value = value;
- String moved_value_alias = moved_value;
-
EXPECT_EQ(std::move(typed_throwing).ValueOrUnchanged(std::move(moved_value)),
moved_value_alias);
- EXPECT_TRUE(mutator->MaybeInplaceMutate(AnyView(value)).IsUnchanged());
- EXPECT_TRUE(mutator->MaybeInplaceMutate<String>(value).IsUnchanged());
- auto unique = mutator->MaybeInplaceMutateIfUniqueExpected(AnyView(value));
- ASSERT_TRUE(unique.is_ok());
- EXPECT_TRUE(std::move(unique).value().IsUnchanged());
- auto typed_unique =
mutator->MaybeInplaceMutateIfUniqueExpected<String>(value);
- ASSERT_TRUE(typed_unique.is_ok());
- EXPECT_TRUE(std::move(typed_unique).value().IsUnchanged());
+ for (InplaceMode inplace_mode : {InplaceMode::kDisallow,
InplaceMode::kAllow}) {
+ Any packed_mode(inplace_mode);
+ EXPECT_EQ(packed_mode.type_index(), TypeIndex::kTVMFFIInt);
+ EXPECT_EQ(packed_mode.cast<int64_t>(), static_cast<int32_t>(inplace_mode));
+ auto result = mutator->MutateExpected(AnyView(value), inplace_mode);
+ ASSERT_TRUE(result.is_ok());
+ EXPECT_TRUE(std::move(result).value().IsUnchanged());
+ auto throwing_result = mutator->Mutate(AnyView(value), inplace_mode);
+ EXPECT_TRUE(throwing_result.IsUnchanged());
+
EXPECT_TRUE(std::move(throwing_result).ValueOrUnchanged(AnyView(value)).same_as(value));
+ }
EXPECT_TRUE(StructuralMutateExpected(Any(value),
never_matches).value().same_as(value));
EXPECT_TRUE(StructuralMapExpected<WalkOrder::kPostOrder>(
Any(value), [](int64_t item) -> Expected<Any> { return
Any(item + 1); })
@@ -206,30 +208,22 @@ TEST(StructuralMutate,
UnchangedProtocolResolvesAtThrowingEntryPoints) {
StructuralMutator wrong_type_mutator(
make_object<WrongTypeMutator>(std::move(replace_int_with_string)));
- auto any_error = wrong_type_mutator->MutateExpected(int64_t{-1});
- ASSERT_TRUE(any_error.is_err());
- EXPECT_EQ(any_error.error().message(), "direct-forward failure");
- auto any_inplace_error =
wrong_type_mutator->MaybeInplaceMutateExpected(int64_t{-1});
- ASSERT_TRUE(any_inplace_error.is_err());
- EXPECT_EQ(any_inplace_error.error().message(), "direct-forward failure");
- auto typed_error = wrong_type_mutator->MutateExpected<int64_t>(int64_t{-1});
- ASSERT_TRUE(typed_error.is_err());
- EXPECT_EQ(typed_error.error().message(), "direct-forward failure");
- auto typed_inplace_error =
wrong_type_mutator->MaybeInplaceMutateExpected<int64_t>(int64_t{-1});
- ASSERT_TRUE(typed_inplace_error.is_err());
- EXPECT_EQ(typed_inplace_error.error().message(), "direct-forward failure");
-
- Expected<UnchangedOr<int64_t>> result =
wrong_type_mutator->MutateExpected<int64_t>(int64_t{1});
- ASSERT_TRUE(result.is_err());
- EXPECT_EQ(result.error().kind(), "TypeError");
- Expected<UnchangedOr<int64_t>> inplace_result =
- wrong_type_mutator->MaybeInplaceMutateExpected<int64_t>(int64_t{1});
- ASSERT_TRUE(inplace_result.is_err());
- EXPECT_EQ(inplace_result.error().kind(), "TypeError");
- Expected<UnchangedOr<int64_t>> unique_result =
-
wrong_type_mutator->MaybeInplaceMutateIfUniqueExpected<int64_t>(int64_t{1});
- ASSERT_TRUE(unique_result.is_err());
- EXPECT_EQ(unique_result.error().kind(), "TypeError");
+ auto narrow_result = [&](int64_t input,
+ InplaceMode inplace_mode) ->
Expected<UnchangedOr<int64_t>> {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<int64_t>, result,
+
wrong_type_mutator->MutateExpected(input, inplace_mode));
+ return result;
+ };
+ for (InplaceMode inplace_mode : {InplaceMode::kDisallow,
InplaceMode::kAllow}) {
+ auto failure = narrow_result(-1, inplace_mode);
+ ASSERT_TRUE(failure.is_err());
+ EXPECT_EQ(failure.error().message(), "direct-forward failure");
+ auto wrong_type = narrow_result(1, inplace_mode);
+ ASSERT_TRUE(wrong_type.is_err());
+ EXPECT_EQ(wrong_type.error().kind(), "TypeError");
+ EXPECT_ANY_THROW(
+ wrong_type_mutator->Mutate(1,
inplace_mode).ValueOrUnchanged(AnyView(1)).cast<int64_t>());
+ }
for (WalkOrder order : {WalkOrder::kPreOrder, WalkOrder::kPostOrder}) {
TVar root("n");
@@ -262,7 +256,7 @@ class TNestedMapHookObj : public Object {
static TVMFFIAny StructuralMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
const auto* self = value.cast<const TNestedMapHookObj*>();
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Any>, mapped,
- mutator->MutateExpected(self->field));
+ mutator->MutateExpected(self->field,
InplaceMode::kDisallow));
if (mapped.UnchangedOrSameAs(Any(self->field))) {
return Unchanged().CopyToTVMFFIAny();
}
@@ -331,16 +325,15 @@ class StructuralMapWithMutateCount : public
StructuralMapEngineBase {
const MutateCount& count() const { return count_; }
- Expected<UnchangedOr<Any>> DefaultMutateExpected(AnyView value) noexcept {
+ Expected<UnchangedOr<Any>> DefaultMutateExpected(AnyView value,
+ InplaceMode inplace_mode)
noexcept {
++count_.value;
- ++count_.mutate_expected;
- return StructuralMapEngineBase::DefaultMutateExpected(value);
- }
-
- Expected<UnchangedOr<Any>> DefaultMaybeInplaceMutateExpected(AnyView value)
noexcept {
- ++count_.value;
- ++count_.maybe_inplace_expected;
- return StructuralMapEngineBase::DefaultMaybeInplaceMutateExpected(value);
+ if (inplace_mode == InplaceMode::kAllow) {
+ ++count_.maybe_inplace_expected;
+ } else {
+ ++count_.mutate_expected;
+ }
+ return StructuralMapEngineBase::DefaultMutateExpected(value, inplace_mode);
}
protected:
@@ -388,13 +381,13 @@ TEST(StructuralMap,
ParentLayerOwnsBothDescentsAndProvidesState) {
EXPECT_EQ(mutator->VarRemapGetExpected(1).error().kind(), "TypeError");
mutator->VarRemapSetExpected(key, Any(Unchanged())).value();
- ASSERT_FALSE(mutator->MutateExpected(String("unmatched")).is_err());
+ ASSERT_FALSE(mutator->MutateExpected(String("unmatched"),
InplaceMode::kDisallow).is_err());
AnyArray rebuild_root{int64_t{1}};
- ASSERT_FALSE(mutator->MutateExpected(rebuild_root).is_err());
+ ASSERT_FALSE(mutator->MutateExpected(rebuild_root,
InplaceMode::kDisallow).is_err());
-
ASSERT_FALSE(mutator->MaybeInplaceMutateExpected(String("unmatched")).is_err());
+ ASSERT_FALSE(mutator->MutateExpected(String("unmatched"),
InplaceMode::kAllow).is_err());
AnyArray inplace_root{int64_t{1}};
- ASSERT_FALSE(mutator->MaybeInplaceMutateExpected(inplace_root).is_err());
+ ASSERT_FALSE(mutator->MutateExpected(inplace_root,
InplaceMode::kAllow).is_err());
EXPECT_GT(engine->count().mutate_expected, 0);
EXPECT_GT(engine->count().maybe_inplace_expected, 0);
@@ -404,8 +397,9 @@ TEST(StructuralMap,
ParentLayerOwnsBothDescentsAndProvidesState) {
TVar var("n");
AnyArray repeated{var, var};
- AnyArray mapped =
-
std::move(mutator->Mutate<AnyArray>(repeated)).ValueOrUnchanged(std::move(repeated));
+ AnyArray mapped = mutator->Mutate(repeated, InplaceMode::kDisallow)
+ .ValueOrUnchanged(AnyView(repeated))
+ .cast<AnyArray>();
EXPECT_EQ(var_callback_count, 2);
EXPECT_FALSE(mapped[0].cast<TVar>().same_as(mapped[1].cast<TVar>()));
}
@@ -414,7 +408,7 @@ TEST(StructuralMutate,
CallbackOwnsMutationAndErrorsStayExpected) {
std::vector<int64_t> trace;
auto mutate_array = [&](const AnyArray& value, StructuralMutateLayer*
mutator) -> Expected<Any> {
EXPECT_EQ(mutator->callback_tag(), 23);
- auto first_result = mutator->MutateExpected(value[0]);
+ auto first_result = mutator->MutateExpected(value[0],
InplaceMode::kDisallow);
if (TVM_FFI_PREDICT_FALSE(first_result.is_err())) {
return Unexpected(std::move(first_result).error());
}
@@ -430,7 +424,9 @@ TEST(StructuralMutate,
CallbackOwnsMutationAndErrorsStayExpected) {
StructuralMutator mutator(make_object<Mutator>(std::move(mutate_array),
std::move(mutate_int)));
AnyArray root{int64_t{1}, int64_t{2}};
- AnyArray mapped =
std::move(mutator->Mutate<AnyArray>(root)).ValueOrUnchanged(std::move(root));
+ AnyArray mapped = mutator->Mutate(root, InplaceMode::kDisallow)
+ .ValueOrUnchanged(AnyView(root))
+ .cast<AnyArray>();
ASSERT_EQ(mapped.size(), 2U);
EXPECT_EQ(mapped[0].cast<int64_t>(), 2);
EXPECT_EQ(mapped[1].cast<int64_t>(), 10);
@@ -468,8 +464,9 @@ TEST(StructuralMutate, CallbackControlsRecursion) {
StructuralMutate(
root,
[](const TPair& pair, StructuralMutatorObj* mutator) ->
Expected<UnchangedOr<ObjectRef>> {
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<ObjectRef>,
lhs_result,
-
mutator->MutateExpected<ObjectRef>(pair->lhs));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ObjectRef>, lhs_result,
+ mutator->MutateExpected(pair->lhs, InplaceMode::kDisallow));
ObjectRef original_lhs = pair->lhs;
ObjectRef lhs =
std::move(lhs_result).ValueOrUnchanged(std::move(original_lhs));
return UnchangedOr<ObjectRef>(TPair(std::move(lhs), pair->rhs));
@@ -493,7 +490,7 @@ TEST(StructuralMutate, SingleCallbackCanDelegateToDefault) {
if (auto integer = value.as<int64_t>()) {
return Any(*integer + 1);
}
- return mutator->DefaultMutateExpected(value);
+ return mutator->DefaultMutateExpected(value, InplaceMode::kDisallow);
};
AnyArray root{int64_t{1}, AnyArray{int64_t{2}}};
AnyArray result = StructuralMutate(root, mutate).cast<AnyArray>();
@@ -503,20 +500,56 @@ TEST(StructuralMutate,
SingleCallbackCanDelegateToDefault) {
}
TEST(StructuralMutate, PreservesUniqueContainerIdentity) {
+ auto increment = [](int64_t value, StructuralMutatorObj*) -> Expected<Any> {
+ return Any(value + 1);
+ };
AnyArray inner{int64_t{1}};
const Object* inner_address = inner.get();
AnyArray root{Any(std::move(inner))};
const Object* root_address = root.get();
- AnyArray mapped =
- StructuralMutate(std::move(root), [](int64_t value,
StructuralMutatorObj*) -> Expected<Any> {
- return Any(value + 1);
- }).cast<AnyArray>();
+ AnyArray mapped = StructuralMutate(std::move(root),
increment).cast<AnyArray>();
AnyArray mapped_inner = mapped[0].cast<AnyArray>();
EXPECT_EQ(mapped.get(), root_address);
EXPECT_EQ(mapped_inner.get(), inner_address);
EXPECT_EQ(mapped_inner[0].cast<int64_t>(), 2);
+
+ using Mutator = StructuralMutateEngine<StructuralMapEngineBase,
decltype(increment)>;
+ StructuralMutator mutator(make_object<Mutator>(increment));
+ AnyArray default_original{int64_t{1}};
+ EXPECT_FALSE(mutator->Mutate(default_original)
+ .ValueOrUnchanged(AnyView(default_original))
+ .cast<AnyArray>()
+ .same_as(default_original));
+ EXPECT_FALSE(std::move(mutator->MutateExpected(default_original))
+ .value()
+ .ValueOrUnchanged(AnyView(default_original))
+ .cast<AnyArray>()
+ .same_as(default_original));
+ EXPECT_FALSE(std::move(mutator->DefaultMutateExpected(default_original))
+ .value()
+ .ValueOrUnchanged(AnyView(default_original))
+ .cast<AnyArray>()
+ .same_as(default_original));
+ EXPECT_EQ(default_original[0].cast<int64_t>(), 1);
+ EXPECT_TRUE(default_original.unique());
+
+ for (InplaceMode inplace_mode : {InplaceMode::kDisallow,
InplaceMode::kAllow}) {
+ for (bool shared : {false, true}) {
+ AnyArray original{int64_t{1}};
+ Any alias = shared ? Any(original) : Any();
+ EXPECT_EQ(original.use_count(), shared ? 2 : 1);
+ AnyArray result = std::move(mutator->MutateExpected(original,
inplace_mode))
+ .value()
+ .ValueOrUnchanged(AnyView(original))
+ .cast<AnyArray>();
+ EXPECT_EQ(result.same_as(original), inplace_mode == InplaceMode::kAllow
&& !shared);
+ EXPECT_EQ(original[0].cast<int64_t>(),
+ inplace_mode == InplaceMode::kAllow && !shared ? 2 : 1);
+ EXPECT_EQ(result[0].cast<int64_t>(), 2);
+ }
+ }
}
TEST(StructuralMutate, RootByValueProtectsSharedParentSubvalue) {
@@ -532,6 +565,26 @@ TEST(StructuralMutate,
RootByValueProtectsSharedParentSubvalue) {
EXPECT_NE(mapped.get(), child_address);
EXPECT_EQ(outer[0].cast<AnyArray>()[0].cast<int64_t>(), 1);
EXPECT_EQ(mapped[0].cast<int64_t>(), 2);
+
+ AnyArray shared_outer = outer; //
NOLINT(performance-unnecessary-copy-initialization)
+ bool denied_unique_child = false;
+ auto mutate = [&](AnyView value, StructuralMutatorObj* mutator,
+ InplaceMode inplace_mode) -> Expected<UnchangedOr<Any>> {
+ if (auto integer = value.as<int64_t>()) return Any(*integer + 1);
+ if (value.as<Object>()->unique() && inplace_mode ==
InplaceMode::kDisallow) {
+ denied_unique_child = true;
+ }
+ return mutator->DefaultMutateExpected(value, inplace_mode);
+ };
+ using Mutator = StructuralMutateEngine<StructuralMapEngineBase,
decltype(mutate)>;
+ StructuralMutator mutator(make_object<Mutator>(mutate));
+ AnyArray rebuilt = std::move(mutator->MutateExpected(outer,
InplaceMode::kAllow))
+ .value()
+ .ValueOrUnchanged(AnyView(outer))
+ .cast<AnyArray>();
+ EXPECT_TRUE(denied_unique_child);
+ EXPECT_EQ(shared_outer[0].cast<AnyArray>()[0].cast<int64_t>(), 1);
+ EXPECT_EQ(rebuilt[0].cast<AnyArray>()[0].cast<int64_t>(), 2);
}
TEST(StructuralMutate, CallbackArityControlsInplaceMutation) {
@@ -539,18 +592,19 @@ TEST(StructuralMutate,
CallbackArityControlsInplaceMutation) {
AnyArray copy_on_write_root{int64_t{1}};
const Object* inplace_root_address = inplace_root.get();
const Object* copy_on_write_root_address = copy_on_write_root.get();
- std::vector<bool> allow_inplace_trace;
+ std::vector<InplaceMode> inplace_mode_trace;
AnyArray inplace_mapped =
StructuralMutate(
std::move(inplace_root),
[&](const AnyArray& value, StructuralMutatorObj* mutator,
- bool allow_inplace) -> Expected<Any> {
- allow_inplace_trace.push_back(allow_inplace);
- return mutator->DefaultMaybeInplaceMutateExpected(value,
allow_inplace);
+ InplaceMode inplace_mode) -> Expected<Any> {
+ inplace_mode_trace.push_back(inplace_mode);
+ EXPECT_GT(value.use_count(), 1);
+ return mutator->DefaultMutateExpected(value, inplace_mode);
},
- [&](int64_t value, StructuralMutatorObj*, bool allow_inplace) ->
Expected<Any> {
- allow_inplace_trace.push_back(allow_inplace);
+ [&](int64_t value, StructuralMutatorObj*, InplaceMode inplace_mode)
-> Expected<Any> {
+ inplace_mode_trace.push_back(inplace_mode);
return Any(value + 1);
})
.cast<AnyArray>();
@@ -559,7 +613,7 @@ TEST(StructuralMutate,
CallbackArityControlsInplaceMutation) {
StructuralMutate(
std::move(copy_on_write_root),
[](const AnyArray& value, StructuralMutatorObj* mutator) ->
Expected<Any> {
- return mutator->DefaultMutateExpected(value);
+ return mutator->DefaultMutateExpected(value,
InplaceMode::kDisallow);
},
[](int64_t value, StructuralMutatorObj*) -> Expected<Any> { return
Any(value + 1); })
.cast<AnyArray>();
@@ -568,7 +622,8 @@ TEST(StructuralMutate,
CallbackArityControlsInplaceMutation) {
EXPECT_NE(copy_on_write_mapped.get(), copy_on_write_root_address);
EXPECT_EQ(inplace_mapped[0].cast<int64_t>(), 2);
EXPECT_EQ(copy_on_write_mapped[0].cast<int64_t>(), 2);
- EXPECT_EQ(allow_inplace_trace, (std::vector<bool>{true, false}));
+ EXPECT_EQ(inplace_mode_trace,
+ (std::vector<InplaceMode>{InplaceMode::kAllow,
InplaceMode::kDisallow}));
}
TEST(StructuralMutate, MatchedVarOwnsRemapConsistency) {
@@ -1106,7 +1161,9 @@ void CheckDynamicParentLayer() {
StructuralMutator mutator(engine);
AnyArray root{int64_t{1}};
- AnyArray mapped =
std::move(mutator->Mutate<AnyArray>(root)).ValueOrUnchanged(std::move(root));
+ AnyArray mapped = mutator->Mutate(root, InplaceMode::kDisallow)
+ .ValueOrUnchanged(AnyView(root))
+ .cast<AnyArray>();
EXPECT_EQ(mapped[0].cast<int64_t>(), 2);
EXPECT_EQ(calls, 1);
EXPECT_GT(engine->count().value, 0);
diff --git a/tests/cpp/testing_object.h b/tests/cpp/testing_object.h
index 5553c6c3..366bb96a 100644
--- a/tests/cpp/testing_object.h
+++ b/tests/cpp/testing_object.h
@@ -254,9 +254,9 @@ class TMutatePairObj : public Object {
const TMutatePairObj* self =
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TMutatePairObj>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<ObjectRef>, lhs,
- mutator->MutateExpected(self->lhs));
+ mutator->MutateExpected(self->lhs,
InplaceMode::kDisallow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<ObjectRef>, rhs,
- mutator->MutateExpected(self->rhs));
+ mutator->MutateExpected(self->rhs,
InplaceMode::kDisallow));
if (lhs.UnchangedOrSameAs(self->lhs) && rhs.UnchangedOrSameAs(self->rhs)) {
return Unchanged().CopyToTVMFFIAny();
}
diff --git a/tests/python/test_structural.py b/tests/python/test_structural.py
index 8c5a7901..197b5157 100644
--- a/tests/python/test_structural.py
+++ b/tests/python/test_structural.py
@@ -329,6 +329,41 @@ def test_structural_visit_default_visit_binding() -> None:
assert nested_trace == ["array", 1, 2]
+def test_structural_mutate_inplace_mode() -> None:
+ assert tvm_ffi.InplaceMode.DISALLOW == 0
+ assert tvm_ffi.InplaceMode.ALLOW == 1
+ assert not tvm_ffi.InplaceMode.DISALLOW
+ assert tvm_ffi.InplaceMode.ALLOW
+ inplace_trace: list[int] = []
+
+ def mutate_with_mode(value: int, mutator: tvm_ffi.StructuralMutator,
inplace_mode: int) -> int:
+ assert isinstance(mutator, tvm_ffi.StructuralMutator)
+ assert type(inplace_mode) is int
+ inplace_trace.append(inplace_mode)
+ return value + 1
+
+ mode_root = tvm_ffi.Array([1])
+ mode_mapped = tvm_ffi.structural_mutate(mode_root, (int, mutate_with_mode))
+ assert inplace_trace == [tvm_ffi.InplaceMode.DISALLOW]
+ assert not mode_mapped.same_as(mode_root)
+ assert list(mode_root) == [1]
+ assert list(mode_mapped) == [2]
+
+ def mutate_array_with_mode(
+ value: tvm_ffi.Array, mutator: tvm_ffi.StructuralMutator,
inplace_mode: int
+ ) -> object:
+ assert type(inplace_mode) is int
+ inplace_trace.append(inplace_mode)
+ return mutator.default_mutate(value)
+
+ owned_root = tvm_ffi.Array([1])
+ owned_mapped = tvm_ffi.structural_mutate(
+ owned_root._move(), (tvm_ffi.Array, mutate_array_with_mode)
+ )
+ assert inplace_trace == [tvm_ffi.InplaceMode.DISALLOW,
tvm_ffi.InplaceMode.ALLOW]
+ assert list(owned_mapped) == [1]
+
+
def test_structural_mutate_callback_owned_recursion_and_errors() -> None:
trace: list[int | str] = []
@@ -372,22 +407,6 @@ def
test_structural_mutate_callback_owned_recursion_and_errors() -> None:
assert list(default_root) == [3, 4]
assert list(default_mapped) == [4, 5]
- inplace_trace: list[bool] = []
-
- def mutate_with_flag(
- value: int, mutator: tvm_ffi.StructuralMutator, allow_inplace: bool
- ) -> int:
- assert isinstance(mutator, tvm_ffi.StructuralMutator)
- inplace_trace.append(allow_inplace)
- return value + 1
-
- flagged_root = tvm_ffi.Array([1])
- flagged_mapped = tvm_ffi.structural_mutate(flagged_root, (int,
mutate_with_flag))
- assert inplace_trace == [False]
- assert not flagged_mapped.same_as(flagged_root)
- assert list(flagged_root) == [1]
- assert list(flagged_mapped) == [2]
-
direct_trace: list[int] = []
def fail_directly(value: int, mutator: tvm_ffi.StructuralMutator) ->
object:
