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 aef2d001 [FEAT][Rust] Support context policies in visit and mutate
dispatch (#803)
aef2d001 is described below
commit aef2d001d1fa53f647e56f11953203b75abe280a
Author: Shushi Hong <[email protected]>
AuthorDate: Sun Sep 20 14:37:19 2026 -0400
[FEAT][Rust] Support context policies in visit and mutate dispatch (#803)
This PR adds `policy = expression` support to `#[dispatch(visit)]` and
`#[dispatch(mutate)]`. Policies share the dispatch object's state and
run when no handler matches or a handler requests default recursion.
Children still go through normal callback dispatch.
It also preserves an outer `Pattern` during recursive visits and avoids
duplicate error frames without stopping at existing null entries.
Mutation policies retain the input's in-place permission and
`UnchangedOr` result.
---
docs/guides/rust_lang_guide.md | 12 +-
rust/tvm-ffi-macros/src/dispatch.rs | 100 +++++++++---
rust/tvm-ffi-macros/src/lib.rs | 5 +
rust/tvm-ffi/src/extra/structural_common.rs | 12 +-
rust/tvm-ffi/src/extra/structural_mutate.rs | 32 +++-
rust/tvm-ffi/src/extra/structural_visit.rs | 32 +++-
rust/tvm-ffi/src/extra/structural_visit/policy.rs | 41 +++--
rust/tvm-ffi/tests/test_structural_mutate.rs | 187 ++++++++++++++++++++--
rust/tvm-ffi/tests/test_structural_visit.rs | 112 ++++++++++---
9 files changed, 448 insertions(+), 85 deletions(-)
diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md
index a4b9709d..5d01101e 100644
--- a/docs/guides/rust_lang_guide.md
+++ b/docs/guides/rust_lang_guide.md
@@ -411,6 +411,9 @@ assert_eq!(depth.max, 2);
Implement `StructuralVisitor` directly to override its low-level `visit`
method.
+Use `#[dispatch(visit, policy = MyPolicy)]` to apply a `ContextPolicy<Self>`
+to default recursion.
+
Two safety notes: mutable `List`/`Dict` contents are snapshotted before
callbacks run, so mutation during traversal cannot invalidate the walk; and
a non-container type with a foreign `__s_visit__` hook is rejected rather
@@ -558,8 +561,8 @@ Closure callbacks are `Fn`; mutable data belongs in the
callback state.
Use `MutateValue<'_, T>` with `default_maybe_inplace_mutate` to forward its
permission, or `default_mutate_with_mode` to restrict it. Borrow through
`value`; mutation
contexts do not expose `current()`, and copy-only default descent takes an
-explicit borrow (`default_mutate(value)`). Generated handlers may take
-`InplaceMode` after `&mut Mutator`; see `MutateValue` for ownership
requirements.
+explicit borrow (`default_mutate(value)`). Generated handlers can read the mode
+through `Mutator::inplace_mode()`; see `MutateValue` for ownership
requirements.
`#[dispatch(mutate)]` groups typed `mutate_*` callbacks. The dispatch object
owns its mutable pass state, while `Mutator` supplies recursion and the current
@@ -595,7 +598,10 @@ assert_eq!(mutated.iter().collect::<Vec<_>>(), vec![2, 3]);
assert_eq!(increment.integers, 2);
```
-For a named custom recursion policy, implement `StructuralMutator` and pass
+Use `#[dispatch(mutate, policy = MyPolicy)]` to apply a
`MutContextPolicy<Self>`
+to default recursion.
+
+For low-level custom recursion, implement `StructuralMutator` and pass
`&mut` it to `structural_mutate`. `InplaceValue` is an engine-issued
capability: callers cannot construct it from a read-only `StructuralView`.
Override
`dispatch_maybe_inplace_mutate` to opt into default container reuse;
diff --git a/rust/tvm-ffi-macros/src/dispatch.rs
b/rust/tvm-ffi-macros/src/dispatch.rs
index b4528d10..94aeda86 100644
--- a/rust/tvm-ffi-macros/src/dispatch.rs
+++ b/rust/tvm-ffi-macros/src/dispatch.rs
@@ -21,8 +21,8 @@ use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{quote, quote_spanned};
use syn::{
- parse_macro_input, FnArg, ImplItem, ImplItemMethod, ItemImpl, Meta,
NestedMeta, PathArguments,
- Type,
+ parse_macro_input, Expr, FnArg, ImplItem, ImplItemMethod, ItemImpl, Meta,
NestedMeta,
+ PathArguments, Token, Type,
};
use crate::utils::get_tvm_ffi_crate;
@@ -31,7 +31,7 @@ pub(crate) fn dispatch(attr: TokenStream, item: TokenStream)
-> TokenStream {
let args = parse_macro_input!(attr as DispatchArgs);
let item_impl = parse_macro_input!(item as ItemImpl);
- match expand(&item_impl, args.mode) {
+ match expand(&item_impl, args) {
Ok(generated) => quote!(#item_impl #generated).into(),
Err(error) => {
let error = error.to_compile_error();
@@ -42,6 +42,7 @@ pub(crate) fn dispatch(attr: TokenStream, item: TokenStream)
-> TokenStream {
struct DispatchArgs {
mode: DispatchMode,
+ policy: Option<Expr>,
}
#[derive(Clone, Copy)]
@@ -97,25 +98,32 @@ impl syn::parse::Parse for DispatchArgs {
`dispatch(mutate)`",
));
};
- if !input.is_empty() {
- let message = if matches!(mode, DispatchMode::Mutate) {
- "`dispatch(mutate)` takes no further arguments; the definition
region is \
- available through `Mutator::region()`"
- .to_owned()
- } else {
- format!(
- "`dispatch({})` takes no further arguments; a handler that
needs the \
- definition-region state declares a trailing
`DefRegionKind` argument",
- mode.name()
- )
- };
- return Err(input.error(message));
- }
- Ok(DispatchArgs { mode })
+ let policy = if input.is_empty() {
+ None
+ } else {
+ input.parse::<Token![,]>()?;
+ if !matches!(mode, DispatchMode::Visit | DispatchMode::Mutate) {
+ return Err(input.error(
+ "`policy` is supported by `dispatch(visit)` and
`dispatch(mutate)`",
+ ));
+ }
+ let name: syn::Ident = input.parse()?;
+ if name != "policy" {
+ return Err(syn::Error::new(name.span(), "expected `policy =
expression`"));
+ }
+ input.parse::<Token![=]>()?;
+ let policy = input.parse()?;
+ if !input.is_empty() {
+ input.parse::<Token![,]>()?;
+ }
+ Some(policy)
+ };
+ Ok(DispatchArgs { mode, policy })
}
}
-fn expand(item_impl: &ItemImpl, mode: DispatchMode) ->
syn::Result<TokenStream2> {
+fn expand(item_impl: &ItemImpl, args: DispatchArgs) ->
syn::Result<TokenStream2> {
+ let DispatchArgs { mode, policy } = args;
if item_impl.trait_.is_some() {
return Err(syn::Error::new_spanned(
item_impl,
@@ -163,6 +171,55 @@ fn expand(item_impl: &ItemImpl, mode: DispatchMode) ->
syn::Result<TokenStream2>
let self_type = &item_impl.self_ty;
let (impl_generics, _, where_clause) = item_impl.generics.split_for_impl();
let impl_cfg_attrs = presence_attrs(&item_impl.attrs)?;
+ let (policy_state_impl, policy_method) = if let Some(policy) = policy {
+ let (module, state_trait, method) = match mode {
+ DispatchMode::Visit => (
+ quote!(#tvm_ffi::extra::structural_visit),
+ quote!(VisitCallbackState),
+ quote! {
+ fn default_visit_children(
+ &mut self,
+ value: &#tvm_ffi::StructuralView,
+ kind: #tvm_ffi::DefRegionKind,
+ ) -> #tvm_ffi::Result<Option<#tvm_ffi::VisitInterrupt>> {
+ let policy = #policy;
+
#tvm_ffi::extra::structural_visit::default_visit_with_policy(
+ self, &policy, value, kind,
+ )
+ }
+ },
+ ),
+ DispatchMode::Mutate => (
+ quote!(#tvm_ffi::extra::structural_mutate),
+ quote!(MutateCallbackState),
+ quote! {
+ fn on_default_mutate(
+ &mut self,
+ value: #tvm_ffi::MutateValue<'_>,
+ kind: #tvm_ffi::DefRegionKind,
+ ) -> #tvm_ffi::Result<#tvm_ffi::Any> {
+ let policy = #policy;
+
#tvm_ffi::extra::structural_mutate::default_mutate_with_policy(
+ self, &policy, value, kind,
+ )
+ }
+ },
+ ),
+ _ => unreachable!(),
+ };
+ (
+ quote! {
+ #(#[#impl_cfg_attrs])*
+ impl #impl_generics #module::#state_trait<Self> for #self_type
#where_clause {
+ fn callback_state(&self) -> &Self { self }
+ fn callback_state_mut(&mut self) -> &mut Self { self }
+ }
+ },
+ method,
+ )
+ } else {
+ (quote!(), quote!())
+ };
let ordering_errors = handlers
.iter()
.enumerate()
@@ -205,6 +262,8 @@ fn expand(item_impl: &ItemImpl, mode: DispatchMode) ->
syn::Result<TokenStream2>
impl #impl_generics
#tvm_ffi::extra::structural_visit::StructuralVisitor
for #self_type #where_clause
{
+ #policy_method
+
#[inline]
#[allow(unreachable_code, unused_variables)]
fn visit(
@@ -239,6 +298,8 @@ fn expand(item_impl: &ItemImpl, mode: DispatchMode) ->
syn::Result<TokenStream2>
impl #impl_generics
#tvm_ffi::extra::structural_mutate::MutateDispatch
for #self_type #where_clause
{
+ #policy_method
+
#[inline(always)]
#[allow(unreachable_code, unused_variables)]
fn dispatch_mutate(
@@ -267,6 +328,7 @@ fn expand(item_impl: &ItemImpl, mode: DispatchMode) ->
syn::Result<TokenStream2>
Ok(quote! {
#(#ordering_errors)*
+ #policy_state_impl
#(#[#impl_cfg_attrs])*
#dispatch_impl
diff --git a/rust/tvm-ffi-macros/src/lib.rs b/rust/tvm-ffi-macros/src/lib.rs
index 655ef8fc..7cb3805d 100644
--- a/rust/tvm-ffi-macros/src/lib.rs
+++ b/rust/tvm-ffi-macros/src/lib.rs
@@ -26,6 +26,11 @@ mod object_macros;
mod utils;
/// Generate `walk`, `map`, `visit`, or `mutate` dispatch from an inherent
impl.
+///
+/// `#[dispatch(visit, policy = expr)]` and `#[dispatch(mutate, policy =
expr)]`
+/// use a `ContextPolicy<Self>` or `MutContextPolicy<Self>` for default
recursion.
+/// The expression runs at each default descent and may read `self` for
configuration;
+/// mutable traversal state stays on `self`. Policy tuples compose as usual.
#[proc_macro_error]
#[proc_macro_attribute]
pub fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream {
diff --git a/rust/tvm-ffi/src/extra/structural_common.rs
b/rust/tvm-ffi/src/extra/structural_common.rs
index 5fb9bfd9..604841c2 100644
--- a/rust/tvm-ffi/src/extra/structural_common.rs
+++ b/rust/tvm-ffi/src/extra/structural_common.rs
@@ -45,6 +45,7 @@ pub(crate) fn with_visit_error_context(error: Error, raw:
TVMFFIAny) -> Error {
))
.ok()?;
}
+ let node = StructuralView::from_raw(raw).cast::<object::ObjectRef>()?;
let mut previous = error.extra_context();
let mut nodes = Vec::new();
if let Some(prior) = previous.as_ref() {
@@ -60,6 +61,16 @@ pub(crate) fn with_visit_error_context(error: Error, raw:
TVMFFIAny) -> Error {
.ok()?
.try_as::<i64>()?;
let get_item = Function::get_global("ffi.ListGetItem").ok()?;
+ if size > 0 {
+ let last = get_item
+ .call_tuple((records.clone(), size - 1))
+ .ok()?
+ .try_as::<object::ObjectRef>();
+ // Callback, policy and default descent can report the
same frame.
+ if last.is_some_and(|last| last.same_as(&node)) {
+ return None;
+ }
+ }
for i in 0..size {
nodes.push(get_item.call_tuple((records.clone(),
i)).ok()?);
}
@@ -69,7 +80,6 @@ pub(crate) fn with_visit_error_context(error: Error, raw:
TVMFFIAny) -> Error {
.ok()?;
}
}
- let node = StructuralView::from_raw(raw).cast::<object::ObjectRef>()?;
nodes.push(Any::from(node));
let records = Function::get_global("ffi.List")
.ok()?
diff --git a/rust/tvm-ffi/src/extra/structural_mutate.rs
b/rust/tvm-ffi/src/extra/structural_mutate.rs
index 366d4924..ed972af4 100644
--- a/rust/tvm-ffi/src/extra/structural_mutate.rs
+++ b/rust/tvm-ffi/src/extra/structural_mutate.rs
@@ -766,6 +766,16 @@ pub trait MutateDispatch: Sized {
) -> Option<Result<Any>> {
self.dispatch_mutate(value.as_value(), mutator)
}
+
+ #[doc(hidden)]
+ fn on_default_mutate(&mut self, value: MutateValue<'_>, kind:
DefRegionKind) -> Result<Any> {
+ default_mutate_driver(
+ self,
+ value.value.raw(),
+ kind,
+ value.permit(value.inplace_mode()),
+ )
+ }
}
impl<D: MutateDispatch> IntoMutator<ByMutateDispatch> for D {
@@ -1022,7 +1032,8 @@ struct DirectMutateCallbacks<'a, Link, Marker> {
_marker: PhantomData<fn(Marker)>,
}
-trait MutateCallbackState<State> {
+#[doc(hidden)]
+pub trait MutateCallbackState<State> {
fn callback_state(&self) -> &State;
fn callback_state_mut(&mut self) -> &mut State;
}
@@ -1724,6 +1735,10 @@ pub trait StructuralMutator: Sized {
}
impl<D: MutateDispatch> StructuralMutator for D {
+ fn on_default_mutate(&mut self, value: MutateValue<'_>, kind:
DefRegionKind) -> Result<Any> {
+ MutateDispatch::on_default_mutate(self, value, kind)
+ }
+
#[inline(always)]
fn dispatch_mutate(
&mut self,
@@ -1775,6 +1790,21 @@ impl<D: MutateDispatch> StructuralMutator for D {
}
}
+#[doc(hidden)]
+pub fn default_mutate_with_policy<D: MutateDispatch + MutateCallbackState<D>>(
+ dispatch: &mut D,
+ policy: &impl MutContextPolicy<D>,
+ value: MutateValue<'_>,
+ kind: DefRegionKind,
+) -> Result<Any> {
+ policy::mutate_with_policy(
+ &mut policy::MutationDescent { driver: dispatch },
+ policy,
+ value,
+ kind,
+ )
+}
+
// Closure callback chains use a type-erased context driver so one concrete
// function signature can recurse through the complete chain.
#[inline(always)]
diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs
b/rust/tvm-ffi/src/extra/structural_visit.rs
index 15109a32..06ab7896 100644
--- a/rust/tvm-ffi/src/extra/structural_visit.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit.rs
@@ -564,7 +564,8 @@ struct DirectVisitCallbacks<'a, Link, Marker> {
_marker: PhantomData<fn(Marker)>,
}
-trait VisitCallbackState<State> {
+#[doc(hidden)]
+pub trait VisitCallbackState<State> {
fn callback_state(&self) -> &State;
fn callback_state_mut(&mut self) -> &mut State;
}
@@ -1014,6 +1015,16 @@ fn default_user_visit_children<V: StructuralVisitor>(
)
}
+#[doc(hidden)]
+pub fn default_visit_with_policy<V: StructuralVisitor + VisitCallbackState<V>>(
+ visitor: &mut V,
+ policy: &impl ContextPolicy<V>,
+ value: &StructuralView,
+ kind: DefRegionKind,
+) -> Result<Option<VisitInterrupt>> {
+ policy::visit_with_policy(&mut policy::VisitDescent { visitor }, policy,
value, kind)
+}
+
fn try_visit_callbacks<State, Link, Marker>(
driver: &mut impl VisitContextDriver<State>,
callback_ptr: *const Link,
@@ -1219,10 +1230,9 @@ impl<V: StructuralVisitor> ChildVisit for
UserChildren<'_, V> {
if child.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
return Ok(());
}
- match self
- .visitor
- .visit(&StructuralView::from_raw(child), def_region_kind)
- {
+ match with_visit_region(def_region_kind, |kind| {
+ self.visitor.visit(&StructuralView::from_raw(child), kind)
+ }) {
Ok(None) => Ok(()),
Ok(Some(interrupt)) => Err(NativeHalt::Interrupt(interrupt.value)),
Err(error) => Err(with_value_context(NativeHalt::Error(error),
child)),
@@ -1880,6 +1890,18 @@ fn visit_result_from_any(value: Any) -> NativeResult {
}
}
+fn with_visit_region<T>(
+ kind: DefRegionKind,
+ callback: impl FnOnce(DefRegionKind) -> Result<T>,
+) -> Result<T> {
+ let active = active_structural_visitor()?;
+ with_visitor_def_region(active, kind, || {
+ // SAFETY: the active invocation keeps this thread's ABI visitor alive.
+ let kind = def_region_from_raw(unsafe { (*active).def_region_mode })?;
+ callback(kind)
+ })
+}
+
fn with_visitor_def_region<T>(
visitor: StructuralVisitorHandle,
kind: DefRegionKind,
diff --git a/rust/tvm-ffi/src/extra/structural_visit/policy.rs
b/rust/tvm-ffi/src/extra/structural_visit/policy.rs
index 8b8e3d4c..7ffcd2aa 100644
--- a/rust/tvm-ffi/src/extra/structural_visit/policy.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit/policy.rs
@@ -26,8 +26,8 @@ use super::*;
/// `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. Save and restore scoped state around
descent.
+/// tuples may nest. Policies receive `&self`; mutable pass data belongs in the
+/// context's state. Save and restore scoped state around descent.
pub trait ContextPolicy<State> {
/// Customize default descent for the current value.
///
@@ -82,15 +82,17 @@ pub(super) fn visit_with_policy<State>(
value: &StructuralView,
def_region_kind: DefRegionKind,
) -> Result<Option<VisitInterrupt>> {
- policy.default_visit(
- value,
- &mut VisitContext {
- driver,
- current: StructuralView::from_raw(value.raw()),
- def_region_kind,
- _not_send_sync: PhantomData,
- },
- )
+ with_visit_region(def_region_kind, |def_region_kind| {
+ policy.default_visit(
+ value,
+ &mut VisitContext {
+ driver,
+ current: StructuralView::from_raw(value.raw()),
+ def_region_kind,
+ _not_send_sync: PhantomData,
+ },
+ )
+ })
}
struct NextPolicy<'a, State, Policy> {
@@ -223,17 +225,12 @@ impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
NativeVisit
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,
- )
- })
+ visit_with_policy(
+ &mut WalkDescent::<_, _, PRE_ORDER> { visitor: self },
+ &*policy,
+ value,
+ kind,
+ )
}
}
diff --git a/rust/tvm-ffi/tests/test_structural_mutate.rs
b/rust/tvm-ffi/tests/test_structural_mutate.rs
index 15d36c17..bd22e593 100644
--- a/rust/tvm-ffi/tests/test_structural_mutate.rs
+++ b/rust/tvm-ffi/tests/test_structural_mutate.rs
@@ -23,11 +23,12 @@ use tvm_ffi::function::FunctionObj;
use tvm_ffi::object::ObjectRef;
use tvm_ffi::{
dispatch, structural_map, structural_mutate, structural_visit,
structural_walk, Any, AnyView,
- Array, DefRegionKind, DefaultMutContextPolicy, Error, FieldGetter,
Function, InplaceMode,
- InplaceValue, IntoMapper, Map, MapDispatch, MapWithContextPolicy,
MutContextPolicy,
- MutateCallbacks, MutateContext, MutateValue, Mutator, Object, ObjectArc,
ObjectRefCore, Result,
- String as FfiString, StructuralMutator, StructuralVarRemap,
StructuralView, TypeIndex,
- Unchanged, UnchangedOr, VisitContext, WalkOrder, WalkResult, RUNTIME_ERROR,
+ Array, DefRegionKind, DefaultContextPolicy, DefaultMutContextPolicy,
Error, FieldGetter,
+ Function, InplaceMode, InplaceValue, IntoMapper, Map, MapDispatch,
MapWithContextPolicy,
+ MutContextPolicy, MutateCallbacks, MutateContext, MutateValue, Mutator,
Object, ObjectArc,
+ ObjectRefCore, Result, String as FfiString, StructuralMutator,
StructuralVarRemap,
+ StructuralView, StructuralVisitor, TypeIndex, Unchanged, UnchangedOr,
VisitCallbacks,
+ VisitContext, VisitInterrupt, WalkOrder, WalkResult, RUNTIME_ERROR,
};
struct IncrementIntegers;
@@ -743,6 +744,29 @@ fn
reflected_object_without_shallow_copy_is_rejected_even_when_unchanged() {
#[test]
fn callback_errors_preserve_message_and_add_object_context() {
+ struct Delegate(Error);
+ #[dispatch(visit, policy = (DefaultContextPolicy, DefaultContextPolicy))]
+ impl Delegate {
+ fn visit_any(
+ &mut self,
+ value: &StructuralView,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ if value.cast::<i64>().is_some() {
+ return Err(self.0.clone());
+ }
+ self.default_visit_children(value, kind)
+ }
+ }
+ #[dispatch(mutate, policy = (DefaultMutContextPolicy,
DefaultMutContextPolicy))]
+ impl Delegate {
+ fn mutate_any(&mut self, value: MutateValue<'_>, ctx: &mut Mutator) ->
Result<Any> {
+ if value.cast::<i64>().is_some() {
+ return Err(self.0.clone());
+ }
+ ctx.default_maybe_inplace_mutate(self, value)
+ }
+ }
let child = reflected_object();
let root = call_global(
"ffi.MakeObjectFromPackedArgs",
@@ -753,7 +777,7 @@ fn
callback_errors_preserve_message_and_add_object_context() {
],
);
let payload =
ObjectRef::try_from(Any::from(Array::new(vec![7i64]))).unwrap();
- let records = call_global("ffi.List", &[child.clone()]);
+ let records = call_global("ffi.List", &[Any::new()]);
let context = call_global(
"ffi.MakeObjectFromPackedArgs",
&[
@@ -800,6 +824,15 @@ fn
callback_errors_preserve_message_and_add_object_context() {
.err()
.unwrap(),
);
+ let mut mapper = MapWithContextPolicy::new(
+ (|_: i64| -> Result<i64> { Err(source.clone()) }).into_mapper(),
+ (DefaultMutContextPolicy, DefaultMutContextPolicy),
+ );
+ errors.push(
+ structural_map(root.clone(), &mut mapper, order)
+ .err()
+ .unwrap(),
+ );
}
errors.push(
structural_mutate(
@@ -817,6 +850,34 @@ fn
callback_errors_preserve_message_and_add_object_context() {
.err()
.unwrap(),
);
+ let mut visitor = VisitCallbacks::new(
+ (),
+ |value: &StructuralView, ctx: &mut VisitContext<'_, ()>| {
+ if value.cast::<i64>().is_some() {
+ return Err(source.clone());
+ }
+ ctx.visit_children()
+ },
+ )
+ .with_policy((DefaultContextPolicy, DefaultContextPolicy));
+ errors.push(structural_visit(&root, &mut visitor).err().unwrap());
+ let mut mutator =
+ MutateCallbacks::new((), |value: MutateValue<'_>, ctx: &mut
MutateContext| {
+ if value.cast::<i64>().is_some() {
+ return Err(source.clone());
+ }
+ ctx.default_maybe_inplace_mutate(value)
+ })
+ .with_policy((DefaultMutContextPolicy, DefaultMutContextPolicy));
+ errors.push(structural_mutate(root.clone(), &mut mutator).err().unwrap());
+ let mut delegate = Delegate(source.clone());
+ errors.push(structural_visit(&root, &mut delegate).err().unwrap());
+ errors.push(
+ structural_mutate(root.clone(), &mut delegate)
+ .err()
+ .unwrap(),
+ );
+ let seeded = errors[0].clone();
for error in errors {
assert_eq!(error.message(), "callback failed");
assert!(error.backtrace().contains("origin"));
@@ -831,6 +892,13 @@ fn
callback_errors_preserve_message_and_add_object_context() {
"reverse_visit_pattern",
));
let size = i64::try_from(call_global("ffi.ListSize",
&[records.clone()])).unwrap();
+ assert_eq!(size, 3);
+ assert_eq!(
+ call_global("ffi.ListGetItem", &[records.clone(),
0_i64.into()]).try_as::<()>(),
+ Some(())
+ );
+ let innermost = call_global("ffi.ListGetItem", &[records.clone(),
1_i64.into()]);
+ assert_eq!(any_object_pointer(&innermost), any_object_pointer(&child));
let outermost = call_global("ffi.ListGetItem", &[records, (size -
1).into()]);
assert_eq!(any_object_pointer(&outermost), any_object_pointer(&root));
let paths = call_global(
@@ -857,6 +925,28 @@ fn
callback_errors_preserve_message_and_add_object_context() {
);
assert_eq!(source.backtrace(), "origin");
+ // Keep nonconsecutive occurrences: child -> root -> child is a distinct
path.
+ let error = structural_map(
+ child.clone(),
+ |_: i64| -> Result<i64> { Err(seeded.clone()) },
+ WalkOrder::PreOrder,
+ )
+ .err()
+ .unwrap();
+ let context = Any::from(error.extra_context().unwrap());
+ let records = Any::from(reflected_field::<ObjectRef>(
+ &context,
+ "reverse_visit_pattern",
+ ));
+ assert_eq!(
+ i64::try_from(call_global("ffi.ListSize",
&[records.clone()])).unwrap(),
+ 4
+ );
+ for (i, expected) in [&child, &root, &child].into_iter().enumerate() {
+ let node = call_global("ffi.ListGetItem", &[records.clone(), (i as i64
+ 1).into()]);
+ assert_eq!(any_object_pointer(&node), any_object_pointer(expected));
+ }
+
// The failing node may be a pre-order replacement or a post-order rebuilt
parent.
for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
let root = Array::new(vec![0i64]);
@@ -1944,6 +2034,13 @@ impl MutContextPolicy<PolicyState> for RecordPolicy {
}
}
+#[dispatch(mutate, policy = (ArrayPolicy, RecordPolicy))]
+impl PolicyState {
+ fn mutate_integer(&mut self, value: i64) -> i64 {
+ self.map_integer(value)
+ }
+}
+
#[test]
fn mutation_policies_share_state_and_preserve_callback_order() {
let root = || Array::new(vec![Any::from(Array::new(vec![1_i64])),
Any::from(2_i64)]);
@@ -2012,6 +2109,21 @@ fn
mutation_policies_share_state_and_preserve_callback_order() {
assert_eq!(i64::try_from(array_item(&output, 1)).unwrap(), 3);
assert_eq!(mutator.state().events, pre);
assert_eq!(mutator.state().depth, 0);
+
+ let mut dispatch = PolicyState::default();
+ let output = structural_mutate(root(), &mut dispatch).unwrap();
+ assert_eq!(i64::try_from(array_item(&output, 1)).unwrap(), 3);
+ assert_eq!(
+ i64::try_from(array_item(&array_item(&output, 0), 0)).unwrap(),
+ 2
+ );
+ assert_eq!(dispatch.depth, 0);
+ assert_eq!(
+ dispatch.events,
+ pre.into_iter()
+ .filter(|(tag, _)| *tag != "callback")
+ .collect::<Vec<_>>()
+ );
}
#[test]
@@ -2061,6 +2173,8 @@ fn
mutation_policy_continuations_preserve_ownership_and_markers() {
struct Ownership {
increment: bool,
retained: Option<Any>,
+ mode: InplaceMode,
+ retain: bool,
}
#[dispatch(map)]
impl Ownership {
@@ -2095,8 +2209,24 @@ fn
mutation_policy_continuations_preserve_ownership_and_markers() {
Ok(result)
}
}
- for entry in 0..3 {
- // pre-order map, post-order map, mutation callback
+ #[dispatch(mutate, policy = (
+ Control { mode: self.mode, retain: self.retain },
+ DefaultMutContextPolicy,
+ ))]
+ impl Ownership {
+ fn mutate_integer(&mut self, value: i64) -> UnchangedOr<i64> {
+ self.map_integer(value)
+ }
+ fn mutate_array(
+ &mut self,
+ value: MutateValue<'_, Array<i64>>,
+ mutator: &mut Mutator,
+ ) -> Result<UnchangedOr<Any>> {
+ mutator.default_maybe_inplace_mutate_result(self, value)
+ }
+ }
+ for entry in 0..4 {
+ // pre-order map, post-order map, closure mutation, generated mutation
for case in 0..4 {
// unique, shared, alias retained by policy, forced copy
for increment in [false, true] {
@@ -2117,6 +2247,8 @@ fn
mutation_policy_continuations_preserve_ownership_and_markers() {
let state = Ownership {
increment,
retained: None,
+ mode: policy.0.mode,
+ retain: policy.0.retain,
};
let (output, state) = if entry < 2 {
let mut mapper = MapWithContextPolicy::new(state, policy);
@@ -2131,6 +2263,10 @@ fn
mutation_policy_continuations_preserve_ownership_and_markers() {
)
.unwrap();
(output, mapper.into_state())
+ } else if entry == 3 {
+ let mut dispatch = state;
+ let output = structural_mutate(root, &mut
dispatch).unwrap();
+ (output, dispatch)
} else {
let mut mutator = MutateCallbacks::new(
state,
@@ -2165,7 +2301,7 @@ fn
mutation_policy_continuations_preserve_ownership_and_markers() {
fn mutation_policy_regions_retargeting_and_error_restore() {
use DefRegionKind::{None as Use, Pattern, Simple};
#[derive(Default)]
- struct Regions(Vec<(i64, DefRegionKind)>);
+ struct Regions(Vec<(i64, DefRegionKind)>, DefRegionKind);
#[dispatch(map)]
impl Regions {
fn map_integer(&mut self, x: i64, kind: DefRegionKind) -> i64 {
@@ -2222,6 +2358,25 @@ fn
mutation_policy_regions_retargeting_and_error_restore() {
ctx.default_maybe_inplace_mutate_result(value)
}
}
+ #[dispatch(mutate, policy = (Redirect(self.1), Observe))]
+ impl Regions {
+ fn mutate_any(
+ &mut self,
+ value: MutateValue<'_>,
+ mutator: &mut Mutator,
+ ) -> Result<UnchangedOr<Any>> {
+ if let Some(x) = value.cast::<i64>() {
+ self.map_integer(x, mutator.def_region_kind());
+ }
+ if value
+ .as_node::<tvm_ffi::collections::array::ArrayObj>()
+ .is_some()
+ {
+ self.0.push((200, mutator.def_region_kind()));
+ }
+ mutator.default_maybe_inplace_mutate_result(self, value)
+ }
+ }
assert_eq!(
unsafe { tvm_ffi::tvm_ffi_sys::TVMFFITestingDummyTarget() },
0
@@ -2281,6 +2436,20 @@ fn
mutation_policy_regions_retargeting_and_error_restore() {
mutator.state().0,
vec![(1, Simple), (2, Pattern), (99, Pattern), (3, Use)]
);
+ let mut dispatch = Regions(vec![], requested);
+ let output = structural_mutate(false, &mut dispatch).unwrap();
+ assert_eq!(i64::try_from(array_item(&output, 0)).unwrap(), 1);
+ assert_eq!(dispatch.0, vec![(100, Pattern), (1, Pattern)]);
+ dispatch.0.clear();
+ let graph = Function::get_global("testing.make_visit_region_graph")
+ .unwrap()
+ .call_tuple((false,))
+ .unwrap();
+ structural_mutate(graph, &mut dispatch).unwrap();
+ assert_eq!(
+ dispatch.0,
+ vec![(1, Simple), (2, Pattern), (99, Pattern), (3, Use)]
+ );
}
}
diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs
b/rust/tvm-ffi/tests/test_structural_visit.rs
index 4901b305..19896901 100644
--- a/rust/tvm-ffi/tests/test_structural_visit.rs
+++ b/rust/tvm-ffi/tests/test_structural_visit.rs
@@ -109,6 +109,18 @@ fn
composed_policies_share_array_scope_with_visit_and_walk_callbacks() {
// Only unmatched arrays enter the default policies in this visit.
assert_eq!(visitor.state().descent_depths, vec![1, 2]);
+ #[dispatch(visit, policy = (ArrayScope, RecordDescent))]
+ impl CollectIntegers {
+ fn visit_integer(&mut self, value: i64) {
+ self.record(value);
+ }
+ }
+ let mut visitor = CollectIntegers::default();
+ assert!(structural_visit(&root, &mut visitor).unwrap().is_none());
+ assert_eq!(visitor.integers, expected);
+ assert_eq!(visitor.descent_depths, vec![1, 2]);
+ assert_eq!(visitor.depth, 0);
+
#[dispatch(walk)]
impl CollectIntegers {
fn walk_integer(&mut self, value: i64) -> WalkResult {
@@ -194,9 +206,31 @@ fn
policy_continuation_scopes_regions_and_restores_after_halts() {
}
}
use DefRegionKind::{None as Use, Pattern, Simple};
+ #[dispatch(visit, policy = (Scope(Simple, 4), (Scope(Pattern, 3),
Scope(Use, 2))))]
+ impl Probe {
+ fn visit_any(
+ &mut self,
+ value: &StructuralView,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ if let Some(value) = value.cast::<i64>() {
+ self.record(value, kind)
+ } else {
+ self.default_visit_children(value, kind)
+ }
+ }
+ }
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)] {
+ for (entry, order) in [
+ None,
+ None,
+ Some(WalkOrder::PreOrder),
+ Some(WalkOrder::PostOrder),
+ ]
+ .into_iter()
+ .enumerate()
+ {
let state = Probe {
outcome,
seen: vec![],
@@ -212,6 +246,10 @@ fn
policy_continuation_scopes_regions_and_restores_after_halts() {
.unwrap()
.is_none());
(result, walker.into_state())
+ } else if entry == 0 {
+ let mut visitor = state;
+ let result = structural_visit(&root, &mut visitor);
+ (result, visitor)
} else {
let mut visitor =
VisitCallbacks::new(state, |value: i64, ctx: &mut
VisitContext<'_, Probe>| {
@@ -591,37 +629,61 @@ fn
policy_regions_compose_with_field_flags_and_function_hooks() {
}
}
-#[test]
-fn walk_policy_preserves_reflected_pattern_before_default_descent() {
- let root = visit_region_graph(false);
- struct Reenter(DefRegionKind);
- impl ContextPolicy<PolicyRegionTrace> for Reenter {
- fn default_visit(
- &self,
- value: &StructuralView,
- 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(),
+struct ReenterRegion(DefRegionKind);
+impl ContextPolicy<PolicyRegionTrace> for ReenterRegion {
+ fn default_visit(
+ &self,
+ value: &StructuralView,
+ 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(),
+ }
+ }
+}
+
+#[dispatch(visit, policy = ReenterRegion(DefRegionKind::None))]
+impl PolicyRegionTrace {
+ fn visit_any(
+ &mut self,
+ value: &StructuralView,
+ kind: DefRegionKind,
+ ) -> Result<Option<VisitInterrupt>> {
+ if let Some(integer) = value.cast::<i64>() {
+ self.walk_integer(integer, kind);
+ if integer != 2 && integer != 3 {
+ return Ok(None);
}
}
+ self.default_visit_children(value, DefRegionKind::None)
}
+}
+
+#[test]
+fn policies_preserve_reflected_pattern_before_default_descent() {
+ let root = visit_region_graph(false);
use DefRegionKind::{None as Use, Pattern, Simple};
+ let mut visitor = PolicyRegionTrace::default();
+ assert!(structural_visit(&root, &mut visitor).unwrap().is_none());
+ assert_eq!(
+ visitor.0,
+ vec![(1, Simple), (2, Pattern), (99, Pattern), (3, Use), (100, Use)]
+ );
for requested in [Use, Simple] {
for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
let mut walker =
- WalkWithContextPolicy::new(PolicyRegionTrace::default(),
Reenter(requested));
+ WalkWithContextPolicy::new(PolicyRegionTrace::default(),
ReenterRegion(requested));
assert!(structural_walk(&root, &mut walker, order)
.unwrap()
.is_none());
