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 1cb982fe [REFACTOR][EXTRA] Allow StructuralMapMutator to carry state
via Parent (#749)
1cb982fe is described below
commit 1cb982feab6708db86d4ae821a032d8fd27ee7a0
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 5 16:09:55 2026 -0400
[REFACTOR][EXTRA] Allow StructuralMapMutator to carry state via Parent
(#749)
Downstream mutators sometimes need to carry state alongside traversal,
scoped around the subtree they descend into. StructuralMapMutator cannot
express that today: there is no base slot a caller can supply, and its
dynamic counterpart is confined to a source file, so downstream cannot
layer over it either.
This change parameterizes StructuralMapEngine and StructuralMapDynEngine
on a Parent layer. The named StructuralMap and StructuralMapExpected
entry points remain unchanged and continue to use
StructuralMapEngineBase.
A Parent declares StateTupleType and a protected StateTuple() returning
it by value, declares paired raw and Expected descents when overriding
traversal, and accepts the concrete engine vtable in its constructor.
Typed callbacks receive (value, state...) with an optional trailing
def-region kind. Descent uses dependent this-> lookup, matching
StructuralWalkEngine, so composed deeper-layer declarations remain
eligible; explicit Parent qualification would pin the call at that
layer. Shared variable remapping remains owned by
StructuralMapEngineBase.
The dynamic engine moves into the public header with the same Parent
protocol. Compile-time remap classification removes the remap machinery
only for final, statically non-remappable object types. Positive,
nullable, non-final, and erased cases retain the runtime metadata query
so typed callbacks preserve dynamic-engine behavior for None, redeclared
kinds, and missing metadata.
---
include/tvm/ffi/extra/structural_mutate.h | 452 +++++++++++++++++++++++-------
src/ffi/extra/structural_mutate.cc | 212 +-------------
tests/cpp/extra/test_structural_mutate.cc | 126 +++++++++
3 files changed, 477 insertions(+), 313 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index b782b161..b3ab3b13 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -58,10 +58,10 @@ class StructuralMutatorObj;
* \param value The borrowed value to mutate.
* \return Raw ``TVMFFIAny`` containing the mutated value or an Error.
*
- * \note The hook is exception-free like \ref FStructuralVisit. Representable
failures must be
- * returned as an Error. Hook implementations should use non-throwing
accessors when the
- * engine's type dispatch has already established the type; allocation
failure and violated
- * container invariants remain fatal.
+ * \note The hook is exception-free. Representable failures must be returned
as an Error. Hook
+ * implementations should use non-throwing accessors when the engine's
type dispatch has
+ * already established the type; allocation failure and violated
container invariants
+ * remain fatal.
*/
using FStructuralMutate = TVMFFIAny (*)(StructuralMutatorObj* mutator, AnyView
value) noexcept;
@@ -686,40 +686,31 @@ TVM_FFI_COLD_CODE inline TVMFFIAny
SMutateDeclaredTypeErrorRaw() noexcept {
TVM_FFI_UNSAFE_S_MUTATE_ASSIGN_OR_RETURN_SKIP_CHECK_IMPL_(
\
TVM_FFI_STR_CONCAT(tvm_ffi_mutate_result_, __COUNTER__), Type, Name,
ResultExpr)
-/// \cond Doxygen_Suppress
-/// \endcond
-
-/*!
- * \brief Structural mutator that invokes typed callbacks during recursive
mapping.
- *
- * \tparam order Callback placement relative to child mapping.
- * \tparam Dispatch Callback dispatcher.
- * \sa StructuralMapCallbackChain
- */
-/*!
- * \brief A runtime table of Function callbacks, usable as a single link.
- *
- * The typed links are matched at compile time from their argument type. The
Python-driven API
- * instead carries a runtime list keyed by type index, so it appears to the
mutator as one link
- * that performs its own lookup. That keeps a single traversal for both
dispatch strategies.
- */
+} // namespace details
/*!
- * \brief Base of both the static and the dynamic StructuralMapMutator.
- *
- * Carries the identity-substitution environment they share. The dispatch
thunks downcast only
- * as far as this class, so both reuse them regardless of how they store their
callbacks.
+ * \brief Shared state and variable-remap dispatch for both structural-map
engines.
*
+ * This base owns the identity-substitution environment used for FreeVar and
DAG identities.
+ * Its variable-remap vtable thunks downcast to this common subobject,
allowing the typed and
+ * dynamic engines to share the environment even when a Parent layer sits
above this base.
*/
-class StructuralMapMutatorBaseObj : public StructuralMutatorObj {
+class StructuralMapEngineBase : public StructuralMutatorObj {
public:
- explicit StructuralMapMutatorBaseObj(const StructuralMutatorVTable* vtable)
+ /*! \brief Empty callback-state protocol used when no custom Parent layer is
present. */
+ using StateTupleType = std::tuple<>;
+
+ /*! \brief Construct the shared engine base with the concrete engine's
vtable. */
+ explicit StructuralMapEngineBase(const StructuralMutatorVTable* vtable)
: StructuralMutatorObj(vtable) {}
protected:
+ /*! \brief Return the empty state tuple exposed to typed map callbacks. */
+ TVM_FFI_INLINE StateTupleType StateTuple() const noexcept { return {}; }
+
/// \cond Doxygen_Suppress
// Out of line so its strings stay out of the per-node dispatch function,
which TryLink inlines
- // into. Shared by both map mutators: the typed one here and the dynamic one
in the .cc.
+ // into. Shared by the typed and dynamic engines below.
TVM_FFI_COLD_CODE static Expected<Any> SMutateDescentTypeError() noexcept {
return Unexpected(Error("TypeError", "structural mutate: descent changed
the node type", ""));
}
@@ -732,8 +723,8 @@ class StructuralMapMutatorBaseObj : public
StructuralMutatorObj {
* \return Raw ``TVMFFIAny`` containing the owning replacement, FFI None, or
Error.
*/
static TVMFFIAny DispatchVarRemapGet(StructuralMutatorObj* mutator, AnyView
var) noexcept {
- auto* self = static_cast<StructuralMapMutatorBaseObj*>(mutator);
- return ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapGetImpl(var));
+ auto* self = static_cast<StructuralMapEngineBase*>(mutator);
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapGetImpl(var));
}
/*!
@@ -745,38 +736,8 @@ class StructuralMapMutatorBaseObj : public
StructuralMutatorObj {
*/
static TVMFFIAny DispatchVarRemapSet(StructuralMutatorObj* mutator, AnyView
var,
AnyView mapped_value) noexcept {
- auto* self = static_cast<StructuralMapMutatorBaseObj*>(mutator);
- return ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapSetImpl(var,
mapped_value));
- }
-
- /*!
- * \brief Invoke a matched callback with optional def-region context.
- * \tparam Callback Callable returning a value implicitly convertible to
``Expected<Any>``.
- * \tparam Value Type of the converted value passed to the callback.
- * \param callback The matched callback.
- * \param value The converted value passed to the callback.
- * \param kind The active def-region kind.
- * \return The callback result normalized to ``Expected<Any>``.
- *
- * The callback may return a different type than the one that selected the
link; only the
- * caller holding a field's static type can check that.
- */
- template <typename Callback, typename Value>
- TVM_FFI_INLINE static Expected<Any> InvokeCallbackLink(Callback& callback,
Value&& value,
- TVMFFIDefRegionKind
kind) {
- using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
- static_assert(std::is_convertible_v<typename FuncInfo::RetType,
Expected<Any>>,
- "StructuralMap callbacks must return a replacement value,
Error, Unexpected, "
- "or Expected<U> implicitly convertible to Expected<Any>");
- try {
- if constexpr (FuncInfo::num_args == 1) {
- return callback(std::forward<Value>(value));
- } else {
- return callback(std::forward<Value>(value), kind);
- }
- } catch (const Error& err) {
- return Unexpected(err);
- }
+ auto* self = static_cast<StructuralMapEngineBase*>(mutator);
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapSetImpl(var,
mapped_value));
}
/*!
@@ -843,35 +804,57 @@ class StructuralMapMutatorBaseObj : public
StructuralMutatorObj {
};
/*!
- * \brief Structural mutator that invokes statically typed callbacks during
recursive mapping.
+ * \brief Callback-dispatched structural mutator with a state-carrying Parent
layer.
*
* Each callback is an ordinary callable and the engine selects it on its
first argument's
* type, so selection is a compile-time-known ``as<TSub>()`` on the input node.
*
+ * ``Parent`` derives from ``StructuralMapEngineBase``, publishes
``StateTupleType``, accepts
+ * and forwards the mutator vtable in its constructor, and provides a protected
+ * ``StateTuple() const noexcept``. A Parent that overrides expected descent
must hand-write
+ * the matching protected raw redirect so neither entry path silently bypasses
the layer. Each
+ * callback receives every tuple entry positionally, followed optionally by
+ * ``TVMFFIDefRegionKind``. Descent calls use ``this->``, matching
+ * ``StructuralWalkEngine``: lookup happens at instantiation in the Parent's
class scope and is
+ * not virtual dispatch. This deliberately leaves composed deeper-layer
declarations eligible;
+ * spelling the calls as ``Parent::member`` would instead pin lookup at that
qualified layer.
+ * A Parent must not hide other engine-internal ``this->`` members because
matching ABI-vtable
+ * paths deliberately terminate at ``StructuralMapEngineBase``.
+ *
+ * \tparam Parent Mutator layer that supplies descent and callback state
through the complete
+ * protocol above.
* \tparam order Callback placement relative to child mapping.
* \tparam Callbacks The callbacks, tested in declaration order.
*/
-template <WalkOrder order, typename... Callbacks>
-class StructuralMapMutatorObj : public StructuralMapMutatorBaseObj {
+template <typename Parent, WalkOrder order, typename... Callbacks>
+class StructuralMapEngine : public Parent {
public:
+ static_assert(std::is_base_of_v<StructuralMapEngineBase, Parent>,
+ "StructuralMap Parent must derive from
StructuralMapEngineBase");
+ /*! \brief Tuple of state references supplied by the Parent layer. */
+ using StateTupleType = typename Parent::StateTupleType;
+
/*!
* \brief Construct a callback-aware mutator that owns its callbacks.
* \param callbacks The typed callback links, tested in declaration order.
*/
- explicit StructuralMapMutatorObj(Callbacks... callbacks)
- : StructuralMapMutatorBaseObj(VTable()),
callbacks_(std::move(callbacks)...) {}
+ explicit StructuralMapEngine(Callbacks... callbacks)
+ : Parent(VTable()), callbacks_(std::move(callbacks)...) {}
private:
+ using ExpectedUnsafe = details::ExpectedUnsafe;
+ using AnyUnsafe = details::AnyUnsafe;
+
/*!
* \brief Return the shared callback-aware mutator vtable.
* \return Pointer to the immutable mutator vtable for this specialization.
*/
static const StructuralMutatorVTable* VTable() {
static const StructuralMutatorVTable vtable{
- &StructuralMapMutatorObj::DispatchMutate,
- &StructuralMapMutatorObj::DispatchMaybeInplaceMutate,
- &StructuralMapMutatorObj::DispatchVarRemapGet,
- &StructuralMapMutatorObj::DispatchVarRemapSet,
+ &StructuralMapEngine::DispatchMutate,
+ &StructuralMapEngine::DispatchMaybeInplaceMutate,
+ &StructuralMapEngine::DispatchVarRemapGet,
+ &StructuralMapEngine::DispatchVarRemapSet,
};
return &vtable;
}
@@ -884,7 +867,7 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
*/
static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
AnyView value) noexcept {
- auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
+ auto* self = static_cast<StructuralMapEngine*>(mutator);
return self->MaybeInplaceMutateImplRaw(value);
}
@@ -895,10 +878,36 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
* \return Raw ``TVMFFIAny`` containing the mutated value or Error.
*/
static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
- auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
+ auto* self = static_cast<StructuralMapEngine*>(mutator);
return self->MutateImplRaw(value);
}
+ template <typename Callback, typename Value, size_t... Is>
+ TVM_FFI_INLINE Expected<Any> InvokeTypedCallbackLink(Callback& callback,
Value&& value,
+
std::index_sequence<Is...>) noexcept {
+ using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
+ static_assert(std::is_convertible_v<typename FuncInfo::RetType,
Expected<Any>>,
+ "StructuralMap callbacks must return a replacement value,
Error, Unexpected, "
+ "or Expected<U> implicitly convertible to Expected<Any>");
+ static_assert(
+ FuncInfo::num_args == 1 + sizeof...(Is) || FuncInfo::num_args == 2 +
sizeof...(Is),
+ "StructuralMap callback takes (value, state...) with an optional
trailing "
+ "definition-region kind");
+ try {
+ static_assert(std::is_same_v<decltype(this->StateTuple()),
StateTupleType>,
+ "Parent::StateTuple() must return Parent::StateTupleType
by value");
+ StateTupleType states = this->StateTuple();
+ if constexpr (FuncInfo::num_args == 1 + sizeof...(Is)) {
+ return callback(std::forward<Value>(value), std::get<Is>(states)...);
+ } else {
+ return callback(std::forward<Value>(value), std::get<Is>(states)...,
+ this->def_region_kind());
+ }
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+ }
+
/*!
* \brief Test one link against \p value and, if it matches, mutate the node
through it.
*
@@ -911,14 +920,13 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
*/
template <bool kMaybeInplace, typename Callback>
TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value,
Expected<Any>* out) noexcept {
- using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
- static_assert(FuncInfo::num_args == 1 || FuncInfo::num_args == 2,
- "StructuralMap callbacks must take one argument (value) or
two arguments "
- "(value, def-region kind)");
+ using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
+ static_assert(FuncInfo::num_args >= 1,
+ "StructuralMap callback must take at least a value
argument");
using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
- // Deliberately duplicated by StructuralMapDynMutatorObj::TryLink in
structural_mutate.cc,
+ // Deliberately duplicated by StructuralMapDynEngine::TryLink below,
// which differs only in how a link is found and called; keep the two in
step.
//
// The match test and the matched-node path live together rather than in
separate functions:
@@ -934,11 +942,33 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
if (!matched.has_value()) return false;
}
- // A FreeVar or DAG node maps once and every later occurrence reuses that
result, so if this
- // node already has a cached remap entry, return it instead of mutating it
again.
- const bool remappable = IsRemappableIdentity(value.type_index());
+ // A final statically non-remappable type discards the remap path at
optimization time.
+ // Every other case uses runtime metadata: nullable refs may match None,
non-final subclasses
+ // may redeclare the kind, and metadata may be absent.
+ const bool remappable = [&]() {
+ if constexpr (std::is_base_of_v<ObjectRef, TSub>) {
+ using TNode = typename TSub::ContainerType;
+ if constexpr (TNode::_type_final &&
+ TNode::_type_s_eq_hash_kind != kTVMFFISEqHashKindFreeVar
&&
+ TNode::_type_s_eq_hash_kind !=
kTVMFFISEqHashKindDAGNode) {
+ return false;
+ }
+ }
+ if constexpr (std::is_pointer_v<TSub> &&
+ std::is_base_of_v<Object,
std::remove_cv_t<std::remove_pointer_t<TSub>>>) {
+ using TNode = std::remove_cv_t<std::remove_pointer_t<TSub>>;
+ constexpr bool kFinalNonRemappable =
+ TNode::_type_final && TNode::_type_s_eq_hash_kind !=
kTVMFFISEqHashKindFreeVar &&
+ TNode::_type_s_eq_hash_kind != kTVMFFISEqHashKindDAGNode;
+ if constexpr (kFinalNonRemappable) return false;
+ }
+ return this->IsRemappableIdentity(value.type_index());
+ }();
+
+ // A FreeVar or DAG node maps once and every later occurrence reuses that
result, so if
+ // this node already has a cached remap entry, return it instead of
mutating it again.
if (remappable) {
- Expected<Any> mapped = VarRemapGetExpected(value);
+ Expected<Any> mapped = this->VarRemapGetExpected(value);
if (mapped.is_err()) {
*out = std::move(mapped);
return true;
@@ -949,22 +979,22 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
}
}
- const TVMFFIDefRegionKind kind = def_region_kind();
+ using StateIndices =
std::make_index_sequence<std::tuple_size_v<StateTupleType>>;
if constexpr (order == WalkOrder::kPreOrder) {
// Pre-order: the callback rewrites this node first, then descent runs
over whatever it
// produced, so a replacement subtree is itself mapped.
Expected<Any> callback_result = [&]() -> Expected<Any> {
if constexpr (std::is_same_v<TSub, AnyView>) {
- return InvokeCallbackLink(callback, value, kind);
+ return InvokeTypedCallbackLink(callback, value, StateIndices{});
} else if constexpr (std::is_same_v<TSub, Any>) {
- return InvokeCallbackLink(callback, Any(value), kind);
+ return InvokeTypedCallbackLink(callback, Any(value), StateIndices{});
} else {
// Reuses the conversion the match already performed.
- return InvokeCallbackLink(callback, *std::move(matched), kind);
+ return InvokeTypedCallbackLink(callback, *std::move(matched),
StateIndices{});
}
}();
if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
- UpdateVisitErrorContext(callback_result, value);
+ this->UpdateVisitErrorContext(callback_result, value);
*out = std::move(callback_result);
return true;
}
@@ -980,22 +1010,22 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
if (mapped_data->type_index == input_data.type_index &&
mapped_data->zero_padding == input_data.zero_padding &&
mapped_data->v_int64 == input_data.v_int64) {
- return DefaultMaybeInplaceMutateExpected(value);
+ return this->DefaultMaybeInplaceMutateExpected(value);
}
const Object* mapped_obj = mapped_value.as<Object>();
bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
- return can_inplace ? DefaultMaybeInplaceMutateExpected(mapped_value)
- : DefaultMutateExpected(mapped_value);
+ return can_inplace ?
this->DefaultMaybeInplaceMutateExpected(mapped_value)
+ : this->DefaultMutateExpected(mapped_value);
} else {
- return DefaultMutateExpected(mapped_value);
+ return this->DefaultMutateExpected(mapped_value);
}
}();
if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
} else {
// Post-order: children are mapped first and the callback sees the
rebuilt node, so it
// observes its operands already substituted.
- Expected<Any> descended =
- kMaybeInplace ? DefaultMaybeInplaceMutateExpected(value) :
DefaultMutateExpected(value);
+ Expected<Any> descended = kMaybeInplace ?
this->DefaultMaybeInplaceMutateExpected(value)
+ :
this->DefaultMutateExpected(value);
if (TVM_FFI_PREDICT_FALSE(descended.is_err())) {
*out = std::move(descended);
return true;
@@ -1006,29 +1036,29 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
const Any& mapped_value = ExpectedUnsafe::GetData(descended);
*out = [&]() -> Expected<Any> {
if constexpr (std::is_same_v<TSub, AnyView>) {
- return InvokeCallbackLink(callback, AnyView(mapped_value), kind);
+ return InvokeTypedCallbackLink(callback, AnyView(mapped_value),
StateIndices{});
} else if constexpr (std::is_same_v<TSub, Any>) {
- return InvokeCallbackLink(callback, Any(mapped_value), kind);
+ return InvokeTypedCallbackLink(callback, Any(mapped_value),
StateIndices{});
} else {
// Re-converted rather than reusing the match: the callback is
invoked on the node
// descent handed back, and must only see the type it asked for.
Default mutation is
// required to preserve the type, so failing here means some hook
broke that.
std::optional<TSub> descended_sub = mapped_value.template as<TSub>();
if (TVM_FFI_PREDICT_FALSE(!descended_sub.has_value())) {
- return SMutateDescentTypeError();
+ return this->SMutateDescentTypeError();
}
- return InvokeCallbackLink(callback, *std::move(descended_sub), kind);
+ return InvokeTypedCallbackLink(callback, *std::move(descended_sub),
StateIndices{});
}
}();
if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
- UpdateVisitErrorContext(*out, mapped_value);
+ this->UpdateVisitErrorContext(*out, mapped_value);
return true;
}
}
// Bind this node's identity to its final result, so every later
occurrence reuses it.
if (remappable) {
- Expected<void> set_result = VarRemapSetExpected(value,
ExpectedUnsafe::GetData(*out));
+ Expected<void> set_result = this->VarRemapSetExpected(value,
ExpectedUnsafe::GetData(*out));
if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
*out = Unexpected(std::move(set_result).error());
}
@@ -1056,7 +1086,7 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
if (TryLinks<false>(value, &out, std::index_sequence_for<Callbacks...>{}))
{
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return DefaultMutateRaw(value);
+ return this->DefaultMutateRaw(value);
}
/*!
@@ -1069,14 +1099,230 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
if (TryLinks<true>(value, &out, std::index_sequence_for<Callbacks...>{})) {
return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return DefaultMaybeInplaceMutateRaw(value);
+ return this->DefaultMaybeInplaceMutateRaw(value);
}
/*! \brief The callback links, tested in declaration order. */
std::tuple<Callbacks...> callbacks_;
};
-} // namespace details
+/*!
+ * \brief Structural mutator whose links are runtime ffi.Functions keyed by
type index.
+ *
+ * This is the dynamic counterpart of \ref StructuralMapEngine. It remains a
distinct
+ * straight-line implementation because a post-order dynamic link must
remember the runtime
+ * type index selected before descent and recheck the rebuilt node against
that same index.
+ * Moving it into this header makes the same Parent layering available to
downstream dynamic
+ * mutators; only its runtime link table and invocation differ from the typed
engine.
+ * Its ``Parent`` follows the same descent-layer protocol, while runtime
``Function`` callbacks
+ * retain their existing ``(value)`` or ``(value, TVMFFIDefRegionKind)`` ABI.
+ *
+ * \tparam Parent Mutator layer that supplies descent and callback state
through the same
+ * ``this->``-bound protocol documented on \ref
StructuralMapEngine.
+ * \tparam order Callback placement relative to child mapping.
+ */
+template <typename Parent, WalkOrder order>
+class StructuralMapDynEngine : public Parent {
+ public:
+ static_assert(std::is_base_of_v<StructuralMapEngineBase, Parent>,
+ "StructuralMap Parent must derive from
StructuralMapEngineBase");
+ /*!
+ * \brief Construct a dynamic map engine with the default Parent constructor.
+ * \param callbacks Runtime links invoked as ``callback(value)``.
+ * \param callbacks_with_def_region_kind Runtime links invoked with the
active region kind.
+ */
+ StructuralMapDynEngine(Array<Tuple<int32_t, Function>> callbacks,
+ Array<Tuple<int32_t, Function>>
callbacks_with_def_region_kind)
+ : Parent(VTable()),
+ callbacks_(std::move(callbacks)),
+
callbacks_with_def_region_kind_(std::move(callbacks_with_def_region_kind)) {}
+
+ private:
+ using ExpectedUnsafe = details::ExpectedUnsafe;
+ using AnyUnsafe = details::AnyUnsafe;
+
+ /*! \brief Return the shared dynamic-engine mutator vtable. */
+ static const StructuralMutatorVTable* VTable() {
+ static const StructuralMutatorVTable vtable{
+ &StructuralMapDynEngine::DispatchMutate,
+ &StructuralMapDynEngine::DispatchMaybeInplaceMutate,
+ &StructuralMapDynEngine::DispatchVarRemapGet,
+ &StructuralMapDynEngine::DispatchVarRemapSet,
+ };
+ return &vtable;
+ }
+
+ /*! \brief Dispatch optional in-place mutation through the ABI vtable. */
+ static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
+ AnyView value) noexcept {
+ return
static_cast<StructuralMapDynEngine*>(mutator)->MaybeInplaceMutateImplRaw(value);
+ }
+
+ /*! \brief Dispatch non-in-place mutation through the ABI vtable. */
+ static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
+ return static_cast<StructuralMapDynEngine*>(mutator)->MutateImplRaw(value);
+ }
+
+ /*!
+ * \brief Find the first runtime link registered for \p type_index.
+ * \param type_index The input node's runtime type index.
+ * \param with_kind Set when the matched link also takes a def-region kind.
+ * \param link_type_index Set to the registered type index the link matched
on, so post-order
+ * traversal can recheck the descended node against the same target.
+ * \return The matched Function, or nullopt when no link applies.
+ */
+ Optional<Function> FindLink(int32_t type_index, bool* with_kind,
+ int32_t* link_type_index) const noexcept {
+ for (const Tuple<int32_t, Function>& entry : callbacks_) {
+ if (details::RuntimeTypeIndexMatch(type_index, entry.get<0>())) {
+ *with_kind = false;
+ *link_type_index = entry.get<0>();
+ return entry.get<1>();
+ }
+ }
+ for (const Tuple<int32_t, Function>& entry :
callbacks_with_def_region_kind_) {
+ if (details::RuntimeTypeIndexMatch(type_index, entry.get<0>())) {
+ *with_kind = true;
+ *link_type_index = entry.get<0>();
+ return entry.get<1>();
+ }
+ }
+ return std::nullopt;
+ }
+
+ /*!
+ * \brief Invoke a matched runtime link with its requested arguments.
+ *
+ * The caller reads the live def-region kind at invocation time.
``CallExpected`` uses the
+ * exception-free safe-call path and represents raised errors as
``Unexpected``.
+ */
+ TVM_FFI_INLINE static Expected<Any> InvokeLink(const Function& fn, bool
with_kind, AnyView target,
+ TVMFFIDefRegionKind kind)
noexcept {
+ return with_kind ? fn.CallExpected<Any>(target, kind) :
fn.CallExpected<Any>(target);
+ }
+
+ /*!
+ * \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.
+ * \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>
+ TVM_FFI_INLINE bool TryLink(AnyView value, Expected<Any>* out) noexcept {
+ // Step for step the same walk as StructuralMapEngine::TryLink, and
deliberately so: only
+ // link detection and invocation differ. Everything below except finding
and calling the
+ // link is shared semantics, so a change to either copy belongs in both.
+ bool with_kind = false;
+ int32_t link_type_index = TypeIndex::kTVMFFINone;
+ // A local, so descending into a matching child cannot change what this
node invokes.
+ Optional<Function> matched = FindLink(value.type_index(), &with_kind,
&link_type_index);
+ if (!matched.has_value()) return false;
+
+ // --- identity remap, entry half -----------------------------------------
+ // A FreeVar or DAG node maps once and every later occurrence reuses that
result.
+ const bool remappable = this->IsRemappableIdentity(value.type_index());
+ if (remappable) {
+ Expected<Any> mapped = this->VarRemapGetExpected(value);
+ if (mapped.is_err()) {
+ *out = std::move(mapped);
+ return true;
+ }
+ if (ExpectedUnsafe::GetData(mapped).type_index() !=
TypeIndex::kTVMFFINone) {
+ *out = std::move(mapped);
+ return true;
+ }
+ }
+
+ // --- callback and descent, in walk order --------------------------------
+ if constexpr (order == WalkOrder::kPreOrder) {
+ // Pre-order: the callback rewrites this node first, then descent runs
over what it made.
+ Expected<Any> callback_result =
+ InvokeLink(*matched, with_kind, value, this->def_region_kind());
+ if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
+ this->UpdateVisitErrorContext(callback_result, value);
+ *out = std::move(callback_result);
+ return true;
+ }
+ Any mapped_value = ExpectedUnsafe::GetData(callback_result);
+ *out = [&]() -> Expected<Any> {
+ if constexpr (kMaybeInplace) {
+ const TVMFFIAny* mapped_data =
AnyUnsafe::TVMFFIAnyPtrFromAny(mapped_value);
+ const TVMFFIAny input_data = value.CopyToTVMFFIAny();
+ if (mapped_data->type_index == input_data.type_index &&
+ mapped_data->zero_padding == input_data.zero_padding &&
+ mapped_data->v_int64 == input_data.v_int64) {
+ return this->DefaultMaybeInplaceMutateExpected(value);
+ }
+ const Object* mapped_obj = mapped_value.as<Object>();
+ bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
+ return can_inplace ?
this->DefaultMaybeInplaceMutateExpected(mapped_value)
+ : this->DefaultMutateExpected(mapped_value);
+ } else {
+ return this->DefaultMutateExpected(mapped_value);
+ }
+ }();
+ if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
+ } else {
+ // Post-order: children are mapped first, so the callback sees the
rebuilt node.
+ Expected<Any> descended = kMaybeInplace ?
this->DefaultMaybeInplaceMutateExpected(value)
+ :
this->DefaultMutateExpected(value);
+ if (TVM_FFI_PREDICT_FALSE(descended.is_err())) {
+ *out = std::move(descended);
+ return true;
+ }
+ const Any& mapped_value = ExpectedUnsafe::GetData(descended);
+ // Selection used the input node. Recheck the descended node against
that same registered
+ // target before invoking the saved link.
+ if (TVM_FFI_PREDICT_FALSE(
+ !details::RuntimeTypeIndexMatch(mapped_value.type_index(),
link_type_index))) {
+ *out = this->SMutateDescentTypeError();
+ this->UpdateVisitErrorContext(*out, mapped_value);
+ return true;
+ }
+ // WithDefRegionKind restores its state through RAII, so this late read
is equivalent to
+ // the typed engine's invocation-time read even after recursive descent.
+ *out = InvokeLink(*matched, with_kind, mapped_value,
this->def_region_kind());
+ if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
+ this->UpdateVisitErrorContext(*out, mapped_value);
+ return true;
+ }
+ }
+
+ // --- identity remap, exit half ------------------------------------------
+ // Bind this node's identity to its final result for later occurrences.
+ if (remappable) {
+ Expected<void> set_result = this->VarRemapSetExpected(value,
ExpectedUnsafe::GetData(*out));
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ *out = Unexpected(std::move(set_result).error());
+ }
+ }
+ return true;
+ }
+
+ /*! \brief Mutate a value, invoking the first matching runtime link. */
+ TVM_FFI_INLINE TVMFFIAny MutateImplRaw(AnyView value) noexcept {
+ Expected<Any> out{Any()};
+ if (TryLink<false>(value, &out)) {
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
+ }
+ return this->DefaultMutateRaw(value);
+ }
+
+ /*! \brief Optionally mutate a value in place through the first matching
runtime link. */
+ TVM_FFI_INLINE TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
+ Expected<Any> out{Any()};
+ if (TryLink<true>(value, &out)) {
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
+ }
+ return this->DefaultMaybeInplaceMutateRaw(value);
+ }
+
+ /*! \brief Runtime links invoked without def-region context. */
+ Array<Tuple<int32_t, Function>> callbacks_;
+ /*! \brief Runtime links invoked with def-region context. */
+ Array<Tuple<int32_t, Function>> callbacks_with_def_region_kind_;
+};
/*!
* \brief Map a structured value graph and invoke typed replacement callbacks.
@@ -1133,7 +1379,7 @@ class StructuralMapMutatorObj : public
StructuralMapMutatorBaseObj {
template <WalkOrder order, typename... Callbacks>
Expected<Any> StructuralMapExpected(AnyView root, Callbacks&&... callbacks)
noexcept {
static_assert(sizeof...(Callbacks) != 0, "StructuralMap requires at least
one callback");
- using Mutator = details::StructuralMapMutatorObj<order,
std::decay_t<Callbacks>...>;
+ using Mutator = StructuralMapEngine<StructuralMapEngineBase, order,
std::decay_t<Callbacks>...>;
StructuralMutator
mutator(make_object<Mutator>(std::forward<Callbacks>(callbacks)...));
return mutator->MaybeInplaceMutateIfUniqueExpected(root);
}
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index 4627d3a7..f094cafc 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -47,223 +47,15 @@ namespace details {
* \param order Integer value of \ref WalkOrder.
* \return The mapped owning value, or an Error.
*/
-/*!
- * \brief Structural mutator whose links are runtime ffi.Functions keyed by
type index.
- *
- * The dynamic counterpart of \ref StructuralMapMutatorObj. It selects a link
by scanning
- * registered type indices where the static version does a compile-time
``as<TSub>()``. The two
- * are kept apart, rather than sharing one template with a mode flag, so each
version reads
- * straight through; only the identity remap they share lives in the common
base. Keeping this
- * one in the .cc also keeps the link table out of the public header.
- *
- * \tparam order Callback placement relative to child mapping.
- */
-template <WalkOrder order>
-class StructuralMapDynMutatorObj : public StructuralMapMutatorBaseObj {
- public:
- StructuralMapDynMutatorObj(Array<Tuple<int32_t, Function>> callbacks,
- Array<Tuple<int32_t, Function>>
callbacks_with_def_region_kind)
- : StructuralMapMutatorBaseObj(VTable()),
- callbacks_(std::move(callbacks)),
-
callbacks_with_def_region_kind_(std::move(callbacks_with_def_region_kind)) {}
-
- private:
- static const StructuralMutatorVTable* VTable() {
- static const StructuralMutatorVTable vtable{
- &StructuralMapDynMutatorObj::DispatchMutate,
- &StructuralMapDynMutatorObj::DispatchMaybeInplaceMutate,
- &StructuralMapDynMutatorObj::DispatchVarRemapGet,
- &StructuralMapDynMutatorObj::DispatchVarRemapSet,
- };
- return &vtable;
- }
-
- static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
- AnyView value) noexcept {
- return
static_cast<StructuralMapDynMutatorObj*>(mutator)->MaybeInplaceMutateImplRaw(value);
- }
-
- static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
- return
static_cast<StructuralMapDynMutatorObj*>(mutator)->MutateImplRaw(value);
- }
-
- /*!
- * \brief Find the first link registered for \p type_index.
- *
- * \param type_index The input node's runtime type index.
- * \param with_kind Set when the matched link also takes a def-region kind.
- * \param link_type_index Set to the registered type index the link matched
on, so a post-order
- * walk can recheck the descended node against the same target.
- * \return The matched Function, or nullopt when no link applies.
- */
- Optional<Function> FindLink(int32_t type_index, bool* with_kind,
- int32_t* link_type_index) const noexcept {
- for (const Tuple<int32_t, Function>& entry : callbacks_) {
- if (RuntimeTypeIndexMatch(type_index, entry.get<0>())) {
- *with_kind = false;
- *link_type_index = entry.get<0>();
- return entry.get<1>();
- }
- }
- for (const Tuple<int32_t, Function>& entry :
callbacks_with_def_region_kind_) {
- if (RuntimeTypeIndexMatch(type_index, entry.get<0>())) {
- *with_kind = true;
- *link_type_index = entry.get<0>();
- return entry.get<1>();
- }
- }
- return std::nullopt;
- }
-
- /*!
- * \brief Invoke a matched link, threading the live def-region kind through.
- *
- * \p kind is passed rather than stashed at selection time: the engine only
knows the true
- * def-region kind of the node handed back after descent (post-order) or of
the matched node
- * itself (pre-order).
- */
- TVM_FFI_INLINE static Expected<Any> InvokeLink(const Function& fn, bool
with_kind, AnyView target,
- TVMFFIDefRegionKind kind)
noexcept {
- // CallExpected is exception-free: it goes through the safe-call path and
returns any raised
- // error as Unexpected, unlike a directly invoked C++ callback.
- return with_kind ? fn.CallExpected<Any>(target, kind) :
fn.CallExpected<Any>(target);
- }
-
- /*!
- * \brief Test the link table against \p value and mutate through the first
match.
- *
- * \tparam kMaybeInplace 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>
- TVM_FFI_INLINE bool TryLink(AnyView value, Expected<Any>* out) noexcept {
- // Step for step the same walk as StructuralMapMutatorObj::TryLink, and
deliberately so:
- // only link detection and invocation differ between the two, and keeping
them as separate
- // straight-line copies lets each specialize on its own selection strategy
and keeps both
- // readable. Everything below except finding and calling the link is
shared semantics, so a
- // change to either copy belongs in both.
-
- bool with_kind = false;
- int32_t link_type_index = TypeIndex::kTVMFFINone;
- // A local, so descending into a matching child cannot change what this
node invokes.
- Optional<Function> matched = FindLink(value.type_index(), &with_kind,
&link_type_index);
- if (!matched.has_value()) return false;
-
- // --- identity remap, entry half -----------------------------------------
- // A FreeVar or DAG node maps once and every later occurrence reuses that
result, so if this
- // node already has a cached remap entry, return it instead of mutating it
again.
- const bool remappable = IsRemappableIdentity(value.type_index());
- if (remappable) {
- Expected<Any> mapped = VarRemapGetExpected(value);
- if (mapped.is_err()) {
- *out = std::move(mapped);
- return true;
- }
- if (ExpectedUnsafe::GetData(mapped).type_index() !=
TypeIndex::kTVMFFINone) {
- *out = std::move(mapped);
- return true;
- }
- }
-
- // --- callback and descent, in walk order --------------------------------
- const TVMFFIDefRegionKind kind = def_region_kind();
- if constexpr (order == WalkOrder::kPreOrder) {
- // Pre-order: the callback rewrites this node first, then descent runs
over what it made.
- Expected<Any> callback_result = InvokeLink(*matched, with_kind, value,
kind);
- if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
- UpdateVisitErrorContext(callback_result, value);
- *out = std::move(callback_result);
- return true;
- }
- // Own the callback's result: it is the only reference from here on.
- Any mapped_value = ExpectedUnsafe::GetData(callback_result);
- // Each descent names the node it actually ran on in the error context.
- *out = [&]() -> Expected<Any> {
- if constexpr (kMaybeInplace) {
- const TVMFFIAny* mapped_data =
AnyUnsafe::TVMFFIAnyPtrFromAny(mapped_value);
- const TVMFFIAny input_data = value.CopyToTVMFFIAny();
- if (mapped_data->type_index == input_data.type_index &&
- mapped_data->zero_padding == input_data.zero_padding &&
- mapped_data->v_int64 == input_data.v_int64) {
- return DefaultMaybeInplaceMutateExpected(value);
- }
- const Object* mapped_obj = mapped_value.as<Object>();
- bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
- return can_inplace ? DefaultMaybeInplaceMutateExpected(mapped_value)
- : DefaultMutateExpected(mapped_value);
- } else {
- return DefaultMutateExpected(mapped_value);
- }
- }();
- if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
- } else {
- // Post-order: children are mapped first, so the callback sees the
rebuilt node.
- Expected<Any> descended =
- kMaybeInplace ? DefaultMaybeInplaceMutateExpected(value) :
DefaultMutateExpected(value);
- if (TVM_FFI_PREDICT_FALSE(descended.is_err())) {
- *out = std::move(descended);
- return true;
- }
- // Held by reference, not moved out: the error path below names this
node, so it has to
- // survive the callback.
- const Any& mapped_value = ExpectedUnsafe::GetData(descended);
- // The link was selected on the input node, and the callback must only
see the type it
- // registered for. The typed mutator gets this from its
`mapped_value.as<TSub>()`; here the
- // registered type index is the same target, so recheck against it.
- if (TVM_FFI_PREDICT_FALSE(
- !RuntimeTypeIndexMatch(mapped_value.type_index(),
link_type_index))) {
- *out = SMutateDescentTypeError();
- UpdateVisitErrorContext(*out, mapped_value);
- return true;
- }
- *out = InvokeLink(*matched, with_kind, mapped_value, kind);
- if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
- UpdateVisitErrorContext(*out, mapped_value);
- return true;
- }
- }
-
- // Bind this node's identity to its final result, so every later
occurrence reuses it.
- if (remappable) {
- Expected<void> set_result = VarRemapSetExpected(value,
ExpectedUnsafe::GetData(*out));
- if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
- *out = Unexpected(std::move(set_result).error());
- }
- }
- return true;
- }
-
- TVM_FFI_INLINE TVMFFIAny MutateImplRaw(AnyView value) noexcept {
- Expected<Any> out{Any()};
- if (TryLink<false>(value, &out)) {
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
- }
- return DefaultMutateRaw(value);
- }
-
- TVM_FFI_INLINE TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
- Expected<Any> out{Any()};
- if (TryLink<true>(value, &out)) {
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
- }
- return DefaultMaybeInplaceMutateRaw(value);
- }
-
- Array<Tuple<int32_t, Function>> callbacks_;
- Array<Tuple<int32_t, Function>> callbacks_with_def_region_kind_;
-};
-
Expected<Any> StructuralMapExpected(
AnyView root, const Array<Tuple<int32_t, Function>>& callbacks,
const Array<Tuple<int32_t, Function>>& callbacks_with_def_region_kind, int
order) noexcept {
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
- using Mutator = StructuralMapDynMutatorObj<WalkOrder::kPreOrder>;
+ using Mutator = StructuralMapDynEngine<StructuralMapEngineBase,
WalkOrder::kPreOrder>;
StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
return mutator->MaybeInplaceMutateIfUniqueExpected(root);
} else {
- using Mutator = StructuralMapDynMutatorObj<WalkOrder::kPostOrder>;
+ using Mutator = StructuralMapDynEngine<StructuralMapEngineBase,
WalkOrder::kPostOrder>;
StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
return mutator->MaybeInplaceMutateIfUniqueExpected(root);
}
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 162cb38d..8e77a8d9 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -44,6 +44,98 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TMutatePairObj::RegisterReflection(); }
Expected<Any> Increment(int64_t value) { return Any(value + 1); }
+struct MutateCount {
+ int value = 0;
+ int mutate_raw = 0;
+ int mutate_expected = 0;
+ int maybe_inplace_raw = 0;
+ int maybe_inplace_expected = 0;
+};
+
+class StructuralMapWithMutateCount : public StructuralMapEngineBase {
+ public:
+ using StateTupleType = std::tuple<const MutateCount&, const int&>;
+
+ explicit StructuralMapWithMutateCount(const StructuralMutatorVTable* vtable)
+ : StructuralMapEngineBase(vtable) {}
+
+ const MutateCount& count() const { return count_; }
+
+ Expected<Any> DefaultMutateExpected(AnyView value) noexcept {
+ ++count_.value;
+ ++count_.mutate_expected;
+ return StructuralMapEngineBase::DefaultMutateExpected(value);
+ }
+
+ Expected<Any> DefaultMaybeInplaceMutateExpected(AnyView value) noexcept {
+ ++count_.value;
+ ++count_.maybe_inplace_expected;
+ return StructuralMapEngineBase::DefaultMaybeInplaceMutateExpected(value);
+ }
+
+ protected:
+ TVMFFIAny DefaultMutateRaw(AnyView value) noexcept {
+ ++count_.mutate_raw;
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(DefaultMutateExpected(value));
+ }
+
+ TVMFFIAny DefaultMaybeInplaceMutateRaw(AnyView value) noexcept {
+ ++count_.maybe_inplace_raw;
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(DefaultMaybeInplaceMutateExpected(value));
+ }
+
+ StateTupleType StateTuple() const noexcept { return StateTupleType(count_,
marker_); }
+
+ private:
+ MutateCount count_;
+ int marker_ = 17;
+};
+
+TEST(StructuralMap, ParentLayerOwnsBothDescentsAndProvidesState) {
+ std::vector<int> callback_counts;
+ auto identity = [&](const AnyArray& value, const MutateCount& live_count,
const int& live_marker,
+ TVMFFIDefRegionKind kind) -> Expected<Any> {
+ EXPECT_EQ(live_marker, 17);
+ EXPECT_EQ(kind, kTVMFFIDefRegionKindNone);
+ callback_counts.push_back(live_count.value);
+ return Any(value);
+ };
+ int var_callback_count = 0;
+ auto map_var = [&](const TVarObj* value, const MutateCount& live_count,
+ const int& live_marker) -> Expected<Any> {
+ EXPECT_EQ(live_marker, 17);
+ EXPECT_GT(live_count.value, 0);
+ ++var_callback_count;
+ return Any(TVar(value->name + "-mapped"));
+ };
+ using Mutator = StructuralMapEngine<StructuralMapWithMutateCount,
WalkOrder::kPostOrder,
+ decltype(identity), decltype(map_var)>;
+ auto engine = make_object<Mutator>(std::move(identity), std::move(map_var));
+ StructuralMutator mutator(engine);
+
+ ASSERT_FALSE(mutator->MutateExpected(String("unmatched")).is_err());
+ AnyArray rebuild_root{int64_t{1}};
+ ASSERT_FALSE(mutator->MutateExpected(rebuild_root).is_err());
+
+
ASSERT_FALSE(mutator->MaybeInplaceMutateExpected(String("unmatched")).is_err());
+ AnyArray inplace_root{int64_t{1}};
+ ASSERT_FALSE(mutator->MaybeInplaceMutateExpected(inplace_root).is_err());
+
+ EXPECT_GT(engine->count().mutate_raw, 0);
+ EXPECT_GT(engine->count().mutate_expected, 0);
+ EXPECT_GT(engine->count().maybe_inplace_raw, 0);
+ EXPECT_GT(engine->count().maybe_inplace_expected, 0);
+ EXPECT_EQ(callback_counts.size(), 2U);
+ EXPECT_GT(callback_counts[0], 0);
+ EXPECT_GT(callback_counts[1], callback_counts[0]);
+
+ TVar var("n");
+ AnyArray repeated{var, var};
+ AnyArray mapped = mutator->MutateExpected(repeated).value().cast<AnyArray>();
+ EXPECT_EQ(var_callback_count, 1);
+ EXPECT_TRUE(mapped[0].cast<TVar>().same_as(mapped[1].cast<TVar>()));
+}
+
template <WalkOrder order>
void CheckNestedArrayMapOrder(const std::vector<std::string>& expected_trace) {
AnyArray inner_array{int64_t{1}};
@@ -134,6 +226,16 @@ TEST(StructuralMap,
RegisteredMutateHookUsesAssignOrReturn) {
EXPECT_FALSE(nullable_mapped->lhs.defined());
EXPECT_TRUE(nullable_mapped->rhs.same_as(rhs));
+ int nullable_var_callbacks = 0;
+ nullable_mapped =
+ StructuralMap<WalkOrder::kPostOrder>(nullable, [&](const TVar& value) ->
Expected<Any> {
+ ++nullable_var_callbacks;
+ return value.defined() ? Any(value) : Any(ObjectRef(nullptr));
+ }).cast<TMutatePair>();
+ EXPECT_EQ(nullable_var_callbacks, 2);
+ EXPECT_FALSE(nullable_mapped->lhs.defined());
+ EXPECT_TRUE(nullable_mapped->rhs.same_as(rhs));
+
Expected<Any> wrong_type = StructuralMapExpected<WalkOrder::kPostOrder>(
root, [](const TVarObj*) -> Expected<Any> { return Any(int64_t{1}); });
ASSERT_TRUE(wrong_type.is_err());
@@ -426,4 +528,28 @@ TEST(StructuralMapDyn, ReusesRemapResultForRepeatedVar) {
EXPECT_TRUE(arr[0].cast<TVar>().same_as(arr[1].cast<TVar>()));
}
+template <WalkOrder order>
+void CheckDynamicParentLayer() {
+ int64_t calls = 0;
+ Function increment = Function::FromTyped([&](int64_t value) -> Any {
+ ++calls;
+ return Any(value + 1);
+ });
+ using Mutator = StructuralMapDynEngine<StructuralMapWithMutateCount, order>;
+ auto engine = make_object<Mutator>(
+ Array<Tuple<int32_t, Function>>{Tuple<int32_t,
Function>(TypeIndex::kTVMFFIInt, increment)},
+ Array<Tuple<int32_t, Function>>());
+ StructuralMutator mutator(engine);
+
+ AnyArray mapped = mutator->Mutate(AnyArray{int64_t{1}}).cast<AnyArray>();
+ EXPECT_EQ(mapped[0].cast<int64_t>(), 2);
+ EXPECT_EQ(calls, 1);
+ EXPECT_GT(engine->count().value, 0);
+}
+
+TEST(StructuralMapDyn, ParentLayerRunsThroughHeaderDefinedEngine) {
+ CheckDynamicParentLayer<WalkOrder::kPreOrder>();
+ CheckDynamicParentLayer<WalkOrder::kPostOrder>();
+}
+
} // namespace
