This is an automated email from the ASF dual-hosted git repository.
tlopex 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 874c85c5 [FEAT][Rust] Add composable policies for structural visit and
walk (#794)
874c85c5 is described below
commit 874c85c537fc6638dafe584224a770bb008b7c8f
Author: Shushi Hong <[email protected]>
AuthorDate: Wed Sep 16 18:10:09 2026 -0400
[FEAT][Rust] Add composable policies for structural visit and walk (#794)
Add `VisitPolicy<State>` to customize default recursion independently of
callbacks, corresponding to the C++ Parent extension mechanism. Attach
policies through `VisitCallbacks::with_policy` or `WalkWithPolicy`;
nested tuples compose policies sharing the callback state.
Add `VisitContext::default_visit_children(value, kind)` to continue the
policy chain with an explicit value and definition region, bypassing the
target's callback while its children re-enter full dispatch. Preserve
Pattern precedence and restore the enclosing region on return, error, or
interrupt. Keep `visit_children()` as the current-value/current-region
shorthand.
Tests cover composition, shared state, generated walk dispatch, callback
ordering, Skip, policy interrupts/errors, retargeted descent, and region
propagation through reflected fields and Function-valued hooks.
---
docs/guides/rust_lang_guide.md | 4 +
rust/tvm-ffi/src/extra/structural_visit.rs | 175 +++++--
rust/tvm-ffi/src/extra/structural_visit/policy.rs | 272 ++++++++++
rust/tvm-ffi/src/lib.rs | 7 +-
rust/tvm-ffi/tests/test_structural_visit.rs | 602 +++++++++++++++++++++-
src/ffi/testing/testing.cc | 44 ++
6 files changed, 1063 insertions(+), 41 deletions(-)
diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md
index 3d700085..2bc2764c 100644
--- a/docs/guides/rust_lang_guide.md
+++ b/docs/guides/rust_lang_guide.md
@@ -366,6 +366,10 @@ Callbacks are `Fn`; mutable data belongs in the visitor
state. A catch-all
callback must call `visit_children()` explicitly, and interrupt values must be
returned explicitly because `?` only propagates errors.
+`VisitCallbacks::with_policy` and `WalkWithPolicy` customize default recursion
+with reusable `VisitPolicy<State>` policies. See the `VisitPolicy` API
+documentation for policy composition and state access.
+
For a named implementation, `#[dispatch(visit)]` generates
`StructuralVisitor` from `visit_*` methods. Matching handlers own recursion;
unmatched values use default child traversal:
diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs
b/rust/tvm-ffi/src/extra/structural_visit.rs
index b6cef1c8..0fc337ad 100644
--- a/rust/tvm-ffi/src/extra/structural_visit.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit.rs
@@ -57,7 +57,7 @@ use crate::function::Function;
use crate::object::{Object, ObjectArc, ObjectCore};
use crate::reflection::TypeAttrColumn;
use crate::tvm_ffi_sys::TVMFFIFieldFlagBitMask::{
- kTVMFFIFieldFlagBitMaskSEqHashDefSimple,
kTVMFFIFieldFlagBitMaskSEqHashDefPattern,
+ kTVMFFIFieldFlagBitMaskSEqHashDefPattern,
kTVMFFIFieldFlagBitMaskSEqHashDefSimple,
kTVMFFIFieldFlagBitMaskSEqHashIgnore,
};
use crate::tvm_ffi_sys::{
@@ -141,13 +141,9 @@ pub enum DefRegionKind {
const _: () = {
assert!(DefRegionKind::None as i32 ==
TVMFFIDefRegionKind::kTVMFFIDefRegionKindNone as i32);
assert!(
- DefRegionKind::Pattern as i32
- == TVMFFIDefRegionKind::kTVMFFIDefRegionKindPattern as i32
- );
- assert!(
- DefRegionKind::Simple as i32
- == TVMFFIDefRegionKind::kTVMFFIDefRegionKindSimple as i32
+ DefRegionKind::Pattern as i32 ==
TVMFFIDefRegionKind::kTVMFFIDefRegionKindPattern as i32
);
+ assert!(DefRegionKind::Simple as i32 ==
TVMFFIDefRegionKind::kTVMFFIDefRegionKindSimple as i32);
};
/// Interrupt state of a traversal, mirroring C++ `ffi.VisitInterrupt`.
@@ -228,6 +224,9 @@ impl From<Error> for NativeHalt {
type NativeResult = std::result::Result<(), NativeHalt>;
+mod policy;
+pub use policy::{DefaultVisitPolicy, VisitPolicy, WalkWithPolicy};
+
/// State and recursive operations available to a visit callback.
///
/// A matched callback owns traversal of its value. Recursive operations
@@ -299,11 +298,51 @@ impl<State> VisitContext<'_, State> {
self.driver.visit_raw(raw, def_region_kind)
}
- /// Visit the current value's children using registered hooks or reflected
- /// structural fields. The current value itself is not dispatched again.
+ /// Apply default descent without dispatching the current value again.
+ /// A callback enters its configured policy; within a policy this continues
+ /// with the next policy, then registered hooks or reflected fields.
pub fn visit_children(&mut self) -> Result<Option<VisitInterrupt>> {
- self.driver
- .visit_children_raw(self.current.raw(), self.def_region_kind)
+ self.default_visit_children_raw(self.current.raw(),
self.def_region_kind)
+ }
+
+ /// Apply default descent to `value` under an explicit definition region.
+ ///
+ /// Like [`Self::visit_children`], this continues with the next policy (or
+ /// enters the configured policy from a callback), then registered hooks or
+ /// reflected fields. It does not dispatch callbacks for `value` itself;
+ /// its children re-enter the full callback engine.
+ ///
+ /// An enclosing [`DefRegionKind::Pattern`] cannot be downgraded. The
current
+ /// value and region of this context are unchanged after the call,
including
+ /// when descent returns an error or interrupt.
+ pub fn default_visit_children<T>(
+ &mut self,
+ value: &T,
+ def_region_kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>>
+ where
+ for<'x> AnyView<'x>: From<&'x T>,
+ {
+ self.default_visit_children_raw(raw_of(AnyView::from(value)),
def_region_kind)
+ }
+
+ fn default_visit_children_raw(
+ &mut self,
+ raw: TVMFFIAny,
+ def_region_kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
+ return Ok(None);
+ }
+ let kind = if self.def_region_kind == DefRegionKind::Pattern {
+ DefRegionKind::Pattern
+ } else {
+ def_region_kind
+ };
+ let active = active_structural_visitor()?;
+ // Keep the ABI visitor in sync while policies run, not only inside
+ // hooks: a subsequent visit_with must also preserve a pattern region.
+ with_visitor_def_region(active, kind, ||
self.driver.visit_children_raw(raw, kind))
}
}
@@ -473,7 +512,8 @@ macro_rules! impl_visit_chain_link {
impl_callback_chain_tuple_arities!(impl_visit_chain_link);
/// A reusable callback visitor with shared user state.
-pub struct VisitCallbacks<State, Link, Marker> {
+pub struct VisitCallbacks<State, Link, Marker, Policy = DefaultVisitPolicy> {
+ policy: Option<Rc<Policy>>,
state: State,
callbacks: Rc<Link>,
_marker: PhantomData<fn(Marker)>,
@@ -488,12 +528,27 @@ where
Self {
state,
callbacks: Rc::new(callbacks),
+ policy: None,
_marker: PhantomData,
}
}
}
-impl<State, Link, Marker> VisitCallbacks<State, Link, Marker> {
+impl<State, Link, Marker, Policy> VisitCallbacks<State, Link, Marker, Policy> {
+ /// Set the default-recursion policy while retaining the callbacks and
state.
+ /// A matched callback enters the policy only when it calls
`visit_children()`.
+ pub fn with_policy<P: VisitPolicy<State>>(
+ self,
+ policy: P,
+ ) -> VisitCallbacks<State, Link, Marker, P> {
+ VisitCallbacks {
+ state: self.state,
+ callbacks: self.callbacks,
+ policy: Some(Rc::new(policy)),
+ _marker: PhantomData,
+ }
+ }
+
/// Shared access to the callback state.
pub fn state(&self) -> &State {
&self.state
@@ -521,7 +576,9 @@ trait VisitCallbackState<State> {
fn callback_state_mut(&mut self) -> &mut State;
}
-impl<State, Link, Marker> VisitCallbackState<State> for VisitCallbacks<State,
Link, Marker> {
+impl<State, Link, Marker, Policy> VisitCallbackState<State>
+ for VisitCallbacks<State, Link, Marker, Policy>
+{
fn callback_state(&self) -> &State {
&self.state
}
@@ -979,19 +1036,23 @@ pub trait StructuralVisitor: Sized {
value: &VisitValue,
def_region_kind: DefRegionKind,
) -> Result<Option<VisitInterrupt>> {
- let raw = value.raw();
- let context = std::ptr::from_mut(&mut *self).cast::<c_void>();
- let result = visit_children_raw(
- raw,
- &mut UserChildren { visitor: self },
- context,
- def_region_kind,
- )
- .map_err(|halt| with_value_context(halt, raw));
- finish(result)
+ default_user_visit_children(self, value, def_region_kind)
}
}
+fn default_user_visit_children<V: StructuralVisitor>(
+ visitor: &mut V,
+ value: &VisitValue,
+ def_region_kind: DefRegionKind,
+) -> Result<Option<VisitInterrupt>> {
+ let raw = value.raw();
+ let context = std::ptr::from_mut(&mut *visitor).cast::<c_void>();
+ finish(
+ visit_children_raw(raw, &mut UserChildren { visitor }, context,
def_region_kind)
+ .map_err(|halt| with_value_context(halt, raw)),
+ )
+}
+
fn try_visit_callbacks<State, Link, Marker>(
driver: &mut impl VisitContextDriver<State>,
callback_ptr: *const Link,
@@ -1015,9 +1076,10 @@ where
}
}
-impl<State, Link, Marker> StructuralVisitor for VisitCallbacks<State, Link,
Marker>
+impl<State, Link, Marker, Policy> StructuralVisitor for VisitCallbacks<State,
Link, Marker, Policy>
where
Link: VisitChainLink<State, Marker>,
+ Policy: VisitPolicy<State>,
{
fn visit(
&mut self,
@@ -1027,6 +1089,22 @@ where
let callback_ptr = Rc::as_ptr(&self.callbacks);
try_visit_callbacks::<State, Link, Marker>(self, callback_ptr, value,
def_region_kind)
}
+
+ fn default_visit_children(
+ &mut self,
+ value: &VisitValue,
+ def_region_kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ let Some(policy) = self.policy.as_ref().map(Rc::clone) else {
+ return default_user_visit_children(self, value, def_region_kind);
+ };
+ policy::visit_with_policy(
+ &mut policy::VisitDescent { visitor: self },
+ &*policy,
+ value,
+ def_region_kind,
+ )
+ }
}
impl<Link, Marker> StructuralVisitor for DirectVisitCallbacks<'_, Link, Marker>
@@ -1085,8 +1163,32 @@ where
/// Internal callback protocol used by [`IntoWalker`].
#[doc(hidden)]
-pub trait NativeVisit {
+pub trait NativeVisit: Sized {
+ const CUSTOM_DESCENT: bool = false;
+
fn visit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) ->
Result<WalkResult>;
+
+ fn default_visit_children<const PRE_ORDER: bool>(
+ &mut self,
+ value: &VisitValue,
+ def_region_kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ default_walk_children::<Self, PRE_ORDER>(self, value, def_region_kind)
+ }
+}
+
+fn default_walk_children<V: NativeVisit, const PRE_ORDER: bool>(
+ visitor: &mut V,
+ value: &VisitValue,
+ def_region_kind: DefRegionKind,
+) -> Result<Option<VisitInterrupt>> {
+ let context = std::ptr::from_mut(&mut *visitor).cast::<c_void>();
+ finish(visit_children_raw(
+ value.raw(),
+ &mut WalkChildren::<V, PRE_ORDER> { visitor },
+ context,
+ def_region_kind,
+ ))
}
/// Action applied to each child found by the shared traversal.
@@ -1147,12 +1249,19 @@ fn visit_raw<V: NativeVisit, const PRE_ORDER: bool>(
}
}
- let context = std::ptr::from_mut(&mut *visitor).cast::<c_void>();
- let children = &mut WalkChildren::<V, PRE_ORDER> {
- visitor: &mut *visitor,
- };
- if let Err(halt) = visit_children_raw(value, children, context,
def_region_kind) {
- return Err(with_value_context(halt, value));
+ if V::CUSTOM_DESCENT {
+ match visitor.default_visit_children::<PRE_ORDER>(&visit_value,
def_region_kind) {
+ Ok(None) => {}
+ Ok(Some(interrupt)) => return
Err(NativeHalt::Interrupt(interrupt.value)),
+ Err(error) => return Err(with_value_context(error.into(), value)),
+ }
+ } else {
+ // Preserve the raw-result path for ordinary walkers.
+ let context = std::ptr::from_mut(&mut *visitor).cast::<c_void>();
+ let children = &mut WalkChildren::<V, PRE_ORDER> { visitor };
+ if let Err(halt) = visit_children_raw(value, children, context,
def_region_kind) {
+ return Err(with_value_context(halt, value));
+ }
}
if PRE_ORDER {
@@ -1567,7 +1676,7 @@ unsafe fn runtime_walk<V: NativeVisit, const PRE_ORDER:
bool>(
if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
return Ok(());
}
- if raw.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ if !V::CUSTOM_DESCENT && raw.type_index <
TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 {
let visitor = &mut *context.cast::<V>();
if PRE_ORDER {
match visitor.visit(&VisitValue::from_raw(raw), def_region_kind) {
diff --git a/rust/tvm-ffi/src/extra/structural_visit/policy.rs
b/rust/tvm-ffi/src/extra/structural_visit/policy.rs
new file mode 100644
index 00000000..e2372ee2
--- /dev/null
+++ b/rust/tvm-ffi/src/extra/structural_visit/policy.rs
@@ -0,0 +1,272 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+//! Reusable customization of default structural descent.
+
+use super::*;
+
+/// A reusable default-recursion policy sharing state with traversal callbacks.
+///
+/// `visit_children()` on this policy's context continues with the next policy
+/// (or the built-in hooks and reflected fields). `visit()` re-enters the full
+/// callback engine for a child. A tuple `(outer, inner)` composes two
policies;
+/// tuples may nest. Policies are shared during recursive calls, so mutable
data
+/// belongs in the context's state.
+pub trait VisitPolicy<State> {
+ /// Customize default descent for the current value.
+ ///
+ /// Return interrupts explicitly, and restore any scoped state before
+ /// returning an interrupt or error. Walk invokes this between its pre- and
+ /// post-order callback positions; visit invokes it only on callback miss
+ /// or when a matched callback requests default descent.
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ visitor: &mut VisitContext<'_, State>,
+ ) -> Result<Option<VisitInterrupt>>;
+}
+
+/// Default descent through registered hooks or reflected structural fields.
+pub struct DefaultVisitPolicy;
+
+impl<State> VisitPolicy<State> for DefaultVisitPolicy {
+ fn default_visit(
+ &self,
+ _value: &VisitValue,
+ visitor: &mut VisitContext<'_, State>,
+ ) -> Result<Option<VisitInterrupt>> {
+ visitor.visit_children()
+ }
+}
+
+impl<State, Outer: VisitPolicy<State>, Inner: VisitPolicy<State>>
VisitPolicy<State>
+ for (Outer, Inner)
+{
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ visitor: &mut VisitContext<'_, State>,
+ ) -> Result<Option<VisitInterrupt>> {
+ let kind = visitor.def_region_kind();
+ visit_with_policy(
+ &mut NextPolicy {
+ driver: &mut *visitor.driver,
+ policy: &self.1,
+ },
+ &self.0,
+ value,
+ kind,
+ )
+ }
+}
+
+pub(super) fn visit_with_policy<State>(
+ driver: &mut dyn VisitContextDriver<State>,
+ policy: &impl VisitPolicy<State>,
+ value: &VisitValue,
+ def_region_kind: DefRegionKind,
+) -> Result<Option<VisitInterrupt>> {
+ policy.default_visit(
+ value,
+ &mut VisitContext {
+ driver,
+ current: VisitValue::from_raw(value.raw()),
+ def_region_kind,
+ _not_send_sync: PhantomData,
+ },
+ )
+}
+
+struct NextPolicy<'a, State, Policy> {
+ driver: &'a mut dyn VisitContextDriver<State>,
+ policy: &'a Policy,
+}
+
+impl<State, Policy: VisitPolicy<State>> VisitContextDriver<State>
+ for NextPolicy<'_, State, Policy>
+{
+ fn state(&self) -> &State {
+ self.driver.state()
+ }
+ fn state_mut(&mut self) -> &mut State {
+ self.driver.state_mut()
+ }
+ fn visit_raw(&mut self, raw: TVMFFIAny, kind: DefRegionKind) ->
Result<Option<VisitInterrupt>> {
+ self.driver.visit_raw(raw, kind)
+ }
+ fn visit_children_raw(
+ &mut self,
+ raw: TVMFFIAny,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ visit_with_policy(self.driver, self.policy,
&VisitValue::from_raw(raw), kind)
+ }
+}
+
+pub(super) struct VisitDescent<'a, V> {
+ pub(super) visitor: &'a mut V,
+}
+
+impl<State, V: StructuralVisitor + VisitCallbackState<State>>
VisitContextDriver<State>
+ for VisitDescent<'_, V>
+{
+ fn state(&self) -> &State {
+ self.visitor.callback_state()
+ }
+ fn state_mut(&mut self) -> &mut State {
+ self.visitor.callback_state_mut()
+ }
+ fn visit_raw(&mut self, raw: TVMFFIAny, kind: DefRegionKind) ->
Result<Option<VisitInterrupt>> {
+ VisitContextDriver::visit_raw(self.visitor, raw, kind)
+ }
+ fn visit_children_raw(
+ &mut self,
+ raw: TVMFFIAny,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ // Bypass the current policy; children still re-enter the complete
visitor.
+ default_user_visit_children(self.visitor, &VisitValue::from_raw(raw),
kind)
+ }
+}
+
+/// A walk dispatcher combined with a reusable default-recursion policy.
+///
+/// The dispatcher is also the state visible through the policy's context. Use
+/// `#[dispatch(walk)]` or implement [`WalkDispatch`] to define its callbacks.
+/// Run repeatedly with [`Self::walk`], or pass this value to
[`structural_walk`].
+pub struct WalkWithPolicy<Walker, Policy> {
+ walker: Walker,
+ policy: Rc<Policy>,
+}
+
+impl<Walker: WalkDispatch, Policy: VisitPolicy<Walker>> WalkWithPolicy<Walker,
Policy> {
+ /// Combine a dispatcher and a default-recursion policy.
+ pub fn new(walker: Walker, policy: Policy) -> Self {
+ Self {
+ walker,
+ policy: Rc::new(policy),
+ }
+ }
+
+ /// Access the dispatcher and its traversal state.
+ pub fn state(&self) -> &Walker {
+ &self.walker
+ }
+
+ /// Mutably access the dispatcher outside an active walk.
+ pub fn state_mut(&mut self) -> &mut Walker {
+ &mut self.walker
+ }
+
+ /// Recover the dispatcher and its state.
+ pub fn into_state(self) -> Walker {
+ self.walker
+ }
+
+ /// Walk a root using this dispatcher and policy.
+ pub fn walk<R>(&mut self, root: &R, order: WalkOrder) ->
Result<Option<VisitInterrupt>>
+ where
+ for<'x> AnyView<'x>: From<&'x R>,
+ {
+ let raw = raw_of(AnyView::from(root));
+ finish(match order {
+ WalkOrder::PreOrder => {
+ run_structural_visitor(raw, self, walk_runtime_vtable::<Self,
true>())
+ }
+ WalkOrder::PostOrder => {
+ run_structural_visitor(raw, self, walk_runtime_vtable::<Self,
false>())
+ }
+ })
+ }
+}
+
+#[doc(hidden)]
+pub enum ByPolicyWalk {}
+
+impl<Walker: WalkDispatch, Policy: VisitPolicy<Walker>>
IntoWalker<ByPolicyWalk>
+ for WalkWithPolicy<Walker, Policy>
+{
+ type Walker = Self;
+ fn into_walker(self) -> Self {
+ self
+ }
+}
+
+impl<Walker: WalkDispatch, Policy: VisitPolicy<Walker>> NativeVisit
+ for WalkWithPolicy<Walker, Policy>
+{
+ const CUSTOM_DESCENT: bool = true;
+
+ fn visit(&mut self, value: &VisitValue, kind: DefRegionKind) ->
Result<WalkResult> {
+ self.walker
+ .dispatch_walk(value, kind)
+ .unwrap_or(Ok(WalkResult::Advance))
+ }
+
+ fn default_visit_children<const PRE_ORDER: bool>(
+ &mut self,
+ value: &VisitValue,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ let policy = Rc::clone(&self.policy);
+ let active = active_structural_visitor()?;
+ // Reflected children arrive directly from the Rust walker. Synchronize
+ // the ABI region before the policy can re-enter callback dispatch.
+ with_visitor_def_region(active, kind, || {
+ visit_with_policy(
+ &mut WalkDescent::<_, _, PRE_ORDER> { visitor: self },
+ &*policy,
+ value,
+ kind,
+ )
+ })
+ }
+}
+
+struct WalkDescent<'a, Walker, Policy, const PRE_ORDER: bool> {
+ visitor: &'a mut WalkWithPolicy<Walker, Policy>,
+}
+
+impl<Walker: WalkDispatch, Policy: VisitPolicy<Walker>, const PRE_ORDER: bool>
+ VisitContextDriver<Walker> for WalkDescent<'_, Walker, Policy, PRE_ORDER>
+{
+ fn state(&self) -> &Walker {
+ &self.visitor.walker
+ }
+ fn state_mut(&mut self) -> &mut Walker {
+ &mut self.visitor.walker
+ }
+ fn visit_raw(&mut self, raw: TVMFFIAny, kind: DefRegionKind) ->
Result<Option<VisitInterrupt>> {
+ if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
+ return Ok(None);
+ }
+ let active = active_structural_visitor()?;
+ let context = std::ptr::from_mut(&mut *self.visitor).cast::<c_void>();
+ finish(with_current_visitor_context(active, context, || {
+ call_visitor(active, raw, kind)
+ }))
+ }
+ fn visit_children_raw(
+ &mut self,
+ raw: TVMFFIAny,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ default_walk_children::<_, PRE_ORDER>(self.visitor,
&VisitValue::from_raw(raw), kind)
+ }
+}
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index 107399ff..dc44cedc 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -54,9 +54,10 @@ pub use crate::extra::structural_mutate::{
MutateContext, MutateDispatch, Mutator, StructuralMutator,
StructuralVarRemap,
};
pub use crate::extra::structural_visit::{
- structural_visit, structural_walk, DefRegionKind, IntoVisitor,
IntoWalkResult, IntoWalker,
- StructuralVisitor, VisitCallbacks, VisitChainLink, VisitContext,
VisitInterrupt, VisitValue,
- WalkChainLink, WalkDispatch, WalkOrder, WalkResult,
+ structural_visit, structural_walk, DefRegionKind, DefaultVisitPolicy,
IntoVisitor,
+ IntoWalkResult, IntoWalker, StructuralVisitor, VisitCallbacks,
VisitChainLink, VisitContext,
+ VisitInterrupt, VisitPolicy, VisitValue, WalkChainLink, WalkDispatch,
WalkOrder, WalkResult,
+ WalkWithPolicy,
};
pub use crate::extra::unchanged::{Unchanged, UnchangedOr};
pub use crate::function::Function;
diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs
b/rust/tvm-ffi/tests/test_structural_visit.rs
index f2a67f47..9e73b206 100644
--- a/rust/tvm-ffi/tests/test_structural_visit.rs
+++ b/rust/tvm-ffi/tests/test_structural_visit.rs
@@ -23,13 +23,608 @@ use tvm_ffi::{
dispatch, get_type_attr, structural_visit, structural_walk, Any, Array,
DLDataType,
DLDataTypeCode, DefRegionKind, Error, FieldGetter, Function, Map, Object,
ObjectRefCore,
Result, String as FfiString, StructuralVisitor, TypeIndex, VisitCallbacks,
VisitContext,
- VisitInterrupt, VisitValue, WalkOrder, WalkResult, RUNTIME_ERROR,
+ VisitInterrupt, VisitPolicy, VisitValue, WalkOrder, WalkResult,
WalkWithPolicy, RUNTIME_ERROR,
};
fn runtime_error(message: &str) -> Error {
Error::new(RUNTIME_ERROR, message, "")
}
+#[test]
+fn composed_policies_share_array_scope_with_visit_and_walk_callbacks() {
+ #[derive(Default)]
+ struct CollectIntegers {
+ depth: usize,
+ integers: Vec<(i64, usize)>,
+ descent_depths: Vec<usize>,
+ }
+
+ impl CollectIntegers {
+ fn record(&mut self, value: i64) {
+ self.integers.push((value, self.depth));
+ }
+ }
+
+ // The outer policy establishes a scope around an array's children.
+ // It does not know what the callbacks will do with the current depth.
+ struct ArrayScope;
+
+ impl VisitPolicy<CollectIntegers> for ArrayScope {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ visitor: &mut VisitContext<'_, CollectIntegers>,
+ ) -> Result<Option<VisitInterrupt>> {
+ if value.cast::<Array<Any>>().is_none() {
+ return visitor.visit_children();
+ }
+ let outer_depth = visitor.state().depth;
+ let depth = outer_depth + 1;
+ visitor.state_mut().depth = depth;
+
+ // Continue to RecordDescent, not directly to the children.
+ let result = visitor.visit_children();
+
+ visitor.state_mut().depth = outer_depth;
+ result
+ }
+ }
+
+ // The inner policy observes the scope set by ArrayScope. Its continuation
+ // reaches the built-in Array hook, whose children re-enter the full
engine.
+ struct RecordDescent;
+
+ impl VisitPolicy<CollectIntegers> for RecordDescent {
+ fn default_visit(
+ &self,
+ _value: &VisitValue,
+ visitor: &mut VisitContext<'_, CollectIntegers>,
+ ) -> Result<Option<VisitInterrupt>> {
+ let state = visitor.state_mut();
+ state.descent_depths.push(state.depth);
+ visitor.visit_children()
+ }
+ }
+
+ // [1, [2, 3], 4]: the final sibling must see the restored outer depth.
+ let root = Array::new(vec![
+ Any::from(1_i64),
+ Any::from(Array::new(vec![2_i64, 3])),
+ Any::from(4_i64),
+ ]);
+ let expected = vec![(1, 1), (2, 2), (3, 2), (4, 1)];
+
+ let mut visitor = VisitCallbacks::new(
+ CollectIntegers::default(),
+ |value: i64, visitor: &mut VisitContext<'_, CollectIntegers>| {
+ visitor.state_mut().record(value);
+ // Matched visit callbacks own recursion: an integer needs no
descent.
+ },
+ )
+ .with_policy((ArrayScope, RecordDescent));
+ assert!(structural_visit(&root, &mut visitor).unwrap().is_none());
+ assert_eq!(visitor.state().integers, expected);
+ assert_eq!(visitor.state().depth, 0);
+ // Only unmatched arrays enter the default policies in this visit.
+ assert_eq!(visitor.state().descent_depths, vec![1, 2]);
+
+ #[dispatch(walk)]
+ impl CollectIntegers {
+ fn walk_integer(&mut self, value: i64) -> WalkResult {
+ self.record(value);
+ WalkResult::Advance
+ }
+ }
+
+ // Reuse the policies with a macro-generated walk dispatcher. Walk manages
+ // recursion, so default policies also run for matched integer leaves.
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut walker =
+ WalkWithPolicy::new(CollectIntegers::default(), (ArrayScope,
RecordDescent));
+ assert!(walker.walk(&root, order).unwrap().is_none());
+ assert_eq!(walker.state().integers, expected);
+ assert_eq!(walker.state().depth, 0);
+ assert_eq!(walker.state().descent_depths, vec![1, 1, 2, 2, 2, 1]);
+ }
+}
+
+// Enter Simple, then Pattern, then attempt to downgrade to None. Verify each
+// enclosing region is restored even when a child interrupts or returns an
error.
+#[test]
+fn policy_continuation_scopes_regions_and_restores_after_halts() {
+ #[derive(Clone, Copy)]
+ enum Outcome {
+ Finish,
+ Interrupt,
+ Error,
+ }
+ struct Probe {
+ outcome: Outcome,
+ seen: Vec<(i64, DefRegionKind)>,
+ regions: Vec<DefRegionKind>,
+ }
+ impl Probe {
+ fn record(&mut self, value: i64, kind: DefRegionKind) ->
Result<Option<VisitInterrupt>> {
+ self.seen.push((value, kind));
+ if value == 1 {
+ match self.outcome {
+ Outcome::Interrupt => return
Ok(Some(VisitInterrupt::with(42_i64))),
+ Outcome::Error => return Err(runtime_error("continuation
failed")),
+ Outcome::Finish => {}
+ }
+ }
+ Ok(None)
+ }
+ }
+ #[dispatch(walk)]
+ impl Probe {
+ fn walk_integer(&mut self, value: i64, kind: DefRegionKind) ->
Result<WalkResult> {
+ Ok(match self.record(value, kind)? {
+ Some(interrupt) => WalkResult::InterruptWith(interrupt.value),
+ None => WalkResult::Advance,
+ })
+ }
+ }
+ struct Scope(DefRegionKind, i64);
+ impl VisitPolicy<Probe> for Scope {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ ctx: &mut VisitContext<'_, Probe>,
+ ) -> Result<Option<VisitInterrupt>> {
+ let Some(array) = value.cast::<Array<i64>>() else {
+ return ctx.visit_children();
+ };
+ let outer = ctx.def_region_kind();
+ ctx.state_mut().regions.push(outer);
+ let result = ctx.default_visit_children(&array, self.0);
+ assert_eq!(ctx.def_region_kind(), outer);
+ // Re-dispatch checks the restored ABI region, not only the
context.
+ assert!(ctx.visit(&self.1)?.is_none());
+ result
+ }
+ }
+ use DefRegionKind::{None as Use, Pattern, Simple};
+ let root = Array::new(vec![1_i64, 5]);
+ for outcome in [Outcome::Finish, Outcome::Interrupt, Outcome::Error] {
+ for order in [None, Some(WalkOrder::PreOrder),
Some(WalkOrder::PostOrder)] {
+ let state = Probe {
+ outcome,
+ seen: vec![],
+ regions: vec![],
+ };
+ let policies = (Scope(Simple, 4), (Scope(Pattern, 3), Scope(Use,
2)));
+ let (result, state) = if let Some(order) = order {
+ let mut walker = WalkWithPolicy::new(state, policies);
+ let result = walker.walk(&root, order);
+ (result, walker.into_state())
+ } else {
+ let mut visitor =
+ VisitCallbacks::new(state, |value: i64, ctx: &mut
VisitContext<'_, Probe>| {
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().record(value, kind)
+ })
+ .with_policy(policies);
+ let result = structural_visit(&root, &mut visitor);
+ (result, visitor.into_state())
+ };
+ match outcome {
+ Outcome::Finish => assert!(result.unwrap().is_none()),
+ Outcome::Interrupt => {
+
assert_eq!(i64::try_from(result.unwrap().unwrap().value).unwrap(), 42)
+ }
+ Outcome::Error => assert!(result
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("continuation failed")),
+ }
+ assert_eq!(state.regions, vec![Use, Simple, Pattern]);
+ let mut expected = vec![(1, Pattern)];
+ if matches!(outcome, Outcome::Finish) {
+ expected.push((5, Pattern)); // Halts must skip the remaining
child.
+ }
+ // Unwind the innermost, middle, and outer scopes.
+ expected.extend([(2, Pattern), (3, Simple), (4, Use)]);
+ assert_eq!(state.seen, expected);
+ }
+ }
+}
+
+fn visit_region_graph(with_hook: bool) -> Any {
+ assert_eq!(
+ unsafe { tvm_ffi::tvm_ffi_sys::TVMFFITestingDummyTarget() },
+ 0
+ );
+ Function::get_global("testing.make_visit_region_graph")
+ .unwrap()
+ .call_tuple((with_hook,))
+ .unwrap()
+}
+
+#[derive(Default)]
+struct PolicyRegionTrace(Vec<(i64, DefRegionKind)>);
+#[dispatch(walk)]
+impl PolicyRegionTrace {
+ fn walk_integer(&mut self, value: i64, kind: DefRegionKind) -> WalkResult {
+ self.0.push((value, kind));
+ WalkResult::Advance
+ }
+}
+
+// Default descent may target a child container without dispatching its own
+// callback, while keeping subsequent policies and child callback dispatch.
+#[test]
+fn policy_continuation_retargets_without_dispatching_the_container() {
+ #[derive(Default)]
+ struct Probe(Vec<(&'static str, i64, DefRegionKind)>, bool);
+
+ #[dispatch(walk)]
+ impl Probe {
+ fn walk_array(&mut self, value: Array<Any>, kind: DefRegionKind) ->
WalkResult {
+ self.0.push(("array callback", value.len() as i64, kind));
+ if self.1 {
+ WalkResult::Skip
+ } else {
+ WalkResult::Advance
+ }
+ }
+ fn walk_integer(&mut self, value: i64, kind: DefRegionKind) ->
WalkResult {
+ self.0.push(("integer callback", value, kind));
+ WalkResult::Advance
+ }
+ }
+
+ struct Redirect;
+ impl VisitPolicy<Probe> for Redirect {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ ctx: &mut VisitContext<'_, Probe>,
+ ) -> Result<Option<VisitInterrupt>> {
+ let Some(array) = value.cast::<Array<Any>>() else {
+ return ctx.visit_children();
+ };
+ assert_eq!(
+ array.len(),
+ 2,
+ "the redirected array must not restart this policy"
+ );
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push(("before", 2, kind));
+ let target = array.get(0).unwrap();
+ let result = ctx.default_visit_children(&target,
DefRegionKind::Pattern);
+ assert_eq!(ctx.current().cast::<Array<Any>>().unwrap().len(), 2);
+ assert_eq!(ctx.def_region_kind(), DefRegionKind::None);
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push(("after", 2, kind));
+ result
+ }
+ }
+
+ struct Observe;
+ impl VisitPolicy<Probe> for Observe {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ ctx: &mut VisitContext<'_, Probe>,
+ ) -> Result<Option<VisitInterrupt>> {
+ if let Some(array) = value.cast::<Array<Any>>() {
+ assert_eq!(ctx.current().cast::<Array<Any>>().unwrap().len(),
1);
+ let kind = ctx.def_region_kind();
+ ctx.state_mut()
+ .0
+ .push(("next policy", array.len() as i64, kind));
+ } else if value.cast::<i64>().is_none() {
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push(("reflected policy", 0, kind));
+ }
+ ctx.visit_children()
+ }
+ }
+
+ let root = Array::new(vec![Any::from(Array::new(vec![7_i64])),
Any::from(99_i64)]);
+ let descent = vec![
+ ("before", 2, DefRegionKind::None),
+ ("next policy", 1, DefRegionKind::Pattern),
+ ("integer callback", 7, DefRegionKind::Pattern),
+ ("after", 2, DefRegionKind::None),
+ ];
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ for skip in [false, true] {
+ let mut walker = WalkWithPolicy::new(Probe(vec![], skip),
(Redirect, Observe));
+ assert!(walker.walk(&root, order).unwrap().is_none());
+ let mut expected = descent.clone();
+ let array_callback = ("array callback", 2, DefRegionKind::None);
+ match order {
+ WalkOrder::PreOrder => {
+ if skip {
+ expected.clear(); // Skip prevents policy entry as
well as child traversal.
+ }
+ expected.insert(0, array_callback);
+ }
+ WalkOrder::PostOrder => expected.push(array_callback),
+ }
+ assert_eq!(walker.state().0, expected);
+ }
+ }
+ let mut visitor = VisitCallbacks::new(
+ Probe::default(),
+ |value: &VisitValue, ctx: &mut VisitContext<'_, Probe>| {
+ let kind = ctx.def_region_kind();
+ if let Some(array) = value.cast::<Array<Any>>() {
+ ctx.state_mut()
+ .0
+ .push(("array callback", array.len() as i64, kind));
+ ctx.visit_children()
+ } else {
+ ctx.state_mut()
+ .0
+ .push(("integer callback", value.cast::<i64>().unwrap(),
kind));
+ Ok(None)
+ }
+ },
+ )
+ .with_policy((Redirect, Observe));
+ assert!(structural_visit(&root, &mut visitor).unwrap().is_none());
+ let mut expected = descent;
+ expected.insert(0, ("array callback", 2, DefRegionKind::None));
+ assert_eq!(visitor.state().0, expected);
+
+ // A callback can also retarget default descent, entering the configured
+ // policy. This target uses reflection instead of the Array hook above.
+ let target = visit_region_graph(false);
+ let mut visitor = VisitCallbacks::new(
+ Probe::default(),
+ |value: &VisitValue, ctx: &mut VisitContext<'_, Probe>| {
+ if value.cast::<i64>() == Some(-1) {
+ assert!(ctx
+ .default_visit_children(&Any::new(),
DefRegionKind::Pattern)?
+ .is_none());
+ let result = ctx.default_visit_children(&target,
DefRegionKind::Pattern);
+ assert_eq!(ctx.current().cast::<i64>(), Some(-1));
+ assert_eq!(ctx.def_region_kind(), DefRegionKind::None);
+ return result;
+ }
+ // A callback on the reflected target itself would fail this cast.
+ let integer = value.cast::<i64>().unwrap();
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push(("integer callback", integer, kind));
+ Ok(None)
+ },
+ )
+ .with_policy(Observe);
+ assert!(structural_visit(&-1_i64, &mut visitor).unwrap().is_none());
+ assert_eq!(
+ visitor.state().0,
+ vec![
+ ("reflected policy", 0, DefRegionKind::Pattern),
+ ("integer callback", 1, DefRegionKind::Pattern),
+ ("integer callback", 2, DefRegionKind::Pattern),
+ ("integer callback", 3, DefRegionKind::Pattern),
+ ]
+ );
+}
+
+#[test]
+fn policy_halts_skip_remaining_policies_and_restore_outer_state() {
+ #[derive(Default)]
+ struct Probe {
+ events: Vec<&'static str>,
+ depth: usize,
+ }
+ #[dispatch(walk)]
+ impl Probe {
+ fn walk_any(&mut self, value: &VisitValue, kind: DefRegionKind) ->
WalkResult {
+ assert!(
+ value.cast::<Array<i64>>().is_some(),
+ "children must not be visited"
+ );
+ assert_eq!(kind, DefRegionKind::None);
+ self.events.push("callback");
+ WalkResult::Advance
+ }
+ }
+ struct Scope;
+ impl VisitPolicy<Probe> for Scope {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ ctx: &mut VisitContext<'_, Probe>,
+ ) -> Result<Option<VisitInterrupt>> {
+ ctx.state_mut().events.push("enter");
+ ctx.state_mut().depth += 1;
+ let result = ctx.default_visit_children(&value.to_owned(),
DefRegionKind::Pattern);
+ assert_eq!(ctx.def_region_kind(), DefRegionKind::None);
+ ctx.state_mut().depth -= 1;
+ ctx.state_mut().events.push("exit");
+ result
+ }
+ }
+ struct Stop(bool);
+ impl VisitPolicy<Probe> for Stop {
+ fn default_visit(
+ &self,
+ _: &VisitValue,
+ ctx: &mut VisitContext<'_, Probe>,
+ ) -> Result<Option<VisitInterrupt>> {
+ assert_eq!(ctx.def_region_kind(), DefRegionKind::Pattern);
+ assert_eq!(ctx.state().depth, 1);
+ ctx.state_mut().events.push("stop");
+ if self.0 {
+ Err(runtime_error("policy failed"))
+ } else {
+ Ok(Some(VisitInterrupt::with(42_i64)))
+ }
+ }
+ }
+ struct Unreachable;
+ impl VisitPolicy<Probe> for Unreachable {
+ fn default_visit(
+ &self,
+ _: &VisitValue,
+ _: &mut VisitContext<'_, Probe>,
+ ) -> Result<Option<VisitInterrupt>> {
+ panic!("the policy after Stop must not run")
+ }
+ }
+ let root = Array::new(vec![1_i64]);
+ for error in [false, true] {
+ for order in [None, Some(WalkOrder::PreOrder),
Some(WalkOrder::PostOrder)] {
+ let policies = (Scope, (Stop(error), Unreachable));
+ let (result, state) = if let Some(order) = order {
+ let mut walker = WalkWithPolicy::new(Probe::default(),
policies);
+ let result = walker.walk(&root, order);
+ (result, walker.into_state())
+ } else {
+ let mut visitor = VisitCallbacks::new(
+ Probe::default(),
+ |value: &VisitValue, ctx: &mut VisitContext<'_, Probe>| {
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().walk_any(value, kind);
+ ctx.visit_children()
+ },
+ )
+ .with_policy(policies);
+ let result = structural_visit(&root, &mut visitor);
+ (result, visitor.into_state())
+ };
+ if error {
+ assert!(result.err().unwrap().to_string().contains("policy
failed"));
+ } else {
+
assert_eq!(i64::try_from(result.unwrap().unwrap().value).unwrap(), 42);
+ }
+ let expected = if order == Some(WalkOrder::PostOrder) {
+ vec!["enter", "stop", "exit"]
+ } else {
+ vec!["callback", "enter", "stop", "exit"]
+ };
+ assert_eq!(state.events, expected);
+ assert_eq!(state.depth, 0);
+ }
+ }
+}
+
+#[test]
+fn policy_regions_compose_with_field_flags_and_function_hooks() {
+ struct SetRootRegion(i32, DefRegionKind);
+ impl VisitPolicy<PolicyRegionTrace> for SetRootRegion {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ ctx: &mut VisitContext<'_, PolicyRegionTrace>,
+ ) -> Result<Option<VisitInterrupt>> {
+ if value.type_index() == self.0 {
+ ctx.default_visit_children(&value.to_owned(), self.1)
+ } else {
+ ctx.visit_children()
+ }
+ }
+ }
+ use DefRegionKind::{None as Use, Pattern, Simple};
+ for with_hook in [false, true] {
+ let root = visit_region_graph(with_hook);
+ if with_hook {
+ // Exercise the Function-valued __s_visit__ path, not an opaque
pointer hook.
+ assert!(
+ Function::try_from(get_type_attr(root.type_index(),
"__s_visit__").unwrap())
+ .is_ok()
+ );
+ }
+ for region in [Use, Simple, Pattern] {
+ let ordinary = if with_hook && region == Use {
+ Simple
+ } else {
+ region
+ };
+ let mut expected = vec![
+ (1, if region == Pattern { Pattern } else { Simple }),
+ (2, Pattern),
+ (3, ordinary),
+ ];
+ if with_hook {
+ expected.push((4, region));
+ }
+ for order in [None, Some(WalkOrder::PreOrder),
Some(WalkOrder::PostOrder)] {
+ let policy = SetRootRegion(root.type_index(), region);
+ let state = if let Some(order) = order {
+ let mut walker =
WalkWithPolicy::new(PolicyRegionTrace::default(), policy);
+ assert!(walker.walk(&root, order).unwrap().is_none());
+ walker.into_state()
+ } else {
+ let mut visitor = VisitCallbacks::new(
+ PolicyRegionTrace::default(),
+ |value: i64, ctx: &mut VisitContext<'_,
PolicyRegionTrace>| {
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push((value, kind));
+ },
+ )
+ .with_policy(policy);
+ assert!(structural_visit(&root, &mut
visitor).unwrap().is_none());
+ visitor.into_state()
+ };
+ assert_eq!(state.0, expected);
+ }
+ }
+ }
+}
+
+#[test]
+fn walk_policy_preserves_reflected_pattern_before_default_descent() {
+ let root = visit_region_graph(false);
+ struct Reenter(DefRegionKind);
+ impl VisitPolicy<PolicyRegionTrace> for Reenter {
+ fn default_visit(
+ &self,
+ value: &VisitValue,
+ ctx: &mut VisitContext<'_, PolicyRegionTrace>,
+ ) -> Result<Option<VisitInterrupt>> {
+ match value.cast::<i64>() {
+ Some(2) => {
+ assert_eq!(ctx.def_region_kind(), DefRegionKind::Pattern);
+ // Re-dispatch immediately: visit_children() must not be
needed
+ // to synchronize the ABI visitor with the reflected field
region.
+ ctx.visit_with(&99_i64, self.0)
+ }
+ Some(3) => {
+ assert_eq!(ctx.def_region_kind(), DefRegionKind::None);
+ // The preceding field's Pattern scope must not leak into
a sibling.
+ ctx.visit(&100_i64)
+ }
+ _ => ctx.visit_children(),
+ }
+ }
+ }
+ use DefRegionKind::{None as Use, Pattern, Simple};
+ for requested in [Use, Simple] {
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut walker = WalkWithPolicy::new(PolicyRegionTrace::default(),
Reenter(requested));
+ assert!(walker.walk(&root, order).unwrap().is_none());
+ let expected = match order {
+ WalkOrder::PreOrder => vec![
+ (1, Simple),
+ (2, Pattern),
+ (99, Pattern),
+ (3, Use),
+ (100, Use),
+ ],
+ WalkOrder::PostOrder => vec![
+ (1, Simple),
+ (99, Pattern),
+ (2, Pattern),
+ (100, Use),
+ (3, Use),
+ ],
+ };
+ assert_eq!(
+ walker.state().0,
+ expected,
+ "order={order:?}, requested={requested:?}"
+ );
+ }
+ }
+}
+
#[test]
fn public_reflection_access_uses_registered_field_and_type_attr() {
// Keep the existing C++ test library linked so its startup registrations
@@ -328,10 +923,7 @@ fn manual_child_visit_can_override_def_region() {
let root = Array::new(vec![7i64, 8]);
let mut probe = ManualRegionVisitor::default();
assert!(structural_visit(&root, &mut probe).unwrap().is_none());
- assert_eq!(
- probe.seen,
- vec![DefRegionKind::Simple, DefRegionKind::None]
- );
+ assert_eq!(probe.seen, vec![DefRegionKind::Simple, DefRegionKind::None]);
}
#[derive(Default)]
diff --git a/src/ffi/testing/testing.cc b/src/ffi/testing/testing.cc
index 742ca9d0..cda1078e 100644
--- a/src/ffi/testing/testing.cc
+++ b/src/ffi/testing/testing.cc
@@ -30,6 +30,7 @@
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/enum.h>
#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/accessor.h>
@@ -146,6 +147,49 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::ObjectDef<TestCxxStrEnumObj>(refl::init(false)).def_convert<TestCxxStrEnum>();
}
+// Structural traversal fixtures shared by language binding tests.
+class TestVisitRegionsObj : public Object {
+ public:
+ int64_t simple = 1;
+ int64_t pattern = 2;
+ int64_t ordinary = 3;
+
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestVisitRegions",
TestVisitRegionsObj, Object);
+};
+
+class TestVisitHookObj : public Object {
+ public:
+ ObjectRef child;
+
+ explicit TestVisitHookObj(ObjectRef child) : child(std::move(child)) {}
+
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.TestVisitHook", TestVisitHookObj,
Object);
+};
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<TestVisitRegionsObj>()
+ .def_ro("simple", &TestVisitRegionsObj::simple,
refl::AttachFieldFlag::SEqHashDefSimple())
+ .def_ro("pattern", &TestVisitRegionsObj::pattern,
refl::AttachFieldFlag::SEqHashDefPattern())
+ .def_ro("ordinary", &TestVisitRegionsObj::ordinary);
+ refl::ObjectDef<TestVisitHookObj>().def_ro("child",
&TestVisitHookObj::child);
+ refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralVisit);
+ refl::TypeAttrDef<TestVisitHookObj>().def(
+ refl::type_attr::kStructuralVisit,
+ [](StructuralVisitor visitor, const TestVisitHookObj* self) {
+ auto result = visitor->WithDefRegionKind(kTVMFFIDefRegionKindSimple,
+ [&] { return
visitor->Visit(self->child); });
+ if (result.has_value()) return result;
+ // This sibling observes the region restored after the scoped child
visit.
+ return visitor->Visit(int64_t{4});
+ });
+ refl::GlobalDef().def("testing.make_visit_region_graph", [](bool with_hook)
-> Any {
+ ObjectRef fields(make_object<TestVisitRegionsObj>());
+ if (with_hook) return ObjectRef(make_object<TestVisitHookObj>(fields));
+ return fields;
+ });
+}
+
class TestObjectBase : public Object {
public:
int64_t v_i64;
