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 214ea3c4 [REFACTOR][EXTRA] Add StructuralVisit, a callback-driven
traversal primitive (#750)
214ea3c4 is described below
commit 214ea3c4e25c685ff7e05d4a7997ca9f5a4d89ac
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 5 19:43:56 2026 -0400
[REFACTOR][EXTRA] Add StructuralVisit, a callback-driven traversal
primitive (#750)
`StructuralWalk` owns descent, so callbacks cannot choose individual
children or control traversal order. Context-carrying passes need a
lower-level primitive where the callback owns the subtree.
This adds a separate public `StructuralVisitEngine<Parent,
Callbacks...>` with declaration-ordered typed dispatch. A matched
callback receives the active `Parent::VisitorObjType*`, owns descent,
and returns the final result for that value; an unmatched value follows
registered or reflected default traversal from the parent layer.
High-level changes:
- add `StructuralVisitExpected` and `StructuralVisit` with exact
`(value, visitor)` callbacks;
- expose callback-owned visiting and default descent through the dynamic
and Python APIs;
- attach visit error context once at each walk/visit engine boundary,
including reflected fallback;
- support raw and typed `Expected` early-return contexts;
- document and test pruning, interruption, errors, and composed visitor
layers.
---
docs/concepts/structural_eq_hash.rst | 48 ++++-
docs/reference/python/index.rst | 1 +
include/tvm/ffi/expected.h | 29 +++
include/tvm/ffi/extra/structural_mutate.h | 47 ++---
include/tvm/ffi/extra/structural_visit.h | 297 ++++++++++++++++++++++--------
python/tvm_ffi/__init__.py | 2 +
python/tvm_ffi/_ffi_api.py | 4 +
python/tvm_ffi/structural.py | 62 +++++++
src/ffi/extra/structural_visit.cc | 33 ++++
tests/cpp/extra/test_structural_visit.cc | 84 +++++++++
tests/python/test_structural.py | 70 +++++++
11 files changed, 564 insertions(+), 113 deletions(-)
diff --git a/docs/concepts/structural_eq_hash.rst
b/docs/concepts/structural_eq_hash.rst
index 6b998918..b7182560 100644
--- a/docs/concepts/structural_eq_hash.rst
+++ b/docs/concepts/structural_eq_hash.rst
@@ -22,8 +22,9 @@ TVM FFI provides ``structural_equal`` and ``structural_hash``
for the
object graph. These compare objects by **content** — recursively walking
fields — rather than by pointer identity.
-The same reflection metadata also drives ``structural_walk`` for analyses and
-``structural_map`` for rewrites. Their low-level engines,
+The same reflection metadata also drives ``structural_walk`` for analyses,
+``structural_visit`` for callback-owned descent, and ``structural_map`` for
rewrites. Their
+low-level engines,
``StructuralVisitor`` and ``StructuralMutator``, let custom object types
participate in the same traversal protocol.
@@ -1016,6 +1017,9 @@ There are two layers of API:
* - :func:`~tvm_ffi.structural_walk`
- Inspect a value graph without replacing values
- Collect information, validate IR, or stop at a match
+ * - :func:`~tvm_ffi.structural_visit`
+ - Give each matching callback control over child traversal
+ - Visit selected children in a chosen order or definition scope
* - :func:`~tvm_ffi.structural_map`
- Recursively replace values and rebuild changed paths
- Rewriting and compiler optimization passes
@@ -1026,11 +1030,11 @@ There are two layers of API:
- Low-level recursive mutation engine
- Implementing custom mutation hooks and identity substitution
-``structural_walk`` and ``structural_map`` construct the corresponding
low-level
-object, install callback-aware dispatch, run it on the root, and return the
final
-result. Applications normally use these two functions directly. Custom object
-hooks receive the low-level visitor or mutator so that recursive calls remain
in
-the same traversal.
+``structural_walk``, ``structural_visit``, and ``structural_map`` construct the
+corresponding low-level object, install callback-aware dispatch, run it on the
root, and
+return the final result. Applications normally use these functions directly.
Custom object
+hooks receive the low-level visitor or mutator so that recursive calls remain
in the same
+traversal.
StructuralVisitor and StructuralMutator
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1040,10 +1044,21 @@ definition-region kind, and any early-interruption
state. Its main operations
are:
- ``visitor.visit(value)`` visits a child with the same visitor.
+- ``visitor.default_visit(value)`` bypasses the active engine callback for that
+ value but still dispatches its registered ``__s_visit__`` hook.
- ``visitor.def_region_kind()`` reports the active definition-region kind.
- ``visitor.with_def_region_kind(kind, callback)`` temporarily changes that
kind
while ``callback`` performs recursive visits.
+.. warning::
+
+ A ``__s_visit__`` hook must not call ``default_visit`` on the same value
+ currently being visited. Doing so re-enters that hook without a recursion
+ guard, causing stack overflow and a process crash. Use ``default_visit`` on
+ a child whose default traversal is wanted. It is also safe for a
+ ``structural_visit`` engine callback to call ``default_visit`` on its
matched
+ value; that bypasses engine callback dispatch for the value.
+
The default visitor dispatches to a type's ``__s_visit__`` hook when present.
Otherwise POD values are leaves and object-backed values are visited through
their reflected structural fields. Array and List have built-in hooks that
@@ -1262,8 +1277,23 @@ type and accepts an optional second
``TVMFFIDefRegionKind`` argument. The
return FoldAdd(add);
});
+``StructuralVisitExpected`` is the callback-driven form. A matched callback
+receives the active visitor, owns descent into its value, and returns the final
+result for that subtree. Returning without calling the visitor prunes the
+subtree. An unmatched value uses default descent:
+
+.. code-block:: cpp
+
+ Expected<Optional<VisitInterrupt>> result = StructuralVisitExpected(
+ root,
+ [&](const Pair& pair, StructuralVisitorObj* visitor)
+ -> Expected<Optional<VisitInterrupt>> {
+ // The callback owns descent: visit lhs, never visit rhs.
+ return visitor->VisitExpected(pair->lhs);
+ });
+
Walk callbacks return ``Expected<WalkResult>``. Map callbacks return
``Expected<Any>`` and must obey the same non-in-place callback contract as the
Python API. For ``Map`` and ``Dict``, both APIs process values and skip keys.
-``StructuralWalk`` and ``StructuralMap`` are the corresponding throwing
-convenience forms.
+``StructuralWalk``, ``StructuralVisit`` and ``StructuralMap`` are the
+corresponding throwing convenience forms.
diff --git a/docs/reference/python/index.rst b/docs/reference/python/index.rst
index 80a66725..57a2be9f 100644
--- a/docs/reference/python/index.rst
+++ b/docs/reference/python/index.rst
@@ -92,6 +92,7 @@ Structural
structural_equal
structural_hash
structural_walk
+ structural_visit
structural_map
diff --git a/include/tvm/ffi/expected.h b/include/tvm/ffi/expected.h
index 6ebd9ab5..6bf6051f 100644
--- a/include/tvm/ffi/expected.h
+++ b/include/tvm/ffi/expected.h
@@ -431,6 +431,35 @@ struct ExpectedUnsafe {
}
};
+/*!
+ * \brief Return proxy used by early-return macros in raw or typed functions.
+ * \tparam T The success type, fixed when the proxy stores its ``Expected<T>``.
+ *
+ * A return statement selects either the raw ``TVMFFIAny`` conversion used by
+ * hooks or the same ``Expected<T>`` type used by typed helpers. Both
+ * conversions are rvalue-qualified because handing off the stored payload is
+ * a single move; an lvalue helper cannot accidentally transfer it twice. As
+ * with other moved-from values, deliberately converting ``std::move(helper)``
+ * twice remains caller error.
+ */
+template <typename T>
+class MaybeReturnHelper {
+ public:
+ TVM_FFI_INLINE explicit MaybeReturnHelper(Expected<T>&& value) noexcept
+ : value_(std::move(value)) {}
+
+ // NOLINTNEXTLINE(google-explicit-constructor,runtime/explicit)
+ TVM_FFI_INLINE operator TVMFFIAny() && noexcept {
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(value_));
+ }
+
+ // NOLINTNEXTLINE(google-explicit-constructor,runtime/explicit)
+ TVM_FFI_INLINE operator Expected<T>() && noexcept { return
std::move(value_); }
+
+ private:
+ Expected<T> value_;
+};
+
} // namespace details
// TypeTraits specialization for Expected<T>
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index b3ab3b13..a1b9c8ea 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -141,6 +141,9 @@ struct StructuralMutatorVTable {
*/
class StructuralMutatorObj : public Object {
public:
+ /*! \brief Callback-facing mutator type used by composed callback-driven
engines. */
+ using MutatorObjType = StructuralMutatorObj;
+
/*!
* \brief Mutate a value through the mutator vtable.
*
@@ -598,16 +601,14 @@ TVM_FFI_INLINE static Expected<Any>
MutateReflectedFieldsExpected(StructuralMuta
namespace details {
/// \cond Doxygen_Suppress
-// Return from the current mutation function if Result is an Error.
-// Append Node to the mutate error context before returning. Node is required:
dropping it
-// silently degrades every error message produced below this frame.
-// A raw pointer Node must be non-null; pass nullable nodes as ObjectRef or
Any so None is skipped.
-#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result)
\
- do {
\
- auto&& tvm_ffi_res_ = (Result);
\
- if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.is_err())) {
\
- return
::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(::std::move(tvm_ffi_res_));
\
- }
\
+// Return from the current raw or same-T Expected mutation function if Result
is an Error.
+// The rvalue-only proxy lets the enclosing return type select the
representation.
+#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result)
\
+ do {
\
+ auto&& tvm_ffi_res_ = (Result);
\
+ if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.is_err())) {
\
+ return
::tvm::ffi::details::MaybeReturnHelper(::std::move(tvm_ffi_res_)); \
+ }
\
} while (0)
/// \endcond
@@ -615,9 +616,9 @@ namespace details {
// Out of line so its strings and Error construction stay out of the hot path
of whatever hook
// body TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN expands into. Same reason as
// BadStructuralMutateHookError.
-TVM_FFI_COLD_CODE inline TVMFFIAny SMutateDeclaredTypeErrorRaw() noexcept {
- return AnyUnsafe::MoveAnyToTVMFFIAny(
- Any(Error("TypeError", "structural mutate result does not match the
declared type", "")));
+TVM_FFI_COLD_CODE inline Expected<Any> SMutateDeclaredTypeError() noexcept {
+ return Unexpected(
+ Error("TypeError", "structural mutate result does not match the declared
type", ""));
}
/// \cond Doxygen_Suppress
@@ -626,7 +627,8 @@ TVM_FFI_COLD_CODE inline TVMFFIAny
SMutateDeclaredTypeErrorRaw() noexcept {
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result);
\
if
(TVM_FFI_PREDICT_FALSE(!::tvm::ffi::details::AnyUnsafe::CheckAnyStrict<Type>( \
::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))) {
\
- return ::tvm::ffi::details::SMutateDeclaredTypeErrorRaw();
\
+ return ::tvm::ffi::details::MaybeReturnHelper(
\
+ ::tvm::ffi::details::SMutateDeclaredTypeError());
\
}
\
Type Name = /* NOLINT(bugprone-macro-parentheses) */
\
::tvm::ffi::details::AnyUnsafe::MoveFromAnyAfterCheck<Type>(
\
@@ -637,16 +639,15 @@ TVM_FFI_COLD_CODE inline TVMFFIAny
SMutateDeclaredTypeErrorRaw() noexcept {
* \brief Unwrap a successful mutation result into a newly declared value or
return its error.
*
* ``Type`` must be concrete; use a type alias when it contains a top-level
comma. A type mismatch
- * returns ``Unexpected(TypeError)`` through the surrounding ``Expected``
function without
- * throwing, reported with a fixed string so a correct hook pays only one
predicted-not-taken
- * branch per field. This macro declares ``Name`` into the enclosing scope and
must be used in a
- * braced block, never as an unbraced control-flow body. A raw pointer node
must be non-null; pass
- * nullable nodes as ``ObjectRef`` or ``Any`` so ``None`` is skipped when
constructing error
- * context.
+ * returns ``TypeError`` through the surrounding raw or ``Expected`` function
without throwing,
+ * reported with a fixed string so a correct hook pays only one
predicted-not-taken branch per
+ * field. Its early returns work from either a raw ``TVMFFIAny`` hook or a
same-T ``Expected<T>``
+ * helper. 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}
- * TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, child,
mutator->MutateExpected(self->child), self);
+ * TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, child,
mutator->MutateExpected(self->child));
* \endcode
*
* \param Type The concrete type of the successful value.
@@ -672,8 +673,8 @@ TVM_FFI_COLD_CODE inline TVMFFIAny
SMutateDeclaredTypeErrorRaw() noexcept {
/*!
* \brief \ref TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN without the type check.
*
- * Same signature and same error propagation; the difference is only what
happens to a
- * successful result that is not of type \p Type.
+ * Same signature, raw-or-same-T early-return support, and error propagation;
the difference is
+ * only what happens to a successful result that is not of type \p Type.
*
* The caller must guarantee the result has the declared type; a mismatch is
undefined behavior
* in a release build, and debug builds catch it with ``TVM_FFI_DCHECK``.
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index e8d567ad..b8d41948 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -138,6 +138,8 @@ struct StructuralVisitorVTable {
*/
class StructuralVisitorObj : public Object {
public:
+ /*! \brief Callback-facing visitor type used by composed callback-driven
engines. */
+ using VisitorObjType = StructuralVisitorObj;
/*! \brief State references made available to callback-aware visitor layers.
*/
using StateTupleType = std::tuple<>;
@@ -200,6 +202,12 @@ class StructuralVisitorObj : public Object {
* \brief Visit using the structural visit behavior registered by
kStructuralVisit for each type,
* or reflected structural fields when no custom behavior is registered.
*
+ * \note Dispatches to the value type's registered ``__s_visit__`` hook,
falling back to
+ * reflected fields only when no hook is registered. Call it on a child, or
on a matched
+ * value from a ``StructuralVisit`` callback to bypass callback dispatch for
that value.
+ * Called on the value whose own hook is running, it re-enters that hook --
there is no way
+ * to request the reflected path from inside a hook.
+ *
* \param value The value to visit.
* \return Expected interrupt state. An error means traversal failed.
*/
@@ -208,27 +216,16 @@ class StructuralVisitorObj : public Object {
static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralVisit);
AnyView attr = column[type_index];
- // Hooks propagate errors untouched; this is the engine dispatching into
`value`, so the
- // node is named here and nowhere else -- exactly one frame per node.
// case 1: Type-specific override registered as an opaque ABI visit
function pointer.
if (attr.type_index() == TypeIndex::kTVMFFIOpaquePtr) {
auto* visit_fn = reinterpret_cast<FStructuralVisit>(attr.cast<void*>());
TVMFFIAny raw = (*visit_fn)(this, value);
- if (TVM_FFI_PREDICT_FALSE(raw.type_index == TypeIndex::kTVMFFIError)) {
- details::UpdateVisitErrorContext(raw, value);
- }
return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>(raw);
}
// case 2: Type-specific override registered as an ffi::Function.
if (attr.type_index() == TypeIndex::kTVMFFIFunction) {
- Expected<Optional<VisitInterrupt>> result =
- attr.cast<Function>().CallExpected<Optional<VisitInterrupt>>(this,
value);
- if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
- Error err = result.error();
- details::UpdateVisitErrorContext(err, value);
- }
- return result;
+ return
attr.cast<Function>().CallExpected<Optional<VisitInterrupt>>(this, value);
}
if (TVM_FFI_PREDICT_FALSE(attr.type_index() != TypeIndex::kTVMFFINone)) {
@@ -491,18 +488,15 @@ enum class WalkOrder : int32_t {
namespace details {
/// \cond Doxygen_Suppress
-// Return from the current ABI visit function if Result stops traversal.
-// Result must evaluate to Expected whose raw storage can be moved to
TVMFFIAny.
-// If Result is an Error, append Node to the visit error context before
returning. Node is
-// required: dropping it silently degrades every error message produced below
this frame.
-// A raw pointer Node must be non-null; pass nullable nodes as ObjectRef or
Any so None is skipped.
-#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result)
\
- do {
\
- auto&& tvm_ffi_res_ = (Result);
\
- if (TVM_FFI_PREDICT_FALSE(
\
-
::tvm::ffi::details::StructuralVisitNeedEarlyReturn(tvm_ffi_res_))) {
\
- return
::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(tvm_ffi_res_)); \
- }
\
+// Return from the current raw or same-T Expected visit function if Result
stops traversal.
+// The rvalue-only proxy lets the enclosing return type select the
representation.
+#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result)
\
+ do {
\
+ auto&& tvm_ffi_res_ = (Result);
\
+ if (TVM_FFI_PREDICT_FALSE(
\
+
::tvm::ffi::details::StructuralVisitNeedEarlyReturn(tvm_ffi_res_))) { \
+ return
::tvm::ffi::details::MaybeReturnHelper(::std::move(tvm_ffi_res_)); \
+ }
\
} while (0)
/// \endcond
@@ -525,43 +519,16 @@ namespace details {
* ``Expected<Optional<VisitInterrupt>>`` storage produced by
* ``details::ExpectedUnsafe::MoveToTVMFFIAny``; the engine propagates it
without
* a runtime type check. This deliberate pair keeps the raw ABI path available
- * without rematerializing a typed ``Expected``. Calls use ``this->``, so
lookup
- * happens at instantiation in the Parent's class scope; it is not virtual
- * dispatch.
- *
- * \code
- * class CountingLayer : public StructuralVisitorObj {
- * public:
- * using StateTupleType = std::tuple<const int&>;
- * explicit CountingLayer(const StructuralVisitorVTable* vtable)
- * : StructuralVisitorObj(vtable) {}
+ * without rematerializing a typed ``Expected``. Engine calls use ``Parent::``
+ * qualification; this is static layer dispatch, not virtual dispatch. A layer
+ * must still define its own raw boilerplate because boilerplate inherited from
+ * a base resolves its unqualified typed call in that base's scope.
*
- * Expected<Optional<VisitInterrupt>> DefaultVisitExpected(AnyView value)
noexcept {
- * ++count_;
- * return StructuralVisitorObj::DefaultVisitExpected(value);
- * }
- *
- * protected:
- * TVMFFIAny DefaultVisitRaw(AnyView value) noexcept {
- * return
details::ExpectedUnsafe::MoveToTVMFFIAny(DefaultVisitExpected(value));
- * }
- * StateTupleType StateTuple() const noexcept { return std::tie(count_); }
- *
- * private:
- * int count_ = 0;
- * };
- *
- * auto callback = [](const ObjectRef&, const int&) -> Expected<WalkResult> {
- * return WalkResult::Advance();
- * };
- * using Engine = StructuralWalkEngine<CountingLayer, WalkOrder::kPreOrder,
- * decltype(callback)>;
- * StructuralVisitor visitor(make_object<Engine>(std::move(callback)));
- * auto result = visitor->VisitExpected(root);
- * \endcode
- *
- * \tparam Parent Visitor layer that supplies descent and callback state
through
- * the complete protocol above.
+ * \tparam Parent Traversal layer extended by the engine. Each layer that
+ * customizes typed descent must define its own ``Default*Raw``
+ * boilerplate; inherited boilerplate resolves its unqualified
+ * typed call in the base layer's scope. Engine protocol and
+ * descent calls are ``Parent::``-qualified.
* \tparam order Callback placement relative to child traversal.
* \tparam Callbacks Callback links, tested in declaration order.
*/
@@ -612,14 +579,14 @@ class StructuralWalkEngine : public Parent {
"StructuralWalk callback takes (value, state...) with an optional
trailing "
"definition-region kind");
try {
- static_assert(std::is_same_v<decltype(this->StateTuple()),
StateTupleType>,
+ static_assert(std::is_same_v<decltype(Parent::StateTuple()),
StateTupleType>,
"Parent::StateTuple() must return Parent::StateTupleType
by value");
- StateTupleType states = this->StateTuple();
+ StateTupleType states = Parent::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());
+ Parent::def_region_kind());
}
} catch (const Error& err) {
return Unexpected(err);
@@ -689,9 +656,11 @@ class StructuralWalkEngine : public Parent {
}
{
- // DefaultVisitExpected already named `value` if a hook it dispatched
failed.
- TVMFFIAny result = this->DefaultVisitRaw(value);
+ TVMFFIAny result = Parent::DefaultVisitRaw(value);
if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitRawNeedEarlyReturn(result))) {
+ if (TVM_FFI_PREDICT_FALSE(result.type_index ==
TypeIndex::kTVMFFIError)) {
+ details::UpdateVisitErrorContext(result, value);
+ }
return result;
}
}
@@ -733,22 +702,6 @@ class StructuralWalkEngine : public Parent {
*
* \sa WalkOrder, WalkResult
*
- * Example:
- *
- * \code
- * int num_adds = 0;
- *
- * Expected<Optional<VisitInterrupt>> result =
StructuralWalkExpected<WalkOrder::kPreOrder>(
- * root,
- * [&](const Add& add) -> Expected<WalkResult> {
- * ++num_adds;
- * return WalkResult::Advance();
- * },
- * [&](const Mul& mul) -> Expected<WalkResult> {
- * return WalkResult::Skip();
- * });
- * \endcode
- *
* \tparam order Whether to invoke the callback before or after visiting
children.
* \tparam Callbacks Callback types.
* \param root The root value to visit.
@@ -789,6 +742,188 @@ Optional<VisitInterrupt> StructuralWalk(AnyView root,
Callbacks&&... callbacks)
return StructuralWalkExpected<order>(root,
std::forward<Callbacks>(callbacks)...).value();
}
+// ---------------------------------------------------------------------------
+// Structural Visit API.
+// ---------------------------------------------------------------------------
+
+/*!
+ * \brief Engine of the callback-dispatched \ref tvm::ffi::StructuralVisit.
+ *
+ * A matched callback owns descent into its value, and its result is final. A
+ * value matching no callback uses the Parent layer's default descent. The
local
+ * typed callback fold preserves declaration-order first match and converts an
+ * ``Error`` thrown by a matched callback into the visit result.
+ *
+ * \tparam Parent Traversal layer extended by the engine. Each layer that
+ * customizes typed descent must define its own ``Default*Raw``
+ * boilerplate; inherited boilerplate resolves its unqualified
+ * typed call in the base layer's scope. Engine protocol and
+ * descent calls are ``Parent::``-qualified.
+ * \tparam Callbacks Callable types whose first parameter selects the
dispatched value type.
+ */
+template <typename Parent, typename... Callbacks>
+class StructuralVisitEngine : public Parent {
+ public:
+ static_assert(std::is_base_of_v<StructuralVisitorObj, Parent>,
+ "StructuralVisit Parent must derive from
StructuralVisitorObj");
+ /*!
+ * \brief Construct a visit engine over a chain of typed callbacks.
+ * \param callbacks Callbacks tested in declaration order; the first match
runs.
+ */
+ explicit StructuralVisitEngine(Callbacks... callbacks)
+ : Parent(VTable()), callbacks_(std::move(callbacks)...) {}
+
+ private:
+ /*!
+ * \brief Return this engine's callback-aware visitor vtable.
+ * \return Pointer to the immutable visitor vtable for this specialization.
+ */
+ static const StructuralVisitorVTable* VTable() {
+ static const StructuralVisitorVTable vtable{
+ &StructuralVisitEngine::DispatchVisit,
+ };
+ return &vtable;
+ }
+
+ /*!
+ * \brief Dispatch from the erased visitor pointer to the concrete engine.
+ * \param self The erased structural visitor object.
+ * \param value The value to visit.
+ * \return Interrupt state, or an error if traversal failed.
+ */
+ static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, AnyView value)
noexcept {
+ return static_cast<StructuralVisitEngine*>(self)->VisitImpl(value);
+ }
+
+ /*!
+ * \brief Visit one value, handing a matched callback ownership of its
descent.
+ * \param value The value to visit.
+ * \return Interrupt state, or an error if traversal failed.
+ */
+ TVMFFIAny VisitImpl(AnyView value) noexcept {
+ if (TVM_FFI_PREDICT_FALSE(value.type_index() == TypeIndex::kTVMFFINone)) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ Expected<Optional<VisitInterrupt>>(std::nullopt));
+ }
+ if (std::optional<Expected<Optional<VisitInterrupt>>> matched =
DispatchCallbacks(value)) {
+ // The matched callback already traversed as much of `value` as it
wanted, so its
+ // result is final and the engine does not descend on its own.
+ Expected<Optional<VisitInterrupt>> result = *std::move(matched);
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ Error err = result.error();
+ details::UpdateVisitErrorContext(err, value);
+ }
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+ // No callback claimed `value`. The Parent layer owns default descent.
+ TVMFFIAny result = Parent::DefaultVisitRaw(value);
+ if (TVM_FFI_PREDICT_FALSE(result.type_index == TypeIndex::kTVMFFIError)) {
+ details::UpdateVisitErrorContext(result, value);
+ }
+ return result;
+ }
+
+ /*! \brief Try one typed callback and preserve Error as an expected result.
*/
+ template <typename Callback>
+ TVM_FFI_INLINE std::optional<Expected<Optional<VisitInterrupt>>>
TryLink(Callback& callback,
+
AnyView value) noexcept {
+ using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
+ static_assert(FuncInfo::num_args == 2, "StructuralVisit callback takes
(value, visitor)");
+ using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
+ using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
+ 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::VisitorObjType>,
+ "second StructuralVisit callback argument must be "
+ "exactly Parent::VisitorObjType*");
+ auto* visitor = static_cast<typename Parent::VisitorObjType*>(this);
+ try {
+ if constexpr (std::is_same_v<TSub, AnyView>) {
+ return callback(value, visitor);
+ } else if constexpr (std::is_same_v<TSub, Any>) {
+ return callback(Any(value), visitor);
+ } else if (auto matched = value.template as<TSub>()) {
+ return callback(*std::move(matched), visitor);
+ }
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+ return std::nullopt;
+ }
+
+ /*! \brief Fold this engine's callback tuple in declaration order. */
+ template <size_t... Is>
+ TVM_FFI_INLINE std::optional<Expected<Optional<VisitInterrupt>>> TryLinks(
+ AnyView value, std::index_sequence<Is...>) noexcept {
+ std::optional<Expected<Optional<VisitInterrupt>>> result;
+ (... || (result = TryLink(std::get<Is>(callbacks_), value)).has_value());
+ return result;
+ }
+
+ /*!
+ * \brief Run the callback chain on \p value.
+ * \param value The value to dispatch on.
+ * \return The matched callback's result, or an empty optional when none
matched.
+ *
+ * \note An unmatched value is reported as such rather than folded into a
"continue"
+ * result: the engine has to tell "the callback chose to stop here" apart
from "no
+ * callback claimed this value".
+ */
+ std::optional<Expected<Optional<VisitInterrupt>>> DispatchCallbacks(AnyView
value) noexcept {
+ return TryLinks(value, std::index_sequence_for<Callbacks...>{});
+ }
+
+ /*! \brief Typed callbacks tested in declaration order, first match wins. */
+ std::tuple<Callbacks...> callbacks_;
+};
+
+/*!
+ * \brief Visit a structured value, letting a matched callback own descent.
+ *
+ * Each callback takes ``(value, StructuralVisitorObj* visitor)`` and returns
+ * ``Expected<Optional<VisitInterrupt>>``. The first argument follows the same
+ * matching rules as ``StructuralWalk``; callbacks are tested in declaration
+ * order and the first match is used.
+ *
+ * A matched callback owns descent into its value, and its result is final.
+ * Returning ``std::nullopt`` completes that subtree, a ``VisitInterrupt``
halts
+ * the traversal, and an ``Error`` fails it. A value matching no callback uses
+ * the visitor's default descent.
+ *
+ * \sa StructuralWalkExpected, StructuralVisitorObj, VisitInterrupt
+ *
+ * \tparam Callbacks Callback types.
+ * \param root The root value to visit.
+ * \param callbacks Callbacks invoked for matching nodes.
+ * \return ``std::nullopt`` if traversal completed, or the interrupt that
halted it.
+ */
+template <typename... Callbacks>
+Expected<Optional<VisitInterrupt>> StructuralVisitExpected(AnyView root,
+ Callbacks&&...
callbacks) noexcept {
+ static_assert(sizeof...(Callbacks) != 0, "StructuralVisit requires at least
one callback");
+ using Engine = StructuralVisitEngine<StructuralVisitorObj,
std::decay_t<Callbacks>...>;
+ StructuralVisitor
visitor(make_object<Engine>(std::forward<Callbacks>(callbacks)...));
+ return visitor->VisitExpected(root);
+}
+
+/*!
+ * \brief Throwing error over \ref tvm::ffi::StructuralVisitExpected.
+ *
+ * See \ref tvm::ffi::StructuralVisitExpected for callback semantics and
traversal behavior.
+ *
+ * \tparam Callbacks Callback types.
+ * \param root The root value to visit.
+ * \param callbacks Callbacks invoked for matching nodes. Each callback takes
+ * ``(value, StructuralVisitorObj*)`` and should return
+ * ``Expected<Optional<VisitInterrupt>>``.
+ * \return ``std::nullopt`` if traversal completed, or the interrupt that
halted it.
+ * \throws Error if traversal or a callback returned an error.
+ */
+template <typename... Callbacks>
+Optional<VisitInterrupt> StructuralVisit(AnyView root, Callbacks&&...
callbacks) {
+ return StructuralVisitExpected(root,
std::forward<Callbacks>(callbacks)...).value();
+}
+
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_
diff --git a/python/tvm_ffi/__init__.py b/python/tvm_ffi/__init__.py
index 59d55d89..1f4e812d 100644
--- a/python/tvm_ffi/__init__.py
+++ b/python/tvm_ffi/__init__.py
@@ -87,6 +87,7 @@ if TYPE_CHECKING or not _is_config_mode():
structural_equal,
structural_hash,
structural_map,
+ structural_visit,
structural_walk,
)
from . import serialization
@@ -188,6 +189,7 @@ __all__ = [
"structural_equal",
"structural_hash",
"structural_map",
+ "structural_visit",
"structural_walk",
"system_lib",
"use_raw_stream",
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/_ffi_api.py
index c19828eb..09fa770d 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/_ffi_api.py
@@ -118,7 +118,9 @@ if TYPE_CHECKING:
def StructuralMutatorVarRemapGet(_0: _StructuralMutator, _1: Any, /) ->
Any: ...
def StructuralMutatorVarRemapSet(_0: _StructuralMutator, _1: Any, _2: Any,
/) -> None: ...
def StructuralMutatorWithDefRegionKind(_0: _StructuralMutator, _1: int,
_2: Callable[..., Any], /) -> Any: ...
+ def StructuralVisit(_0: Any, _1: Sequence[tuple[int, Callable[..., Any]]],
/) -> _VisitInterrupt | None: ...
def StructuralVisitorDefRegionKind(_0: _StructuralVisitor, /) -> int: ...
+ def StructuralVisitorDefaultVisit(_0: _StructuralVisitor, _1: Any, /) ->
_VisitInterrupt | None: ...
def StructuralVisitorVisit(_0: _StructuralVisitor, _1: Any, /) ->
_VisitInterrupt | None: ...
def StructuralVisitorWithDefRegionKind(_0: _StructuralVisitor, _1: int,
_2: Callable[..., Any], /) -> Any: ...
def StructuralWalk(_0: Any, _1: Sequence[tuple[int, Callable[..., Any]]],
_2: Sequence[tuple[int, Callable[..., Any]]], _3: int, /) -> _VisitInterrupt |
None: ...
@@ -216,7 +218,9 @@ __all__ = [
"StructuralMutatorVarRemapGet",
"StructuralMutatorVarRemapSet",
"StructuralMutatorWithDefRegionKind",
+ "StructuralVisit",
"StructuralVisitorDefRegionKind",
+ "StructuralVisitorDefaultVisit",
"StructuralVisitorVisit",
"StructuralVisitorWithDefRegionKind",
"StructuralWalk",
diff --git a/python/tvm_ffi/structural.py b/python/tvm_ffi/structural.py
index 7a91ef6d..60d3b229 100644
--- a/python/tvm_ffi/structural.py
+++ b/python/tvm_ffi/structural.py
@@ -42,6 +42,7 @@ __all__ = [
"structural_equal",
"structural_hash",
"structural_map",
+ "structural_visit",
"structural_walk",
]
@@ -379,6 +380,33 @@ class StructuralVisitor(Object):
"""
return _ffi_api.StructuralVisitorVisit(self, value)
+ def default_visit(self, value: Any) -> VisitInterrupt | None:
+ """Visit ``value`` using its registered or reflected child traversal.
+
+ .. warning::
+ Never call ``default_visit`` on the value whose ``__s_visit__``
hook
+ is currently running. Doing so re-enters the same hook recursively
+ and can exhaust the C stack, crashing the process.
+
+ This bypasses the active engine callback dispatch for ``value`` itself,
+ but still invokes the value type's registered ``__s_visit__`` hook when
+ one exists. Use it for a child whose default traversal is wanted.
+ Recursive children still use this visitor.
+
+ Parameters
+ ----------
+ value
+ Value whose default children should be traversed.
+
+ Returns
+ -------
+ result
+ ``None`` if traversal should continue, otherwise a
+ :class:`VisitInterrupt` carrying the early-exit payload.
+
+ """
+ return _ffi_api.StructuralVisitorDefaultVisit(self, value)
+
def def_region_kind(self) -> DefRegionKind:
"""Low-level API to return the currently active structural def-region
kind.
@@ -619,6 +647,40 @@ def structural_walk(
return _ffi_api.StructuralWalk(root, entries,
entries_with_def_region_kind, order_int)
+def structural_visit(
+ root: Any,
+ callbacks: tuple | Sequence | Callable = (),
+) -> VisitInterrupt | None:
+ """Visit a value structurally with callbacks that own child traversal.
+
+ Each callback receives ``(value, visitor)``. It may call
+ :meth:`StructuralVisitor.visit` for selected children or
+ :meth:`StructuralVisitor.default_visit` for the value's registered/default
+ descent. Returning without either call prunes that value's subtree.
+
+ Parameters
+ ----------
+ root
+ Root value to traverse.
+
+ callbacks
+ A callback, ``(type, callback)`` entry, grouped type entry, or sequence
+ of entries. Entries are tried in order and the first match owns
descent.
+
+ Returns
+ -------
+ result
+ ``None`` if traversal completed, otherwise a :class:`VisitInterrupt`.
+
+ """
+ callback_entries = _normalize_callbacks(callbacks,
api_name="structural_visit")
+ entries: list[tuple[int, Callable[[Any, StructuralVisitor], Any]]] = [
+ (_callback_type_to_type_index(t, api_name="structural_visit"), fn)
+ for t, fn in callback_entries
+ ]
+ return _ffi_api.StructuralVisit(root, entries)
+
+
def structural_map(
root: Any,
callbacks: tuple | Sequence | Callable = (),
diff --git a/src/ffi/extra/structural_visit.cc
b/src/ffi/extra/structural_visit.cc
index 62b251dd..e1f73a5f 100644
--- a/src/ffi/extra/structural_visit.cc
+++ b/src/ffi/extra/structural_visit.cc
@@ -86,6 +86,30 @@ Expected<Optional<VisitInterrupt>> StructuralWalkExpected(
}
}
+/*!
+ * \brief Runtime callback-driven structural visit.
+ * \param root The root value to visit.
+ * \param callbacks Runtime callback entries of ``(type_index,
ffi::Function)`` invoked as
+ * ``callback(value, visitor)``.
+ * \return Expected interrupt state. An error means traversal failed.
+ */
+Expected<Optional<VisitInterrupt>> StructuralVisitExpected(
+ AnyView root, const Array<Tuple<int32_t, Function>>& callbacks) noexcept {
+ auto dispatch = [callbacks](AnyView value,
+ StructuralVisitorObj* visitor) ->
Expected<Optional<VisitInterrupt>> {
+ for (const auto& entry : callbacks) {
+ if (!RuntimeTypeIndexMatch(value.type_index(), entry.template get<0>()))
continue;
+ return entry.template get<1>().CallExpected<Optional<VisitInterrupt>>(
+ value, GetRef<StructuralVisitor>(visitor));
+ }
+ return visitor->DefaultVisitExpected(value);
+ };
+
+ using Visitor = StructuralVisitEngine<StructuralVisitorObj,
decltype(dispatch)>;
+ StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
+ return visitor->VisitExpected(root);
+}
+
/*! \brief Visit entries in a sequence container. */
TVMFFIAny VisitSeqContainer(StructuralVisitorObj* visitor, const SeqBaseObj*
self) noexcept {
for (const Any& item : *self) {
@@ -140,6 +164,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def("ffi.VisitInterrupt", [](Any value) { return
VisitInterrupt(std::move(value)); })
.def_method("ffi.StructuralVisitorVisit", &StructuralVisitorObj::Visit)
+ .def_method("ffi.StructuralVisitorDefaultVisit",
+ [](const StructuralVisitor& visitor, AnyView value) {
+ return visitor->DefaultVisitExpected(value).value();
+ })
.def_method("ffi.StructuralVisitorDefRegionKind",
&StructuralVisitorObj::def_region_kind)
.def_method(
"ffi.StructuralVisitorWithDefRegionKind",
@@ -153,6 +181,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return details::StructuralWalkExpected(root, callbacks,
callbacks_with_def_region_kind,
order)
.value();
+ })
+ .def("ffi.StructuralVisit",
+ [](AnyView root,
+ const Array<Tuple<int32_t, Function>>& callbacks) ->
Optional<VisitInterrupt> {
+ return details::StructuralVisitExpected(root, callbacks).value();
});
refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralVisit);
refl::TypeAttrDef<ArrayObj>().attr(
diff --git a/tests/cpp/extra/test_structural_visit.cc
b/tests/cpp/extra/test_structural_visit.cc
index d579e6a6..82a6bccf 100644
--- a/tests/cpp/extra/test_structural_visit.cc
+++ b/tests/cpp/extra/test_structural_visit.cc
@@ -85,10 +85,13 @@ struct VisitTag {
template <typename Parent = StructuralVisitorObj>
class StructuralWalkWithVisitCount : public Parent {
public:
+ using VisitorObjType = StructuralWalkWithVisitCount;
using StateTupleType = std::tuple<const VisitCount&, const VisitTag&>;
explicit StructuralWalkWithVisitCount(const StructuralVisitorVTable* vtable)
: Parent(vtable) {}
+ int callback_tag() const noexcept { return visit_tag_.value; }
+
TVM_FFI_INLINE Expected<Optional<VisitInterrupt>>
DefaultVisitExpected(AnyView value) noexcept {
++visit_count_.value;
return Parent::DefaultVisitExpected(value);
@@ -108,6 +111,12 @@ class StructuralWalkWithVisitCount : public Parent {
VisitTag visit_tag_;
};
+template <typename Parent>
+class StructuralVisitOuterLayer : public Parent {
+ public:
+ explicit StructuralVisitOuterLayer(const StructuralVisitorVTable* vtable) :
Parent(vtable) {}
+};
+
StructuralVisitor MakeTestVisitor() { return
StructuralVisitor(make_object<TestVisitorObj>()); }
TestVisitorObj* AsTestVisitor(const StructuralVisitor& visitor) {
@@ -576,6 +585,7 @@ TEST(StructuralVisitor, WalkCatchesError) {
Expected<Optional<VisitInterrupt>> result =
StructuralWalkExpected<WalkOrder::kPreOrder>(
root, [&](const ObjectRef&) -> Expected<WalkResult> {
TVM_FFI_THROW(ValueError) << "walk callback threw";
+ return WalkResult::Advance();
});
ASSERT_TRUE(result.is_err());
@@ -666,4 +676,78 @@ TEST(StructuralVisitor, WalkAnyFallback) {
ExpectTrace(trace, {"object-ref", "object-ref"});
}
+// ---------------------------------------------------------------------------
+// StructuralVisit behavior.
+// ---------------------------------------------------------------------------
+
+TEST(StructuralVisit, CallbackDrivenTraversal) {
+ TVarWithDep lhs("lhs", TVarWithDep("pruned-dependency"));
+ TVarWithDep stop("stop");
+ TVarWithDep skipped("skipped");
+ std::vector<std::pair<std::string, TVMFFIDefRegionKind>> trace;
+ Expected<Optional<VisitInterrupt>> result = StructuralVisitExpected(
+ TPair(Array<ObjectRef>{lhs}, Array<ObjectRef>{stop, skipped}),
+ [](const TPairObj* pair,
+ StructuralVisitorObj* visitor) -> Expected<Optional<VisitInterrupt>> {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindRecursive, [&] { return
visitor->VisitExpected(pair->lhs); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindNonRecursive, [&] { return
visitor->VisitExpected(pair->rhs); }));
+ return Optional<VisitInterrupt>(std::nullopt);
+ },
+ [&](const TVarWithDepObj* var,
+ StructuralVisitorObj* visitor) -> Expected<Optional<VisitInterrupt>>
{
+ trace.emplace_back(var->name, visitor->def_region_kind());
+ if (var->name == "stop") {
+ return Optional<VisitInterrupt>(VisitInterrupt(String("found
stop")));
+ }
+ return Optional<VisitInterrupt>(std::nullopt);
+ });
+ ASSERT_TRUE(result.is_ok());
+ ASSERT_TRUE(result.value().has_value());
+ EXPECT_EQ(result.value().value()->value.cast<String>(), "found stop");
+ ASSERT_EQ(trace.size(), 2u);
+ EXPECT_EQ(trace[0], std::make_pair(std::string("lhs"),
kTVMFFIDefRegionKindRecursive));
+ EXPECT_EQ(trace[1], std::make_pair(std::string("stop"),
kTVMFFIDefRegionKindNonRecursive));
+
+ using CallbackLayer = StructuralWalkWithVisitCount<>;
+ using ComposedLayer = StructuralVisitOuterLayer<CallbackLayer>;
+ bool layer_callback_ran = false;
+ bool base_callback_ran = false;
+ auto layer_callback = [&](const TVarWithDepObj*,
+ CallbackLayer* visitor) ->
Expected<Optional<VisitInterrupt>> {
+ EXPECT_EQ(visitor->callback_tag(), 7);
+ layer_callback_ran = true;
+ return Optional<VisitInterrupt>(std::nullopt);
+ };
+ auto base_callback = [&](const TVarObj*, CallbackLayer*) ->
Expected<Optional<VisitInterrupt>> {
+ base_callback_ran = true;
+ return Optional<VisitInterrupt>(std::nullopt);
+ };
+ using ComposedEngine =
+ StructuralVisitEngine<ComposedLayer, decltype(layer_callback),
decltype(base_callback)>;
+ static_assert(std::is_same_v<typename ComposedEngine::VisitorObjType,
CallbackLayer>);
+ StructuralVisitor composed(
+ make_object<ComposedEngine>(std::move(layer_callback),
std::move(base_callback)));
+ ASSERT_FALSE(composed->VisitExpected(Array<ObjectRef>{TVarWithDep("layer"),
TVar("base")})
+ .value()
+ .has_value());
+ EXPECT_TRUE(layer_callback_ran);
+ EXPECT_TRUE(base_callback_ran);
+
+ std::vector<std::string> error_trace;
+ result = StructuralVisitExpected(
+ Array<ObjectRef>{TVar("throw"), TVar("after")},
+ [&](const TVarObj* var, StructuralVisitorObj*) ->
Expected<Optional<VisitInterrupt>> {
+ error_trace.emplace_back(var->name);
+ if (var->name == "throw") {
+ TVM_FFI_THROW(ValueError) << "visit callback threw";
+ }
+ return Optional<VisitInterrupt>(std::nullopt);
+ });
+ ASSERT_TRUE(result.is_err());
+ EXPECT_EQ(result.error().message(), "visit callback threw");
+ ExpectTrace(error_trace, {"throw"});
+}
+
} // namespace
diff --git a/tests/python/test_structural.py b/tests/python/test_structural.py
index 51d4e6f7..9ad0f3ae 100644
--- a/tests/python/test_structural.py
+++ b/tests/python/test_structural.py
@@ -259,6 +259,76 @@ def test_structural_walk_interrupt() -> None:
assert tvm_ffi.structural_equal(result.value, {"found": 2})
+def test_structural_visit_default_visit_binding() -> None:
+ trace: list[int | str] = []
+ root = tvm_ffi.Array([tvm_ffi.Array([1, 2, 3])])
+
+ def visit_array(
+ value: tvm_ffi.Array, visitor: tvm_ffi.StructuralVisitor
+ ) -> tvm_ffi.VisitInterrupt | None:
+ trace.append("array")
+ return visitor.default_visit(value[0])
+
+ def interrupt_on_two(
+ value: int, visitor: tvm_ffi.StructuralVisitor
+ ) -> tvm_ffi.VisitInterrupt | None:
+ assert isinstance(visitor, tvm_ffi.StructuralVisitor)
+ trace.append(value)
+ return tvm_ffi.VisitInterrupt("done") if value == 2 else None
+
+ result = tvm_ffi.structural_visit(
+ root,
+ [
+ (tvm_ffi.Array, visit_array),
+ (int, interrupt_on_two),
+ ],
+ )
+
+ # The root callback asks for the nested Array's registered/default descent.
+ # Its matching Array callback is bypassed, while recursive ints re-enter
+ # this visitor and reach their callback.
+ assert trace == ["array", 1, 2]
+ assert isinstance(result, tvm_ffi.VisitInterrupt)
+ assert result.value == "done"
+
+ direct_trace: list[int] = []
+
+ def fail_directly(
+ value: int, visitor: tvm_ffi.StructuralVisitor
+ ) -> tvm_ffi.VisitInterrupt | None:
+ assert isinstance(visitor, tvm_ffi.StructuralVisitor)
+ direct_trace.append(value)
+ raise ValueError("direct structural visit failure")
+
+ with pytest.raises(ValueError, match="direct structural visit failure"):
+ tvm_ffi.structural_visit(tvm_ffi.Array([1, 2]), [(int, fail_directly)])
+ assert direct_trace == [1]
+
+ nested_trace: list[int | str] = []
+
+ def visit_outer_array(
+ value: tvm_ffi.Array, visitor: tvm_ffi.StructuralVisitor
+ ) -> tvm_ffi.VisitInterrupt | None:
+ nested_trace.append("array")
+ return visitor.default_visit(value[0])
+
+ def fail_nested(
+ value: int, visitor: tvm_ffi.StructuralVisitor
+ ) -> tvm_ffi.VisitInterrupt | None:
+ assert isinstance(visitor, tvm_ffi.StructuralVisitor)
+ nested_trace.append(value)
+ if value == 2:
+ raise ValueError("nested structural visit failure")
+ return None
+
+ with pytest.raises(ValueError, match="nested structural visit failure"):
+ tvm_ffi.structural_visit(
+ tvm_ffi.Array([tvm_ffi.Array([1, 2, 3]), tvm_ffi.Array([4])]),
+ [(tvm_ffi.Array, visit_outer_array), (int, fail_nested)],
+ )
+ assert nested_trace == ["array", 1, 2]
+
+
def test_structural_walk_nested_containers_and_skips_map_keys() -> None:
root = tvm_ffi.Array(
[
