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 6118c1fe [FEAT] Support optional structural visit early returns (#790)
6118c1fe is described below
commit 6118c1febc8eba8c486ce2f639d4c95e118e8bba
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Sep 14 19:54:47 2026 -0400
[FEAT] Support optional structural visit early returns (#790)
Allow exception-first structural visitors to use
`TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN` with `Optional<VisitInterrupt>`.
`VisitReturnHelper` forwards optional interrupts directly and preserves
the existing Expected return proxy for typed helpers and raw hooks.
Includes a focused regression for optional continuation and interrupt
propagation.
---
include/tvm/ffi/extra/structural_visit.h | 41 ++++++++++++++++++++++++--------
tests/cpp/extra/test_structural_visit.cc | 20 ++++++++++++++++
2 files changed, 51 insertions(+), 10 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index a1f1d55a..d6d2ced7 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -306,6 +306,27 @@ class StructuralVisitor : public ObjectRef {
namespace details {
+/*!
+ * \brief Return an optional visit result directly or wrap an Expected for
return conversion.
+ */
+template <typename T>
+TVM_FFI_INLINE auto VisitReturnHelper(T&& result) {
+ if constexpr (std::is_same_v<std::remove_cv_t<std::remove_reference_t<T>>,
+ Optional<VisitInterrupt>>) {
+ return std::forward<T>(result);
+ } else {
+ return ExpectedReturnHelper(std::forward<T>(result));
+ }
+}
+
+/*!
+ * \brief Return true when \p result carries a traversal-stopping interrupt.
+ */
+TVM_FFI_INLINE bool StructuralVisitNeedEarlyReturn(
+ const Optional<VisitInterrupt>& result) noexcept {
+ return result.has_value();
+}
+
/*!
* \brief Return true when \p result already carries a traversal-stopping
state.
* \tparam T The Expected success type.
@@ -501,9 +522,9 @@ namespace details {
* \brief Return from a visit hook if \p Result stops traversal.
*
* Propagates an ``Error`` or a ``VisitInterrupt`` out of the enclosing
function
- * and otherwise falls through. Works from a raw ``TVMFFIAny`` hook and from a
- * typed ``Expected`` helper alike; the rvalue-only proxy lets the return type
- * select the representation.
+ * and otherwise falls through. An ``Optional<VisitInterrupt>`` result is
returned
+ * directly. An ``Expected`` result uses an rvalue-only proxy to select the
return
+ * representation for raw ``TVMFFIAny`` hooks or typed ``Expected`` helpers.
*
* A registered ``__s_visit__`` hook is one line per traversed field followed
by
* the terminal return. A field skipped on purpose is guarded by a condition
and
@@ -521,13 +542,13 @@ namespace details {
*
* \param Result An expression yielding the descent result to inspect.
*/
-#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::ExpectedReturnHelper(::std::move(tvm_ffi_res_)); \
- }
\
+#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::VisitReturnHelper(::std::move(tvm_ffi_res_)); \
+ }
\
} while (0)
} // namespace details
diff --git a/tests/cpp/extra/test_structural_visit.cc
b/tests/cpp/extra/test_structural_visit.cc
index 7c015053..b638c128 100644
--- a/tests/cpp/extra/test_structural_visit.cc
+++ b/tests/cpp/extra/test_structural_visit.cc
@@ -676,6 +676,26 @@ TEST(StructuralVisitor, WalkAnyFallback) {
// StructuralVisit behavior.
// ---------------------------------------------------------------------------
+TEST(StructuralVisit, OptionalEarlyReturn) {
+ bool continued = false;
+ auto visit = [&](Optional<VisitInterrupt> result) ->
Optional<VisitInterrupt> {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(result);
+ continued = true;
+ return std::nullopt;
+ };
+
+ EXPECT_FALSE(visit(std::nullopt).has_value());
+ EXPECT_TRUE(continued);
+
+ continued = false;
+ VisitInterrupt interrupt(String("stop"));
+ Optional<VisitInterrupt> result = visit(interrupt);
+ ASSERT_TRUE(result.has_value());
+ EXPECT_TRUE(result.value().same_as(interrupt));
+ EXPECT_EQ(result.value()->value.cast<String>(), "stop");
+ EXPECT_FALSE(continued);
+}
+
TEST(StructuralVisit, CallbackDrivenTraversal) {
TVarWithDep lhs("lhs", TVarWithDep("pruned-dependency"));
TVarWithDep stop("stop");
