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 4205a643 [REFACTOR][EXTRA] Allow StructuralWalkVisitor to carry state
via Parent (#747)
4205a643 is described below
commit 4205a643fa8d861d08f8aa7dc3dedbf1c7bda5f2
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 5 13:51:17 2026 -0400
[REFACTOR][EXTRA] Allow StructuralWalkVisitor to carry state via Parent
(#747)
Downstream walks sometimes need to carry state alongside traversal,
scoped around the subtree they descend into. `StructuralWalkVisitor`
could not express that because its callback chain was composed before
the visitor existed, so callbacks could not see visitor-layer state.
This change moves callback matching into a public walk engine and adds a
`Parent` layer protocol:
- Preserve the exact Parent-free `StructuralWalkExpected<order>` and
`StructuralWalk<order>` entry points; no additional `WithState` free
function ships.
- Expose `tvm::ffi::StructuralWalkEngine<Parent, order, Callbacks...>`
for layered users and document the complete Parent protocol, including
dependent unqualified descent lookup rather than virtual dispatch.
- Preserve protected raw ABI descent and public expected descent; layers
overriding expected descent provide the paired raw redirect in their own
class body.
- Unpack the Parent's const state tuple positionally into every callback
after the value, supporting multiple and duplicate state types without
lookup rules.
- Allow `TVMFFIDefRegionKind` only as an optional trailing argument and
read it lazily only when requested.
- Cover public engine construction, state observation, raw descent,
skip/interrupt behavior, combined state/region dispatch, and
pre/post-order placement in focused tests.
---
include/tvm/ffi/extra/structural_visit.h | 292 ++++++++++++++++++-------------
src/ffi/extra/structural_visit.cc | 6 +-
tests/cpp/extra/test_structural_visit.cc | 116 ++++++++++++
3 files changed, 290 insertions(+), 124 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index 468bfd4f..e8d567ad 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -138,6 +138,9 @@ struct StructuralVisitorVTable {
*/
class StructuralVisitorObj : public Object {
public:
+ /*! \brief State references made available to callback-aware visitor layers.
*/
+ using StateTupleType = std::tuple<>;
+
/*!
* \brief Visit a value, dispatching through this visitor's vtable.
*
@@ -246,6 +249,26 @@ class StructuralVisitorObj : public Object {
/// \endcond
protected:
+ /*!
+ * \brief Redirect raw ABI descent to \ref DefaultVisitExpected.
+ *
+ * A visitor layer that overrides either descent form must redeclare both in
+ * the same layer so dependent member lookup reaches the paired override. The
+ * raw form is deliberately retained for ABI traversal and permits a layer to
+ * forward raw storage without rematerializing a typed ``Expected``.
+ *
+ * \param value The value to descend into.
+ * \return Raw ``Expected<Optional<VisitInterrupt>>`` storage produced by
+ * ``details::ExpectedUnsafe::MoveToTVMFFIAny``. The caller
+ * reinterprets this storage without a runtime type check.
+ */
+ TVM_FFI_INLINE TVMFFIAny DefaultVisitRaw(AnyView value) noexcept {
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(DefaultVisitExpected(value));
+ }
+
+ /*! \brief Return the state references maintained by this visitor layer. */
+ TVM_FFI_INLINE StateTupleType StateTuple() const noexcept { return {}; }
+
/*!
* \brief Construct a structural visitor from an immutable dispatch vtable.
* \param vtable The non-null dispatch table for this visitor. It must
outlive this object.
@@ -300,6 +323,17 @@ TVM_FFI_INLINE bool StructuralVisitNeedEarlyReturn(const
Expected<T>& result) no
return type_index == TypeIndex::kTVMFFIError || type_index ==
TypeIndex::kTVMFFIVisitInterrupt;
}
+/*!
+ * \brief Return true when raw Expected storage carries a traversal-stopping
state.
+ * \param result Raw ``Expected<Optional<VisitInterrupt>>`` storage returned by
+ * a visitor descent hook. Its type contract is trusted without
+ * a runtime check.
+ * \return Whether \p result carries anything other than successful completion.
+ */
+TVM_FFI_INLINE bool StructuralVisitRawNeedEarlyReturn(const TVMFFIAny& result)
noexcept {
+ return result.type_index != TypeIndex::kTVMFFINone;
+}
+
/*!
* \brief Walk reflected structural fields of object-backed \p obj.
*
@@ -472,23 +506,79 @@ namespace details {
} while (0)
/// \endcond
+} // namespace details
+
/*!
- * \brief Visitor used by callback-dispatched ``StructuralWalk``.
+ * \brief Callback-dispatched walk engine with a state-carrying Parent layer.
+ *
+ * Every callback receives all entries of ``Parent::StateTupleType``
+ * positionally after the value and may optionally take
+ * ``TVMFFIDefRegionKind`` as the final argument.
+ *
+ * A Parent layer derives from ``StructuralVisitorObj``, declares a public
+ * ``StateTupleType``, accepts and forwards ``const StructuralVisitorVTable*``
+ * in its constructor, and provides an at-least-protected
+ * ``StateTuple() const noexcept`` that returns ``StateTupleType`` by value.
The
+ * state references in that tuple must outlive the traversal. A layer
overriding
+ * either ``DefaultVisitExpected`` or ``DefaultVisitRaw`` declares both forms,
+ * at least protected, in that same class. The raw form must return raw
+ * ``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) {}
*
+ * 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 order Callback placement relative to child traversal.
- * \tparam Dispatch Callable returning ``Expected<WalkResult>`` when invoked
with ``AnyView`` and
- * the active def-region kind. User callbacks wrapped by
this dispatcher may
- * accept either ``(value)`` or ``(value, def_region_kind)``.
+ * \tparam Callbacks Callback links, tested in declaration order.
*/
-template <WalkOrder order, typename Dispatch>
-class StructuralWalkVisitorObj : public StructuralVisitorObj {
+template <typename Parent, WalkOrder order, typename... Callbacks>
+class StructuralWalkEngine : public Parent {
public:
+ static_assert(std::is_base_of_v<StructuralVisitorObj, Parent>,
+ "StructuralWalk Parent must derive from StructuralVisitorObj");
+ /*! \brief Tuple of const state references supplied by the Parent layer. */
+ using StateTupleType = typename Parent::StateTupleType;
+
/*!
* \brief Construct a structural walk visitor.
- * \param dispatch The composed dispatcher invoked on each visited node.
+ * \param callbacks The typed callback links, tested in declaration order.
*/
- explicit StructuralWalkVisitorObj(Dispatch dispatch)
- : StructuralVisitorObj(VTable()), dispatch_(std::move(dispatch)) {}
+ explicit StructuralWalkEngine(Callbacks... callbacks)
+ : Parent(VTable()), callbacks_(std::move(callbacks)...) {}
private:
/*!
@@ -497,7 +587,7 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
*/
static const StructuralVisitorVTable* VTable() {
static const StructuralVisitorVTable vtable{
- &StructuralWalkVisitorObj::DispatchVisit,
+ &StructuralWalkEngine::DispatchVisit,
};
return &vtable;
}
@@ -509,7 +599,60 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
* \return Interrupt state, or an error if traversal failed.
*/
static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, AnyView value)
noexcept {
- return static_cast<StructuralWalkVisitorObj*>(self)->VisitImpl(value);
+ return static_cast<StructuralWalkEngine*>(self)->VisitImpl(value);
+ }
+
+ /*! \brief Invoke one matched callback with its declared state, if any. */
+ template <typename Callback, typename Value, size_t... Is>
+ TVM_FFI_INLINE Expected<WalkResult> InvokeCallbackLink(Callback& callback,
Value&& value,
+
std::index_sequence<Is...>) noexcept {
+ using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
+ static_assert(
+ FuncInfo::num_args == 1 + sizeof...(Is) || FuncInfo::num_args == 2 +
sizeof...(Is),
+ "StructuralWalk 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 Try one callback link, storing its result when the value type
matches. */
+ template <typename Callback>
+ TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value,
+ Expected<WalkResult>* out) noexcept {
+ using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
+ static_assert(FuncInfo::num_args >= 1, "StructuralWalk callback requires a
value argument");
+ using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
+ using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
+ using StateIndices =
std::make_index_sequence<std::tuple_size_v<StateTupleType>>;
+ if constexpr (std::is_same_v<TSub, AnyView>) {
+ *out = InvokeCallbackLink(callback, value, StateIndices{});
+ return true;
+ } else if constexpr (std::is_same_v<TSub, Any>) {
+ *out = InvokeCallbackLink(callback, Any(value), StateIndices{});
+ return true;
+ } else if (auto matched = value.template as<TSub>()) {
+ *out = InvokeCallbackLink(callback, *std::move(matched), StateIndices{});
+ return true;
+ }
+ return false;
+ }
+
+ /*! \brief Try callback links in declaration order until the first match. */
+ template <size_t... Is>
+ TVM_FFI_INLINE bool TryLinks(AnyView value, Expected<WalkResult>* out,
+ std::index_sequence<Is...>) noexcept {
+ return (TryLink(std::get<Is>(callbacks_), value, out) || ...);
}
/*!
@@ -523,7 +666,8 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
Expected<Optional<VisitInterrupt>>(std::nullopt));
}
if constexpr (order == WalkOrder::kPreOrder) {
- auto result = dispatch_(value, this->def_region_kind());
+ Expected<WalkResult> result = WalkResult::Advance();
+ TryLinks(value, &result, std::index_sequence_for<Callbacks...>{});
if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitNeedEarlyReturn(result))) {
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
Error err = result.error();
@@ -546,14 +690,15 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
{
// DefaultVisitExpected already named `value` if a hook it dispatched
failed.
- auto result = DefaultVisitExpected(value);
- if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitNeedEarlyReturn(result))) {
- return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ TVMFFIAny result = this->DefaultVisitRaw(value);
+ if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitRawNeedEarlyReturn(result))) {
+ return result;
}
}
if constexpr (order == WalkOrder::kPostOrder) {
- auto result = dispatch_(value, this->def_region_kind());
+ Expected<WalkResult> result = WalkResult::Advance();
+ TryLinks(value, &result, std::index_sequence_for<Callbacks...>{});
if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitNeedEarlyReturn(result))) {
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
Error err = result.error();
@@ -567,110 +712,17 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
Expected<Optional<VisitInterrupt>>(std::nullopt));
}
- /*! \brief Composed dispatch closure invoked once per visited node. */
- Dispatch dispatch_;
+ /*! \brief The callback links, tested in declaration order. */
+ std::tuple<Callbacks...> callbacks_;
};
-/*!
- * \brief Compose typed callbacks into a single per-node dispatcher.
- *
- * Each callback dispatches on its first parameter's type; callbacks are tested
- * in declaration order and the first match runs. Callbacks may take an
optional
- * second ``TVMFFIDefRegionKind`` argument. Nodes that match no callback fall
- * through and traversal continues normally.
- */
-struct StructuralWalkCallbackChain {
- /*!
- * \brief Build a dispatcher closure over a chain of typed callbacks.
- * \tparam Callbacks Callable types whose first parameter selects the
dispatched
- * value type.
- * \param callbacks Callbacks to be tested in order.
- * \return A dispatcher closure of type ``Expected<WalkResult>(AnyView,
- * TVMFFIDefRegionKind)``. Each user callback may take either
- * ``(value)`` or ``(value, def_region_kind)``.
- */
- template <typename... Callbacks>
- static auto FromChain(Callbacks... callbacks) {
- return [=](AnyView x, TVMFFIDefRegionKind kind) mutable ->
Expected<WalkResult> {
- try {
- Optional<Expected<WalkResult>> result;
- // Fold expression: each TryCallLink returns empty Optional on no-match
- // (falsy) or a result on match (truthy); || short-circuits on first
match.
- (... || (result = TryCallLink(callbacks, x, kind)));
- if (result.has_value()) {
- return std::move(result).value();
- }
- return WalkResult::Advance();
- } catch (const Error& err) {
- return Unexpected(err);
- }
- };
- }
-
- private:
- /*!
- * \brief Invoke ``callback`` when ``x`` matches its first parameter type.
- * \tparam Callback Callable whose first parameter selects the value type and
- * whose optional second parameter receives the active
def-region kind.
- * \param callback The callback under test.
- * \param x The value to dispatch on.
- * \param kind The active def-region kind.
- * \return The callback result if it matched, empty ``Optional`` otherwise.
- */
- template <typename Callback>
- static Optional<Expected<WalkResult>> TryCallLink(Callback& callback,
AnyView x,
- TVMFFIDefRegionKind kind) {
- using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
- static_assert(FuncInfo::num_args == 1 || FuncInfo::num_args == 2,
- "StructuralWalk callbacks must take one argument (value) or
two arguments "
- "(value, def-region kind)");
- using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
- using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
- if constexpr (std::is_same_v<TSub, AnyView>) {
- // callback on AnyView
- return InvokeCallback(callback, x, kind);
- } else if constexpr (std::is_same_v<TSub, Any>) {
- // callback on Any
- return InvokeCallback(callback, Any(x), kind);
- } else {
- if (auto opt = x.template as<TSub>()) {
- return InvokeCallback(callback, *std::move(opt), kind);
- }
- }
- return std::nullopt;
- }
-
- /*!
- * \brief Invoke a matched callback with optional def-region context.
- * \tparam Callback Callable returning ``Expected<WalkResult>``.
- * \tparam Value Type of the converted value passed to the callback.
- * \param callback The matched callback to invoke.
- * \param value The converted value.
- * \param kind The active def-region kind.
- * \return The callback result.
- */
- template <typename Callback, typename Value>
- static Expected<WalkResult> InvokeCallback(Callback& callback, Value&& value,
- TVMFFIDefRegionKind kind) {
- using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
- if constexpr (FuncInfo::num_args == 1) {
- return callback(std::forward<Value>(value));
- } else {
- return callback(std::forward<Value>(value), kind);
- }
- }
-};
-
-} // namespace details
-
/*!
* \brief Walk a structured value graph and invoke typed callbacks on selected
values.
*
* The callbacks are invoked only for values matching the first argument type
of
* one of the callbacks. The first callback argument may be ``AnyView``,
``Any``,
* an object reference type, an object pointer type, or another FFI-convertible
- * POD type. A callback may also optionally take a second
``TVMFFIDefRegionKind`` argument
- * to inspect whether the value is being visited in a definition region.
+ * POD type. It may optionally take ``TVMFFIDefRegionKind`` after the value.
* Callbacks are tested in order, and the first match is used.
*
* Each callback should return ``Expected<WalkResult>``; see ``WalkResult``.
@@ -700,9 +752,8 @@ struct StructuralWalkCallbackChain {
* \tparam order Whether to invoke the callback before or after visiting
children.
* \tparam Callbacks Callback types.
* \param root The root value to visit.
- * \param callbacks Callbacks invoked for matching nodes. Each callback may
take
- * either ``(value)`` or ``(value, def_region_kind)`` and
should return
- * ``Expected<WalkResult>``.
+ * \param callbacks Callbacks invoked for matching nodes as ``(value)`` or
+ * ``(value, TVMFFIDefRegionKind)``.
* \return ``std::nullopt`` if traversal completed, or the interrupt returned
by
* a callback.
*
@@ -712,10 +763,8 @@ template <WalkOrder order, typename... Callbacks>
Expected<Optional<VisitInterrupt>> StructuralWalkExpected(AnyView root,
Callbacks&&...
callbacks) noexcept {
static_assert(sizeof...(Callbacks) != 0, "StructuralWalk requires at least
one callback");
- auto dispatch =
-
details::StructuralWalkCallbackChain::FromChain(std::forward<Callbacks>(callbacks)...);
- using Visitor = details::StructuralWalkVisitorObj<order, decltype(dispatch)>;
- StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
+ using Visitor = StructuralWalkEngine<StructuralVisitorObj, order,
std::decay_t<Callbacks>...>;
+ StructuralVisitor
visitor(make_object<Visitor>(std::forward<Callbacks>(callbacks)...));
return visitor->VisitExpected(root);
}
@@ -727,9 +776,8 @@ Expected<Optional<VisitInterrupt>>
StructuralWalkExpected(AnyView root,
* \tparam order Whether to invoke the callback before or after visiting
children.
* \tparam Callbacks Callback types.
* \param root The root value to visit.
- * \param callbacks Callbacks invoked for matching nodes. Each callback may
take
- * either ``(value)`` or ``(value, def_region_kind)`` and
should return
- * ``Expected<WalkResult>``.
+ * \param callbacks Callbacks invoked for matching nodes as ``(value)`` or
+ * ``(value, TVMFFIDefRegionKind)``.
* \return ``std::nullopt`` if traversal completed, or the interrupt returned
by
* a callback.
* \throws Error if traversal or a callback returned an error.
diff --git a/src/ffi/extra/structural_visit.cc
b/src/ffi/extra/structural_visit.cc
index a0ca4d9d..62b251dd 100644
--- a/src/ffi/extra/structural_visit.cc
+++ b/src/ffi/extra/structural_visit.cc
@@ -74,11 +74,13 @@ Expected<Optional<VisitInterrupt>> StructuralWalkExpected(
};
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
- using Visitor = StructuralWalkVisitorObj<WalkOrder::kPreOrder,
decltype(dispatch)>;
+ using Visitor =
+ StructuralWalkEngine<StructuralVisitorObj, WalkOrder::kPreOrder,
decltype(dispatch)>;
StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
return visitor->VisitExpected(root);
} else {
- using Visitor = StructuralWalkVisitorObj<WalkOrder::kPostOrder,
decltype(dispatch)>;
+ using Visitor =
+ StructuralWalkEngine<StructuralVisitorObj, WalkOrder::kPostOrder,
decltype(dispatch)>;
StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
return visitor->VisitExpected(root);
}
diff --git a/tests/cpp/extra/test_structural_visit.cc
b/tests/cpp/extra/test_structural_visit.cc
index eeb71fc6..d579e6a6 100644
--- a/tests/cpp/extra/test_structural_visit.cc
+++ b/tests/cpp/extra/test_structural_visit.cc
@@ -27,6 +27,7 @@
#include <stdexcept>
#include <string>
+#include <tuple>
#include <utility>
#include <vector>
@@ -73,6 +74,40 @@ class TestVisitorObj : public StructuralVisitorObj {
}
};
+struct VisitCount {
+ int value = 0;
+};
+
+struct VisitTag {
+ int value = 7;
+};
+
+template <typename Parent = StructuralVisitorObj>
+class StructuralWalkWithVisitCount : public Parent {
+ public:
+ using StateTupleType = std::tuple<const VisitCount&, const VisitTag&>;
+
+ explicit StructuralWalkWithVisitCount(const StructuralVisitorVTable* vtable)
: Parent(vtable) {}
+
+ TVM_FFI_INLINE Expected<Optional<VisitInterrupt>>
DefaultVisitExpected(AnyView value) noexcept {
+ ++visit_count_.value;
+ return Parent::DefaultVisitExpected(value);
+ }
+
+ protected:
+ TVM_FFI_INLINE TVMFFIAny DefaultVisitRaw(AnyView value) noexcept {
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(DefaultVisitExpected(value));
+ }
+
+ TVM_FFI_INLINE StateTupleType StateTuple() const noexcept {
+ return std::tie(visit_count_, visit_tag_);
+ }
+
+ private:
+ VisitCount visit_count_;
+ VisitTag visit_tag_;
+};
+
StructuralVisitor MakeTestVisitor() { return
StructuralVisitor(make_object<TestVisitorObj>()); }
TestVisitorObj* AsTestVisitor(const StructuralVisitor& visitor) {
@@ -411,6 +446,87 @@ TEST(StructuralVisitor, WalkReceivesDefRegionKind) {
ExpectTrace(use_vars, {"x", "y"});
}
+TEST(StructuralVisitor, WalkParentStatePreservesDescentSkipAndInterrupt) {
+ ObjectRef lhs = TVar("lhs");
+ ObjectRef rhs = TVar("rhs");
+ ObjectRef root = TPair(lhs, rhs);
+ std::vector<int> counts;
+ std::vector<int> tags;
+ std::vector<TVMFFIDefRegionKind> kinds;
+
+ auto pair_callback = [&](const TPairObj*, const VisitCount& count, const
VisitTag& tag,
+ TVMFFIDefRegionKind kind) -> Expected<WalkResult> {
+ counts.push_back(count.value);
+ tags.push_back(tag.value);
+ kinds.push_back(kind);
+ return WalkResult::Advance();
+ };
+ auto fallback_callback = [&](const ObjectRef&, const VisitCount& count,
const VisitTag& tag,
+ TVMFFIDefRegionKind kind) ->
Expected<WalkResult> {
+ counts.push_back(count.value);
+ tags.push_back(tag.value);
+ kinds.push_back(kind);
+ return WalkResult::Advance();
+ };
+ using PreOrderEngine = StructuralWalkEngine<StructuralWalkWithVisitCount<>,
WalkOrder::kPreOrder,
+ decltype(pair_callback),
decltype(fallback_callback)>;
+ StructuralVisitor visitor(
+ make_object<PreOrderEngine>(std::move(pair_callback),
std::move(fallback_callback)));
+ Optional<VisitInterrupt> result = visitor->VisitExpected(root).value();
+
+ EXPECT_FALSE(result.has_value());
+ EXPECT_EQ(counts, (std::vector<int>{0, 1, 2}));
+ EXPECT_EQ(tags, (std::vector<int>{7, 7, 7}));
+ EXPECT_EQ(kinds,
+ (std::vector<TVMFFIDefRegionKind>{kTVMFFIDefRegionKindNone,
kTVMFFIDefRegionKindNone,
+ kTVMFFIDefRegionKindNone}));
+ counts.clear();
+
+ auto interrupt_callback = [&](const ObjectRef& node, const VisitCount& count,
+ const VisitTag&) -> Expected<WalkResult> {
+ counts.push_back(count.value);
+ if (node.same_as(lhs)) return WalkResult::Skip();
+ if (node.same_as(rhs)) return
WalkResult::Interrupt(VisitInterrupt(String("done")));
+ return WalkResult::Advance();
+ };
+ using InterruptEngine = StructuralWalkEngine<StructuralWalkWithVisitCount<>,
WalkOrder::kPreOrder,
+ decltype(interrupt_callback)>;
+ StructuralVisitor
interrupt_visitor(make_object<InterruptEngine>(std::move(interrupt_callback)));
+ Optional<VisitInterrupt> interrupt_result =
interrupt_visitor->VisitExpected(root).value();
+
+ ASSERT_TRUE(interrupt_result.has_value());
+ EXPECT_EQ(interrupt_result.value()->value.cast<String>(), "done");
+ EXPECT_EQ(counts, (std::vector<int>{0, 1, 1}));
+
+ counts.clear();
+ auto postorder_callback = [&](const ObjectRef&, const VisitCount& count,
+ const VisitTag&) -> Expected<WalkResult> {
+ counts.push_back(count.value);
+ return WalkResult::Advance();
+ };
+ using PostOrderEngine = StructuralWalkEngine<StructuralWalkWithVisitCount<>,
+ WalkOrder::kPostOrder,
decltype(postorder_callback)>;
+ StructuralVisitor
postorder_visitor(make_object<PostOrderEngine>(std::move(postorder_callback)));
+ Optional<VisitInterrupt> postorder_result =
postorder_visitor->VisitExpected(root).value();
+ EXPECT_FALSE(postorder_result.has_value());
+ EXPECT_EQ(counts, (std::vector<int>{2, 3, 3}));
+
+ auto error_callback = [&](const ObjectRef& node, const VisitCount&,
+ const VisitTag&) -> Expected<WalkResult> {
+ if (node.same_as(lhs)) {
+ return Unexpected(Error("ValueError", "stateful walk failed", ""));
+ }
+ return WalkResult::Advance();
+ };
+ using ErrorEngine = StructuralWalkEngine<StructuralWalkWithVisitCount<>,
WalkOrder::kPreOrder,
+ decltype(error_callback)>;
+ StructuralVisitor
error_visitor(make_object<ErrorEngine>(std::move(error_callback)));
+ Expected<Optional<VisitInterrupt>> error_result =
error_visitor->VisitExpected(root);
+ ASSERT_TRUE(error_result.is_err());
+ EXPECT_EQ(error_result.error().kind(), "ValueError");
+ EXPECT_EQ(error_result.error().message(), "stateful walk failed");
+}
+
TEST(StructuralVisitor, WalkReturnsError) {
ObjectRef lhs = TVar("lhs");
ObjectRef rhs = TVar("rhs");
