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 fc0b2faa [FEAT][Rust] Add context policies for structural map and
mutation (#800)
fc0b2faa is described below
commit fc0b2faad1c998f1dcdd2ae03fac5855bb28c50d
Author: Shushi Hong <[email protected]>
AuthorDate: Fri Sep 18 15:28:52 2026 -0400
[FEAT][Rust] Add context policies for structural map and mutation (#800)
This pr adds `MutationPolicy<State>` to customize default recursion with
shared callback state. `MutateCallbacks::with_policy` applies policies
when a callback requests default descent or no callback matches;
`MapWithPolicy` applies them between map callback positions. Policies
compose as `(outer, inner)` tuples and consume `MutateValue` while
preserving `UnchangedOr` results.
Continue through the next policy without redispatching the target, while
children re-enter the full callback engine. Scoped definition-region
changes preserve Pattern and restore the ABI context on exit. Keep
in-place permission and ownership checks across continuation, including
retained aliases and explicit Disallow.
---
docs/guides/rust_lang_guide.md | 9 +-
rust/tvm-ffi/src/extra/structural_mutate.rs | 195 +++++++--
rust/tvm-ffi/src/extra/structural_mutate/policy.rs | 280 +++++++++++++
.../extra/structural_mutate/policy/compile_fail.rs | 50 +++
rust/tvm-ffi/src/extra/structural_visit.rs | 7 +-
rust/tvm-ffi/src/extra/structural_visit/policy.rs | 10 +-
rust/tvm-ffi/src/lib.rs | 13 +-
rust/tvm-ffi/tests/test_structural_mutate.rs | 452 ++++++++++++++++++++-
rust/tvm-ffi/tests/test_structural_visit.rs | 17 +-
9 files changed, 976 insertions(+), 57 deletions(-)
diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md
index 6da519a3..66d64daa 100644
--- a/docs/guides/rust_lang_guide.md
+++ b/docs/guides/rust_lang_guide.md
@@ -366,7 +366,7 @@ 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` use `ContextPolicy<State>`
+`VisitCallbacks::with_policy` and `WalkWithContextPolicy` use
`ContextPolicy<State>`
to manage context around default recursion. See the `ContextPolicy` API
documentation for composition and shared state access.
@@ -488,6 +488,13 @@ Within one `structural_map` call, callbacks run at every
occurrence; their
results are not cached. Default recursion manages identity remapping with
the same semantics as C++.
+`MutateCallbacks::with_policy` and `MapWithContextPolicy` use
`MutContextPolicy<State>`
+to manage context around default recursion. Policies consume `MutateValue`
+and return `UnchangedOr<Any>`; see the API documentation for composition and
+scoped definition regions. Pass `MapWithContextPolicy` directly to
`structural_map`;
+it cannot be a callback tuple member or another wrapper's dispatcher. Compose
+policies as `(outer, inner)` within one wrapper.
+
Default recursion uses the same type attributes as C++: it calls
`__s_maybe_inplace_mutate__` for a uniquely owned object when available, or
falls back to `__s_mutate__`. Each hook receives the active Rust-backed
diff --git a/rust/tvm-ffi/src/extra/structural_mutate.rs
b/rust/tvm-ffi/src/extra/structural_mutate.rs
index 80998f0a..04b2cfd5 100644
--- a/rust/tvm-ffi/src/extra/structural_mutate.rs
+++ b/rust/tvm-ffi/src/extra/structural_mutate.rs
@@ -71,6 +71,9 @@ pub use super::StructuralView;
/// Compatibility name for [`StructuralView`].
pub use super::StructuralView as MapValue;
+mod policy;
+pub use policy::{DefaultMutContextPolicy, MapWithContextPolicy,
MutContextPolicy};
+
/// Result type produced by a structural-map callback.
#[doc(hidden)]
pub type MapResult = Result<Any>;
@@ -559,6 +562,24 @@ where
self.def_region_kind
}
+ /// Run recursive operations in a definition region, preserving an outer
Pattern.
+ /// The previous context is restored on return, error, or unwinding.
+ pub fn with_def_region_kind<T>(
+ &mut self,
+ kind: DefRegionKind,
+ callback: impl FnOnce(&mut MutateContext<'_, State, Driver>) ->
Result<T>,
+ ) -> Result<T> {
+ with_mutation_region(kind, |kind| {
+ callback(&mut MutateContext {
+ driver: &mut *self.driver,
+ def_region_kind: kind,
+ inplace_mode: self.inplace_mode,
+ _state: PhantomData,
+ _not_send_sync: PhantomData,
+ })
+ })
+ }
+
/// Mutate a borrowed value through the same callback chain. The value and
/// its descendants begin on the non-in-place path.
#[inline(always)]
@@ -1019,7 +1040,8 @@ macro_rules! impl_mutate_chain_link {
impl_callback_chain_tuple_arities!(impl_mutate_chain_link);
/// A reusable typed-dispatch or callback mutator with shared user state.
-pub struct MutateCallbacks<State, Link, Marker> {
+pub struct MutateCallbacks<State, Link, Marker, Policy =
DefaultMutContextPolicy> {
+ policy: Option<Rc<Policy>>,
state: State,
callbacks: Rc<Link>,
_marker: PhantomData<fn(Marker)>,
@@ -1034,12 +1056,26 @@ where
Self {
state,
callbacks: Rc::new(callbacks),
+ policy: None,
_marker: PhantomData,
}
}
}
-impl<State, Link, Marker> MutateCallbacks<State, Link, Marker> {
+impl<State, Link, Marker, Policy> MutateCallbacks<State, Link, Marker, Policy>
{
+ /// Customize default descent while sharing the callback state.
+ pub fn with_policy<P: MutContextPolicy<State>>(
+ self,
+ policy: P,
+ ) -> MutateCallbacks<State, Link, Marker, P> {
+ MutateCallbacks {
+ policy: Some(Rc::new(policy)),
+ state: self.state,
+ callbacks: self.callbacks,
+ _marker: PhantomData,
+ }
+ }
+
/// Shared access to the callback state.
pub fn state(&self) -> &State {
&self.state
@@ -1067,7 +1103,9 @@ trait MutateCallbackState<State> {
fn callback_state_mut(&mut self) -> &mut State;
}
-impl<State, Link, Marker> MutateCallbackState<State> for
MutateCallbacks<State, Link, Marker> {
+impl<State, Link, Marker, Policy> MutateCallbackState<State>
+ for MutateCallbacks<State, Link, Marker, Policy>
+{
fn callback_state(&self) -> &State {
&self.state
}
@@ -1119,6 +1157,26 @@ pub trait MapDispatch: Sized {
) -> Option<MapResult>;
}
+/// Internal root-mapping protocol used by [`IntoMapper`].
+#[doc(hidden)]
+pub trait NativeMap: Sized {
+ fn map_root(&mut self, root: Any, order: WalkOrder) -> Result<Any>;
+}
+
+impl<D: MapDispatch> NativeMap for D {
+ fn map_root(&mut self, root: Any, order: WalkOrder) -> Result<Any> {
+ run_structural_mutator(
+ root,
+ &mut NativeMapper::<_, DefaultMutContextPolicy> {
+ dispatch: self,
+ order,
+ policy: None,
+ remap: StructuralVarRemap::default(),
+ },
+ )
+ }
+}
+
impl<V: MapDispatch> MapDispatch for &mut V {
#[inline]
fn dispatch_map(
@@ -1134,10 +1192,10 @@ impl<V: MapDispatch> MapDispatch for &mut V {
#[diagnostic::on_unimplemented(
message = "unsupported structural-map callback shape",
label = "this value cannot be used as a structural mapper",
- note = "pass `&mut` a type implementing `MapDispatch`, a supported
closure, or a tuple of callbacks"
+ note = "pass `&mut` a `MapDispatch`, a supported closure or callback
tuple, or a `MapWithContextPolicy`"
)]
pub trait IntoMapper<Marker> {
- type Mapper: MapDispatch;
+ type Mapper: NativeMap;
fn into_mapper(self) -> Self::Mapper;
}
@@ -1582,6 +1640,20 @@ pub trait StructuralMutator: Sized {
self.dispatch_mutate(value.as_value(), def_region_kind)
}
+ #[doc(hidden)]
+ fn dispatch_default_mutate(
+ &mut self,
+ value: MutateValue<'_>,
+ kind: DefRegionKind,
+ ) -> Result<Any> {
+ default_mutate_driver(
+ self,
+ value.value.raw(),
+ kind,
+ value.permit(value.inplace_mode()),
+ )
+ }
+
/// Re-enter this mutator for a borrowed value. The value and all of its
/// descendants use the non-in-place path.
fn mutate<T>(&mut self, value: &T, def_region_kind: DefRegionKind) ->
Result<Any>
@@ -1871,8 +1943,9 @@ where
)
}
-impl<State, Link, Marker> StructuralMutator for MutateCallbacks<State, Link,
Marker>
+impl<State, Link, Marker, Policy> StructuralMutator for MutateCallbacks<State,
Link, Marker, Policy>
where
+ Policy: MutContextPolicy<State>,
Link: MutateChainLink<State, Marker>,
Link::Strategy: MutateCallbackStrategy<State, Link, Marker>,
{
@@ -1915,6 +1988,27 @@ where
.map(Any::from),
}
}
+
+ fn dispatch_default_mutate(
+ &mut self,
+ value: MutateValue<'_>,
+ kind: DefRegionKind,
+ ) -> Result<Any> {
+ match self.policy.clone() {
+ Some(policy) => policy::mutate_with_policy(
+ &mut policy::MutationDescent { driver: self },
+ policy.as_ref(),
+ value,
+ kind,
+ ),
+ None => default_mutate_driver(
+ self,
+ value.value.raw(),
+ kind,
+ value.permit(value.inplace_mode()),
+ ),
+ }
+ }
}
impl<Link, Marker> StructuralMutator for DirectMutateCallbacks<'_, Link,
Marker>
@@ -1989,7 +2083,7 @@ where
#[inline(always)]
fn default_mutate_borrowed(&mut self, value: AnyView<'_>, kind:
DefRegionKind) -> Result<Any> {
- default_mutate_driver(self, *value.as_raw_ffi_any(), kind,
Permit::Copy)
+ user_default_mutate(self, *value.as_raw_ffi_any(), kind, Permit::Copy)
}
fn default_mutate_value(
@@ -1999,7 +2093,7 @@ where
mode: InplaceMode,
) -> Result<Any> {
let permit = value.permit(mode);
- default_mutate_driver(self, value.value.raw(), kind, permit)
+ user_default_mutate(self, value.value.raw(), kind, permit)
}
#[inline(always)]
@@ -2019,6 +2113,16 @@ enum Permit {
MaybeInPlace,
}
+impl Permit {
+ fn inplace_mode(self, raw: TVMFFIAny) -> InplaceMode {
+ if self == Self::MaybeInPlace && object_is_unique(raw) {
+ InplaceMode::Allow
+ } else {
+ InplaceMode::Disallow
+ }
+ }
+}
+
struct MemoEntry {
// Keeps the pointer-valued key alive so its address cannot be reused
// during the same mapping invocation.
@@ -2026,13 +2130,14 @@ struct MemoEntry {
result: Any,
}
-struct NativeMapper<D> {
- dispatch: D,
+struct NativeMapper<'a, D, Policy> {
+ dispatch: &'a mut D,
+ policy: Option<Rc<Policy>>,
order: WalkOrder,
remap: StructuralVarRemap,
}
-impl<D: MapDispatch> NativeMapper<D> {
+impl<D: MapDispatch, Policy: MutContextPolicy<D>> NativeMapper<'_, D, Policy> {
fn map_raw(
&mut self,
raw: TVMFFIAny,
@@ -2045,7 +2150,7 @@ impl<D: MapDispatch> NativeMapper<D> {
// Raw strings, byte-array views, and ObjectRValueRef are deliberately
// excluded because converting those borrowed special values into an
// Any performs normalization rather than a bitwise copy.
- if is_plain_inline(raw.type_index) {
+ if self.policy.is_none() && is_plain_inline(raw.type_index) {
let value = StructuralView::from_raw(raw);
return match self.dispatch.dispatch_map(&value, def_region_kind) {
Some(result) => {
@@ -2661,14 +2766,34 @@ fn def_region_from_raw(kind: i32) ->
Result<DefRegionKind> {
}
}
-impl<D: MapDispatch> MutationDriver for NativeMapper<D> {
+impl<D: MapDispatch, Policy: MutContextPolicy<D>> MutationDriver for
NativeMapper<'_, D, Policy> {
fn dispatch_raw(
&mut self,
raw: TVMFFIAny,
def_region_kind: DefRegionKind,
permit: Permit,
) -> Result<Any> {
- self.map_raw(raw, def_region_kind, permit)
+ with_mutation_region(def_region_kind, |kind| self.map_raw(raw, kind,
permit))
+ }
+
+ fn default_map_current_raw(
+ &mut self,
+ raw: TVMFFIAny,
+ kind: DefRegionKind,
+ permit: Permit,
+ ) -> Result<Any> {
+ match self.policy.clone() {
+ Some(policy) => {
+ let view = StructuralView::from_raw(raw);
+ policy::mutate_with_policy(
+ &mut policy::MutationDescent { driver: self },
+ policy.as_ref(),
+ MutateValue::new(&view, permit.inplace_mode(raw)),
+ kind,
+ )
+ }
+ None => default_mutate_driver(self, raw, kind, permit),
+ }
}
fn var_remap_get_raw(&mut self, raw: TVMFFIAny) -> Result<Option<Any>> {
@@ -2955,6 +3080,18 @@ fn with_mutator_def_region<T>(
}
}
+fn with_mutation_region<T>(
+ kind: DefRegionKind,
+ callback: impl FnOnce(DefRegionKind) -> Result<T>,
+) -> Result<T> {
+ let mutator = active_mutator()?;
+ with_mutator_def_region(mutator, kind, || {
+ // SAFETY: the active invocation keeps this thread's ABI mutator alive.
+ let effective = def_region_from_raw(unsafe {
(*mutator).def_region_mode })?;
+ callback(effective)
+ })
+}
+
#[inline(always)]
fn dispatch_user_raw<U: StructuralMutator>(
mutator: &mut U,
@@ -2962,14 +3099,15 @@ fn dispatch_user_raw<U: StructuralMutator>(
def_region_kind: DefRegionKind,
permit: Permit,
) -> Result<Any> {
- let result = if permit == Permit::MaybeInPlace && object_is_unique(raw) {
- let mut scoped_raw = raw;
- mutator
- .dispatch_maybe_inplace_mutate(InplaceValue::from_raw(&mut
scoped_raw), def_region_kind)
- } else {
- mutator.dispatch_mutate(&StructuralView::from_raw(raw),
def_region_kind)
- };
- result.map_err(|error| with_value_context(error, raw))
+ with_mutation_region(def_region_kind, |kind| {
+ let result = if permit == Permit::MaybeInPlace &&
object_is_unique(raw) {
+ let mut scoped_raw = raw;
+ mutator.dispatch_maybe_inplace_mutate(InplaceValue::from_raw(&mut
scoped_raw), kind)
+ } else {
+ mutator.dispatch_mutate(&StructuralView::from_raw(raw), kind)
+ };
+ result.map_err(|error| with_value_context(error, raw))
+ })
}
fn user_default_mutate<U: StructuralMutator>(
@@ -2978,7 +3116,10 @@ fn user_default_mutate<U: StructuralMutator>(
def_region_kind: DefRegionKind,
permit: Permit,
) -> Result<Any> {
- default_mutate_driver(mutator, raw, def_region_kind, permit)
+ with_mutation_region(def_region_kind, |kind| {
+ let value = StructuralView::from_raw(raw);
+ mutator.dispatch_default_mutate(MutateValue::new(&value,
permit.inplace_mode(raw)), kind)
+ })
}
fn default_mutate_driver<D: MutationDriver>(
@@ -3055,13 +3196,7 @@ where
R: Into<Any>,
H: IntoMapper<M>,
{
- let root = root.into();
- let mut native = NativeMapper {
- dispatch: mapper.into_mapper(),
- order,
- remap: StructuralVarRemap::default(),
- };
- run_structural_mutator(root, &mut native)
+ mapper.into_mapper().map_root(root.into(), order)
}
fn shallow_copy(raw: TVMFFIAny) -> Result<Any> {
diff --git a/rust/tvm-ffi/src/extra/structural_mutate/policy.rs
b/rust/tvm-ffi/src/extra/structural_mutate/policy.rs
new file mode 100644
index 00000000..1ebe7a0d
--- /dev/null
+++ b/rust/tvm-ffi/src/extra/structural_mutate/policy.rs
@@ -0,0 +1,280 @@
+/*
+ * 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.
+ */
+
+//! Context management around default structural mutation.
+
+use super::*;
+
+#[cfg(doctest)]
+mod compile_fail;
+
+/// Default-recursion policy for [`MutateCallbacks::with_policy`] and
[`MapWithContextPolicy`].
+///
+/// `ctx.default_maybe_inplace_mutate_result(value)` continues to the next
policy,
+/// then hooks or reflected fields; children re-enter callback dispatch.
+/// Policies share callback state and compose as `(outer, inner)`.
+/// Restore user state before returning, including on errors; use
+/// [`MutateContext::with_def_region_kind`] for scoped definition regions.
+pub trait MutContextPolicy<State> {
+ /// Customize default descent, preserving the input permission and result
marker.
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut MutateContext<'_, State>,
+ ) -> Result<UnchangedOr<Any>>;
+}
+
+/// Default descent through registered hooks or reflected structural fields.
+pub struct DefaultMutContextPolicy;
+
+impl<State> MutContextPolicy<State> for DefaultMutContextPolicy {
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut MutateContext<'_, State>,
+ ) -> Result<UnchangedOr<Any>> {
+ ctx.default_maybe_inplace_mutate_result(value)
+ }
+}
+
+impl<State, Outer: MutContextPolicy<State>, Inner: MutContextPolicy<State>>
MutContextPolicy<State>
+ for (Outer, Inner)
+{
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut MutateContext<'_, State>,
+ ) -> Result<UnchangedOr<Any>> {
+ let kind = ctx.def_region_kind();
+ mutate_with_policy(
+ &mut NextPolicy {
+ driver: &mut *ctx.driver,
+ policy: &self.1,
+ },
+ &self.0,
+ value,
+ kind,
+ )
+ .and_then(UnchangedOr::from_carrier)
+ }
+}
+
+pub(super) fn mutate_with_policy<State>(
+ driver: &mut dyn MutateContextDriver<State>,
+ policy: &impl MutContextPolicy<State>,
+ value: MutateValue<'_>,
+ kind: DefRegionKind,
+) -> Result<Any> {
+ with_mutation_region(kind, |kind| {
+ let mut ctx = MutateContext {
+ driver,
+ def_region_kind: kind,
+ inplace_mode: value.inplace_mode(),
+ _state: PhantomData,
+ _not_send_sync: PhantomData,
+ };
+ policy.default_mutate(value, &mut ctx).map(Any::from)
+ })
+}
+
+struct NextPolicy<'a, State, Policy> {
+ driver: &'a mut dyn MutateContextDriver<State>,
+ policy: &'a Policy,
+}
+
+impl<State, Policy: MutContextPolicy<State>> MutateContextDriver<State>
+ for NextPolicy<'_, State, Policy>
+{
+ fn state(&self) -> &State {
+ self.driver.state()
+ }
+ fn state_mut(&mut self) -> &mut State {
+ self.driver.state_mut()
+ }
+ fn mutate_borrowed(&mut self, value: AnyView<'_>, kind: DefRegionKind) ->
Result<Any> {
+ self.driver.mutate_borrowed(value, kind)
+ }
+ fn mutate_owned(&mut self, value: Any, kind: DefRegionKind, mode:
InplaceMode) -> Result<Any> {
+ self.driver.mutate_owned(value, kind, mode)
+ }
+ fn default_mutate_borrowed(&mut self, value: AnyView<'_>, kind:
DefRegionKind) -> Result<Any> {
+ let view = StructuralView::from_raw(*value.as_raw_ffi_any());
+ mutate_with_policy(self.driver, self.policy,
MutateValue::borrowed(&view), kind)
+ }
+ fn default_mutate_value(
+ &mut self,
+ mut value: MutateValue<'_>,
+ kind: DefRegionKind,
+ mode: InplaceMode,
+ ) -> Result<Any> {
+ value.mode = value.permit(mode).inplace_mode(value.value.raw());
+ mutate_with_policy(self.driver, self.policy, value, kind)
+ }
+ fn var_remap_get(&mut self, var: &StructuralView) -> Result<Option<Any>> {
+ self.driver.var_remap_get(var)
+ }
+ fn var_remap_set(&mut self, var: &StructuralView, replacement: &Any) ->
Result<()> {
+ self.driver.var_remap_set(var, replacement)
+ }
+}
+
+pub(super) struct MutationDescent<'a, Driver> {
+ pub(super) driver: &'a mut Driver,
+}
+
+impl<State, Driver: MutationDriver + MutateCallbackState<State>>
MutateContextDriver<State>
+ for MutationDescent<'_, Driver>
+{
+ fn state(&self) -> &State {
+ self.driver.callback_state()
+ }
+ fn state_mut(&mut self) -> &mut State {
+ self.driver.callback_state_mut()
+ }
+ fn mutate_borrowed(&mut self, value: AnyView<'_>, kind: DefRegionKind) ->
Result<Any> {
+ with_mutation_region(kind, |kind| {
+ self.driver
+ .dispatch_raw(*value.as_raw_ffi_any(), kind, Permit::Copy)
+ })
+ }
+ fn mutate_owned(&mut self, value: Any, kind: DefRegionKind, mode:
InplaceMode) -> Result<Any> {
+ let raw = *value.as_raw_ffi_any();
+ let result = with_mutation_region(kind, |kind| {
+ self.driver.dispatch_raw(raw, kind, mode.permit())
+ })?;
+ Ok(if is_unchanged(&result) { value } else { result })
+ }
+ fn default_mutate_borrowed(&mut self, value: AnyView<'_>, kind:
DefRegionKind) -> Result<Any> {
+ with_mutation_region(kind, |kind| {
+ default_mutate_driver(self.driver, *value.as_raw_ffi_any(), kind,
Permit::Copy)
+ })
+ }
+ fn default_mutate_value(
+ &mut self,
+ value: MutateValue<'_>,
+ kind: DefRegionKind,
+ mode: InplaceMode,
+ ) -> Result<Any> {
+ let permit = value.permit(mode);
+ with_mutation_region(kind, |kind| {
+ default_mutate_driver(self.driver, value.value.raw(), kind, permit)
+ })
+ }
+ fn var_remap_get(&mut self, var: &StructuralView) -> Result<Option<Any>> {
+ self.driver.var_remap_get_raw(var.raw())
+ }
+ fn var_remap_set(&mut self, var: &StructuralView, replacement: &Any) ->
Result<()> {
+ self.driver.var_remap_set_raw(var.raw(), replacement)
+ }
+}
+
+impl<D, Policy> MutateCallbackState<D> for NativeMapper<'_, D, Policy> {
+ fn callback_state(&self) -> &D {
+ self.dispatch
+ }
+ fn callback_state_mut(&mut self) -> &mut D {
+ self.dispatch
+ }
+}
+
+/// A [`MapDispatch`] with a [`MutContextPolicy`], sharing the dispatcher's
state.
+///
+/// Run with [`Self::map`] or [`structural_map`]. Each node's map callback runs
+/// outside its policy scope; its children run inside.
+/// This mapper cannot be a callback tuple member or another wrapper's
dispatcher.
+/// Compose policies as `(outer, inner)` within one wrapper.
+pub struct MapWithContextPolicy<Mapper, Policy> {
+ mapper: Mapper,
+ policy: Rc<Policy>,
+}
+
+impl<Mapper: MapDispatch, Policy: MutContextPolicy<Mapper>>
MapWithContextPolicy<Mapper, Policy> {
+ /// Combine a dispatcher and a default-recursion policy.
+ pub fn new(mapper: Mapper, policy: Policy) -> Self {
+ Self {
+ mapper,
+ policy: Rc::new(policy),
+ }
+ }
+
+ /// Access the dispatcher and shared state.
+ pub fn state(&self) -> &Mapper {
+ &self.mapper
+ }
+
+ /// Mutably access the dispatcher outside recursive calls.
+ pub fn state_mut(&mut self) -> &mut Mapper {
+ &mut self.mapper
+ }
+
+ /// Recover the dispatcher and its state.
+ pub fn into_state(self) -> Mapper {
+ self.mapper
+ }
+
+ /// Map an owning root with fresh invocation-local identity substitutions.
+ pub fn map<R: Into<Any>>(&mut self, root: R, order: WalkOrder) ->
Result<Any> {
+ run_structural_mutator(
+ root.into(),
+ &mut NativeMapper {
+ dispatch: &mut self.mapper,
+ order,
+ policy: Some(self.policy.clone()),
+ remap: StructuralVarRemap::default(),
+ },
+ )
+ }
+}
+
+impl<Mapper: MapDispatch, Policy: MutContextPolicy<Mapper>> NativeMap
+ for MapWithContextPolicy<Mapper, Policy>
+{
+ fn map_root(&mut self, root: Any, order: WalkOrder) -> Result<Any> {
+ self.map(root, order)
+ }
+}
+
+impl<Mapper: MapDispatch, Policy: MutContextPolicy<Mapper>> NativeMap
+ for &mut MapWithContextPolicy<Mapper, Policy>
+{
+ fn map_root(&mut self, root: Any, order: WalkOrder) -> Result<Any> {
+ self.map(root, order)
+ }
+}
+
+#[doc(hidden)]
+pub enum ByPolicyMap {}
+
+impl<Mapper: MapDispatch, Policy: MutContextPolicy<Mapper>>
IntoMapper<ByPolicyMap>
+ for MapWithContextPolicy<Mapper, Policy>
+{
+ type Mapper = Self;
+ fn into_mapper(self) -> Self {
+ self
+ }
+}
+
+impl<'a, Mapper: MapDispatch, Policy: MutContextPolicy<Mapper>>
IntoMapper<ByPolicyMap>
+ for &'a mut MapWithContextPolicy<Mapper, Policy>
+{
+ type Mapper = Self;
+ fn into_mapper(self) -> Self {
+ self
+ }
+}
diff --git a/rust/tvm-ffi/src/extra/structural_mutate/policy/compile_fail.rs
b/rust/tvm-ffi/src/extra/structural_mutate/policy/compile_fail.rs
new file mode 100644
index 00000000..16bb79be
--- /dev/null
+++ b/rust/tvm-ffi/src/extra/structural_mutate/policy/compile_fail.rs
@@ -0,0 +1,50 @@
+/*
+ * 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.
+ */
+
+//! A policy mapper cannot be used as a callback, even in a one-element tuple.
+//! ```compile_fail,E0277
+//! use tvm_ffi::*;
+//! fn invalid<D: MapDispatch>(mapper: &mut MapWithContextPolicy<D,
DefaultMutContextPolicy>) {
+//! structural_map(1_i64, (mapper,), WalkOrder::PreOrder).unwrap();
+//! }
+//! ```
+//!
+//! Adding another callback must not discard the policy.
+//! ```compile_fail,E0277
+//! use tvm_ffi::*;
+//! fn invalid<D: MapDispatch>(mapper: &mut MapWithContextPolicy<D,
DefaultMutContextPolicy>) {
+//! structural_map(1_i64, (mapper, |s: String| s),
WalkOrder::PreOrder).unwrap();
+//! }
+//! ```
+//!
+//! Nested callback tuples obey the same restriction.
+//! ```compile_fail,E0277
+//! use tvm_ffi::*;
+//! fn invalid<D: MapDispatch>(mapper: &mut MapWithContextPolicy<D,
DefaultMutContextPolicy>) {
+//! structural_map(1_i64, (|s: String| s, ((mapper,),)),
WalkOrder::PostOrder).unwrap();
+//! }
+//! ```
+//!
+//! Compose policies in a tuple instead of nesting mapper wrappers.
+//! ```compile_fail,E0277
+//! use tvm_ffi::*;
+//! fn invalid<D: MapDispatch>(mapper: MapWithContextPolicy<D,
DefaultMutContextPolicy>) {
+//! MapWithContextPolicy::new(mapper, DefaultMutContextPolicy);
+//! }
+//! ```
diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs
b/rust/tvm-ffi/src/extra/structural_visit.rs
index 5cea6591..eedce47c 100644
--- a/rust/tvm-ffi/src/extra/structural_visit.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit.rs
@@ -223,12 +223,7 @@ impl From<Error> for NativeHalt {
type NativeResult = std::result::Result<(), NativeHalt>;
mod policy;
-pub use policy::{ContextPolicy, DefaultContextPolicy, WalkWithPolicy};
-
-/// Compatibility name for [`ContextPolicy`].
-pub use policy::ContextPolicy as VisitPolicy;
-/// Compatibility name for [`DefaultContextPolicy`].
-pub use policy::DefaultContextPolicy as DefaultVisitPolicy;
+pub use policy::{ContextPolicy, DefaultContextPolicy, WalkWithContextPolicy};
/// State and recursive operations available to a visit callback.
///
diff --git a/rust/tvm-ffi/src/extra/structural_visit/policy.rs
b/rust/tvm-ffi/src/extra/structural_visit/policy.rs
index 502f0df4..fed91afc 100644
--- a/rust/tvm-ffi/src/extra/structural_visit/policy.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit/policy.rs
@@ -155,12 +155,12 @@ impl<State, V: StructuralVisitor +
VisitCallbackState<State>> VisitContextDriver
/// 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> {
+pub struct WalkWithContextPolicy<Walker, Policy> {
walker: Walker,
policy: Rc<Policy>,
}
-impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
WalkWithPolicy<Walker, Policy> {
+impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
WalkWithContextPolicy<Walker, Policy> {
/// Combine a dispatcher and a default-recursion policy.
pub fn new(walker: Walker, policy: Policy) -> Self {
Self {
@@ -205,7 +205,7 @@ impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
WalkWithPolicy<Walker,
pub enum ByPolicyWalk {}
impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
IntoWalker<ByPolicyWalk>
- for WalkWithPolicy<Walker, Policy>
+ for WalkWithContextPolicy<Walker, Policy>
{
type Walker = Self;
fn into_walker(self) -> Self {
@@ -214,7 +214,7 @@ impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
IntoWalker<ByPolicyWal
}
impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>> NativeVisit
- for WalkWithPolicy<Walker, Policy>
+ for WalkWithContextPolicy<Walker, Policy>
{
const CUSTOM_DESCENT: bool = true;
@@ -245,7 +245,7 @@ impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>>
NativeVisit
}
struct WalkDescent<'a, Walker, Policy, const PRE_ORDER: bool> {
- visitor: &'a mut WalkWithPolicy<Walker, Policy>,
+ visitor: &'a mut WalkWithContextPolicy<Walker, Policy>,
}
impl<Walker: WalkDispatch, Policy: ContextPolicy<Walker>, const PRE_ORDER:
bool>
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index f6ddd640..eadc138c 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -49,15 +49,16 @@ pub use crate::error::{
};
pub use crate::extra::module::Module;
pub use crate::extra::structural_mutate::{
- structural_map, structural_mutate, CallbackMutator, InplaceMode,
InplaceValue, IntoMapResult,
- IntoMapper, IntoMutator, MapChainLink, MapDispatch, MapValue,
MutateCallbacks, MutateChainLink,
- MutateContext, MutateDispatch, MutateValue, Mutator, StructuralMutator,
StructuralVarRemap,
+ structural_map, structural_mutate, CallbackMutator,
DefaultMutContextPolicy, InplaceMode,
+ InplaceValue, IntoMapResult, IntoMapper, IntoMutator, MapChainLink,
MapDispatch, MapValue,
+ MapWithContextPolicy, MutContextPolicy, MutateCallbacks, MutateChainLink,
MutateContext,
+ MutateDispatch, MutateValue, Mutator, StructuralMutator,
StructuralVarRemap,
};
pub use crate::extra::structural_visit::{
structural_visit, structural_walk, ContextPolicy, DefRegionKind,
DefaultContextPolicy,
- DefaultVisitPolicy, IntoVisitor, IntoWalkResult, IntoWalker,
StructuralVisitor, VisitCallbacks,
- VisitChainLink, VisitContext, VisitInterrupt, VisitPolicy, VisitValue,
WalkChainLink,
- WalkDispatch, WalkOrder, WalkResult, WalkWithPolicy,
+ IntoVisitor, IntoWalkResult, IntoWalker, StructuralVisitor,
VisitCallbacks, VisitChainLink,
+ VisitContext, VisitInterrupt, VisitValue, WalkChainLink, WalkDispatch,
WalkOrder, WalkResult,
+ WalkWithContextPolicy,
};
pub use crate::extra::unchanged::{Unchanged, UnchangedOr};
pub use crate::extra::StructuralView;
diff --git a/rust/tvm-ffi/tests/test_structural_mutate.rs
b/rust/tvm-ffi/tests/test_structural_mutate.rs
index b7647603..70183e9c 100644
--- a/rust/tvm-ffi/tests/test_structural_mutate.rs
+++ b/rust/tvm-ffi/tests/test_structural_mutate.rs
@@ -23,8 +23,9 @@ use tvm_ffi::function::FunctionObj;
use tvm_ffi::object::ObjectRef;
use tvm_ffi::{
dispatch, structural_map, structural_mutate, Any, AnyView, Array,
CallbackMutator,
- DefRegionKind, Error, FieldGetter, Function, InplaceMode, InplaceValue,
Map, MapDispatch,
- MapValue, MutateCallbacks, MutateValue, Mutator, Object, ObjectArc,
ObjectRefCore, Result,
+ DefRegionKind, DefaultMutContextPolicy, Error, FieldGetter, Function,
InplaceMode,
+ InplaceValue, IntoMapper, Map, MapDispatch, MapValue,
MapWithContextPolicy, MutContextPolicy,
+ MutateCallbacks, MutateValue, Mutator, Object, ObjectArc, ObjectRefCore,
Result,
String as FfiString, StructuralMutator, StructuralVarRemap, TypeIndex,
Unchanged, UnchangedOr,
WalkOrder, RUNTIME_ERROR,
};
@@ -1649,3 +1650,450 @@ fn
callback_mutate_panics_resume_and_leave_the_next_run_usable() {
.unwrap();
assert_eq!(mutated.get(0).unwrap(), 2);
}
+
+#[derive(Default)]
+struct PolicyState {
+ depth: usize,
+ events: Vec<(&'static str, usize)>,
+}
+#[dispatch(map)]
+impl PolicyState {
+ fn map_integer(&mut self, value: i64) -> i64 {
+ self.events.push(("integer", self.depth));
+ value + 1
+ }
+ fn map_array(&mut self, value: Array<Any>) -> Any {
+ self.events.push(("callback", self.depth));
+ value.into()
+ }
+}
+
+struct ArrayPolicy;
+impl MutContextPolicy<PolicyState> for ArrayPolicy {
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut CallbackMutator<PolicyState>,
+ ) -> Result<UnchangedOr<Any>> {
+ if value
+ .as_node::<tvm_ffi::collections::array::ArrayObj>()
+ .is_none()
+ {
+ return ctx.default_maybe_inplace_mutate_result(value);
+ }
+ ctx.state_mut().depth += 1;
+ let depth = ctx.state().depth;
+ ctx.state_mut().events.push(("enter", depth));
+ let result = ctx.default_maybe_inplace_mutate_result(value);
+ ctx.state_mut().depth -= 1;
+ let depth = ctx.state().depth;
+ ctx.state_mut().events.push(("exit", depth));
+ result
+ }
+}
+struct RecordPolicy;
+impl MutContextPolicy<PolicyState> for RecordPolicy {
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut CallbackMutator<PolicyState>,
+ ) -> Result<UnchangedOr<Any>> {
+ if value
+ .as_node::<tvm_ffi::collections::array::ArrayObj>()
+ .is_some()
+ {
+ let depth = ctx.state().depth;
+ ctx.state_mut().events.push(("next", depth));
+ }
+ ctx.default_maybe_inplace_mutate_result(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)]);
+ let pre = vec![
+ ("callback", 0),
+ ("enter", 1),
+ ("next", 1),
+ ("callback", 1),
+ ("enter", 2),
+ ("next", 2),
+ ("integer", 2),
+ ("exit", 1),
+ ("integer", 1),
+ ("exit", 0),
+ ];
+ let post = vec![
+ ("enter", 1),
+ ("next", 1),
+ ("enter", 2),
+ ("next", 2),
+ ("integer", 2),
+ ("exit", 1),
+ ("callback", 1),
+ ("integer", 1),
+ ("exit", 0),
+ ("callback", 0),
+ ];
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut mapper = MapWithContextPolicy::new(
+ PolicyState::default(),
+ (ArrayPolicy, (RecordPolicy, DefaultMutContextPolicy)),
+ );
+ let output = structural_map(root(), &mut mapper, order).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!(mapper.state().depth, 0);
+ assert_eq!(
+ &mapper.state().events,
+ if order == WalkOrder::PreOrder {
+ &pre
+ } else {
+ &post
+ }
+ );
+ }
+ let mut mutator = MutateCallbacks::new(
+ PolicyState::default(),
+ (
+ |x: i64, ctx: &mut CallbackMutator<PolicyState>| {
+ let depth = ctx.state().depth;
+ ctx.state_mut().events.push(("integer", depth));
+ x + 1
+ },
+ |value: MutateValue<'_>, ctx: &mut CallbackMutator<PolicyState>| {
+ let depth = ctx.state().depth;
+ ctx.state_mut().events.push(("callback", depth));
+ ctx.default_maybe_inplace_mutate_result(value)
+ },
+ ),
+ )
+ .with_policy((ArrayPolicy, RecordPolicy));
+ let output = structural_mutate(root(), &mut mutator).unwrap();
+ assert_eq!(i64::try_from(array_item(&output, 1)).unwrap(), 3);
+ assert_eq!(mutator.state().events, pre);
+ assert_eq!(mutator.state().depth, 0);
+}
+
+#[test]
+fn map_policy_entries_preserve_descent_and_callback_composition() {
+ struct Stop;
+ impl<State> MutContextPolicy<State> for Stop {
+ fn default_mutate(
+ &self,
+ _: MutateValue<'_>,
+ _: &mut CallbackMutator<State>,
+ ) -> Result<UnchangedOr<Any>> {
+ Ok(UnchangedOr::unchanged())
+ }
+ }
+
+ let root = || Array::new(vec![1_i64]);
+ let first = |value: Any|
Array::<i64>::try_from(value).unwrap().get(0).unwrap();
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ for entry in 0..3 {
+ let mut state = PolicyState::default();
+ let output = {
+ let mut mapper =
+ MapWithContextPolicy::new(&mut state,
(DefaultMutContextPolicy, Stop));
+ match entry {
+ 0 => structural_map(root(), mapper, order),
+ 1 => structural_map(root(), &mut mapper, order),
+ _ => mapper.map(root(), order),
+ }
+ }
+ .unwrap();
+ assert_eq!(first(output), 1);
+ assert_eq!(state.events, vec![("callback", 0)]);
+ }
+
+ let mut dispatch = IncrementIntegers;
+ let callbacks = (|s: FfiString| s, (&mut dispatch,));
+ assert_eq!(first(structural_map(root(), callbacks, order).unwrap()),
2);
+
+ let callbacks = (|x: i64| x + 1, (|s: FfiString| s,));
+ assert_eq!(first(structural_map(root(), callbacks, order).unwrap()),
2);
+ let mapper = MapWithContextPolicy::new(callbacks.into_mapper(), Stop);
+ assert_eq!(first(structural_map(root(), mapper, order).unwrap()), 1);
+ }
+}
+
+#[test]
+fn mutation_policy_continuations_preserve_ownership_and_markers() {
+ struct Ownership {
+ increment: bool,
+ retained: Option<Any>,
+ }
+ #[dispatch(map)]
+ impl Ownership {
+ fn map_integer(&mut self, x: i64) -> UnchangedOr<i64> {
+ if self.increment {
+ UnchangedOr::changed(x + 1)
+ } else {
+ UnchangedOr::unchanged()
+ }
+ }
+ }
+ struct Control {
+ mode: InplaceMode,
+ retain: bool,
+ }
+ impl MutContextPolicy<Ownership> for Control {
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut CallbackMutator<Ownership>,
+ ) -> Result<UnchangedOr<Any>> {
+ let array = value
+ .as_node::<tvm_ffi::collections::array::ArrayObj>()
+ .is_some();
+ if array && self.retain {
+ ctx.state_mut().retained = Some(value.to_owned());
+ }
+ let result = ctx.default_mutate_with_mode_result(value,
self.mode)?;
+ if array && !ctx.state().increment {
+ assert!(result.is_unchanged());
+ }
+ Ok(result)
+ }
+ }
+ for entry in 0..3 {
+ // pre-order map, post-order map, mutation callback
+ for case in 0..4 {
+ // unique, shared, alias retained by policy, forced copy
+ for increment in [false, true] {
+ let root = Array::new(vec![1_i64]);
+ let pointer = array_pointer(&root);
+ let alias = (case == 1).then(|| root.clone());
+ let policy = (
+ Control {
+ mode: if case == 3 {
+ InplaceMode::Disallow
+ } else {
+ InplaceMode::Allow
+ },
+ retain: case == 2,
+ },
+ DefaultMutContextPolicy,
+ );
+ let state = Ownership {
+ increment,
+ retained: None,
+ };
+ let (output, state) = if entry < 2 {
+ let mut mapper = MapWithContextPolicy::new(state, policy);
+ let output = mapper
+ .map(
+ root,
+ if entry == 0 {
+ WalkOrder::PreOrder
+ } else {
+ WalkOrder::PostOrder
+ },
+ )
+ .unwrap();
+ (output, mapper.into_state())
+ } else {
+ let mut mutator = MutateCallbacks::new(
+ state,
+ (
+ |x: i64, ctx: &mut CallbackMutator<Ownership>| {
+ ctx.state_mut().map_integer(x)
+ },
+ |value: MutateValue<'_>, ctx: &mut
CallbackMutator<Ownership>| {
+ ctx.default_maybe_inplace_mutate_result(value)
+ },
+ ),
+ )
+ .with_policy(policy);
+ let output = structural_mutate(root, &mut
mutator).unwrap();
+ (output, mutator.into_state())
+ };
+ let output = Array::<i64>::try_from(output).unwrap();
+ assert_eq!(output.get(0).unwrap(), if increment { 2 } else { 1
});
+ assert_eq!(array_pointer(&output) == pointer, !increment ||
case == 0);
+ if let Some(alias) = alias {
+ assert_eq!(alias.get(0).unwrap(), 1);
+ }
+ if let Some(alias) = state.retained {
+
assert_eq!(Array::<i64>::try_from(alias).unwrap().get(0).unwrap(), 1);
+ }
+ }
+ }
+ }
+}
+
+#[test]
+fn mutation_policy_regions_retargeting_and_error_restore() {
+ use DefRegionKind::{None as Use, Pattern, Simple};
+ #[derive(Default)]
+ struct Regions(Vec<(i64, DefRegionKind)>);
+ #[dispatch(map)]
+ impl Regions {
+ fn map_integer(&mut self, x: i64, kind: DefRegionKind) -> i64 {
+ self.0.push((x, kind));
+ x
+ }
+ fn map_array(&mut self, value: Array<i64>, kind: DefRegionKind) -> Any
{
+ self.0.push((200, kind));
+ value.into()
+ }
+ }
+ struct Redirect(DefRegionKind);
+ impl MutContextPolicy<Regions> for Redirect {
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut CallbackMutator<Regions>,
+ ) -> Result<UnchangedOr<Any>> {
+ if value.cast::<bool>() == Some(false) {
+ // Bypass this container's callback, but enter the next policy
and redispatch its children.
+ return ctx.with_def_region_kind(Pattern, |ctx| {
+ ctx.default_mutate(&Array::new(vec![1_i64]))
+ .map(UnchangedOr::changed)
+ });
+ }
+ if value.cast::<i64>() == Some(2) {
+ assert_eq!(ctx.def_region_kind(), Pattern);
+ ctx.mutate_with(&99_i64, self.0)?;
+ let error: Result<()> = ctx.with_def_region_kind(Use, |ctx| {
+ assert_eq!(ctx.def_region_kind(), Pattern);
+ Err(Error::new(RUNTIME_ERROR, "scoped error", ""))
+ });
+ assert!(error.is_err());
+ assert_eq!(ctx.def_region_kind(), Pattern);
+ }
+ ctx.default_maybe_inplace_mutate_result(value)
+ }
+ }
+ struct Observe;
+ impl MutContextPolicy<Regions> for Observe {
+ fn default_mutate(
+ &self,
+ value: MutateValue<'_>,
+ ctx: &mut CallbackMutator<Regions>,
+ ) -> Result<UnchangedOr<Any>> {
+ if value
+ .as_node::<tvm_ffi::collections::array::ArrayObj>()
+ .is_some()
+ {
+ assert_eq!(ctx.def_region_kind(), Pattern);
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push((100, kind));
+ }
+ ctx.default_maybe_inplace_mutate_result(value)
+ }
+ }
+ assert_eq!(
+ unsafe { tvm_ffi::tvm_ffi_sys::TVMFFITestingDummyTarget() },
+ 0
+ );
+ for requested in [Use, Simple] {
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut mapper =
+ MapWithContextPolicy::new(Regions::default(),
(Redirect(requested), Observe));
+ let output = mapper.map(false, order).unwrap();
+ assert_eq!(i64::try_from(array_item(&output, 0)).unwrap(), 1);
+ let mut expected = vec![(100, Pattern), (1, Pattern)];
+ if order == WalkOrder::PostOrder {
+ expected.push((200, Use)); // The root callback sees the
result after descent.
+ }
+ assert_eq!(mapper.state().0, expected);
+ mapper.state_mut().0.clear();
+ let graph = Function::get_global("testing.make_visit_region_graph")
+ .unwrap()
+ .call_tuple((false,))
+ .unwrap();
+ mapper.map(graph, order).unwrap();
+ let expected = if order == WalkOrder::PreOrder {
+ vec![(1, Simple), (2, Pattern), (99, Pattern), (3, Use)]
+ } else {
+ vec![(1, Simple), (99, Pattern), (2, Pattern), (3, Use)]
+ };
+ assert_eq!(mapper.state().0, expected);
+ }
+ let mut mutator = MutateCallbacks::new(
+ Regions::default(),
+ |value: MutateValue<'_>, ctx: &mut CallbackMutator<Regions>| {
+ if let Some(x) = value.cast::<i64>() {
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push((x, kind));
+ }
+ if value
+ .as_node::<tvm_ffi::collections::array::ArrayObj>()
+ .is_some()
+ {
+ let kind = ctx.def_region_kind();
+ ctx.state_mut().0.push((200, kind));
+ }
+ ctx.default_maybe_inplace_mutate_result(value)
+ },
+ )
+ .with_policy((Redirect(requested), Observe));
+ let output = structural_mutate(false, &mut mutator).unwrap();
+ assert_eq!(i64::try_from(array_item(&output, 0)).unwrap(), 1);
+ assert_eq!(mutator.state().0, vec![(100, Pattern), (1, Pattern)]);
+ mutator.state_mut().0.clear();
+ let graph = Function::get_global("testing.make_visit_region_graph")
+ .unwrap()
+ .call_tuple((false,))
+ .unwrap();
+ structural_mutate(graph, &mut mutator).unwrap();
+ assert_eq!(
+ mutator.state().0,
+ vec![(1, Simple), (2, Pattern), (99, Pattern), (3, Use)]
+ );
+ }
+}
+
+#[test]
+fn mutation_policy_halts_restore_state_and_skip_later_policies() {
+ struct Halt(bool);
+ impl MutContextPolicy<PolicyState> for Halt {
+ fn default_mutate(
+ &self,
+ _: MutateValue<'_>,
+ _: &mut CallbackMutator<PolicyState>,
+ ) -> Result<UnchangedOr<Any>> {
+ if self.0 {
+ Err(Error::new(RUNTIME_ERROR, "stop descent", ""))
+ } else {
+ Ok(UnchangedOr::unchanged())
+ }
+ }
+ }
+ for fail in [false, true] {
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut mapper = MapWithContextPolicy::new(
+ PolicyState::default(),
+ (ArrayPolicy, (Halt(fail), RecordPolicy)),
+ );
+ let result = mapper.map(Array::new(vec![1_i64]), order);
+ assert_eq!(result.is_err(), fail);
+ assert_eq!(mapper.state().depth, 0);
+ assert!(!mapper
+ .state()
+ .events
+ .iter()
+ .any(|(tag, _)| matches!(*tag, "integer" | "next")));
+ if let Err(error) = result {
+ assert!(error.message().contains("stop descent"));
+ }
+ }
+ let mut mutator = MutateCallbacks::new(
+ PolicyState::default(),
+ |x: i64, _: &mut CallbackMutator<PolicyState>| x + 1,
+ )
+ .with_policy((ArrayPolicy, (Halt(fail), RecordPolicy)));
+ assert_eq!(
+ structural_mutate(Array::new(vec![1_i64]), &mut mutator).is_err(),
+ fail
+ );
+ assert_eq!(mutator.state().events, vec![("enter", 1), ("exit", 0)]);
+ assert_eq!(mutator.state().depth, 0);
+ }
+}
diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs
b/rust/tvm-ffi/tests/test_structural_visit.rs
index 2858844f..22137fa6 100644
--- a/rust/tvm-ffi/tests/test_structural_visit.rs
+++ b/rust/tvm-ffi/tests/test_structural_visit.rs
@@ -23,7 +23,8 @@ use tvm_ffi::{
dispatch, get_type_attr, structural_visit, structural_walk, Any, Array,
ContextPolicy,
DLDataType, DLDataTypeCode, DefRegionKind, Error, FieldGetter, Function,
Map, Object,
ObjectRefCore, Result, String as FfiString, StructuralVisitor, TypeIndex,
VisitCallbacks,
- VisitContext, VisitInterrupt, VisitValue, WalkOrder, WalkResult,
WalkWithPolicy, RUNTIME_ERROR,
+ VisitContext, VisitInterrupt, VisitValue, WalkOrder, WalkResult,
WalkWithContextPolicy,
+ RUNTIME_ERROR,
};
fn runtime_error(message: &str) -> Error {
@@ -120,7 +121,7 @@ fn
composed_policies_share_array_scope_with_visit_and_walk_callbacks() {
// 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));
+ WalkWithContextPolicy::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);
@@ -195,7 +196,7 @@ fn
policy_continuation_scopes_regions_and_restores_after_halts() {
};
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 mut walker = WalkWithContextPolicy::new(state, policies);
let result = walker.walk(&root, order);
(result, walker.into_state())
} else {
@@ -332,7 +333,7 @@ fn
policy_continuation_retargets_without_dispatching_the_container() {
];
for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
for skip in [false, true] {
- let mut walker = WalkWithPolicy::new(Probe(vec![], skip),
(Redirect, Observe));
+ let mut walker = WalkWithContextPolicy::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);
@@ -473,7 +474,7 @@ fn
policy_halts_skip_remaining_policies_and_restore_outer_state() {
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 mut walker = WalkWithContextPolicy::new(Probe::default(),
policies);
let result = walker.walk(&root, order);
(result, walker.into_state())
} else {
@@ -548,7 +549,8 @@ fn
policy_regions_compose_with_field_flags_and_function_hooks() {
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);
+ let mut walker =
+
WalkWithContextPolicy::new(PolicyRegionTrace::default(), policy);
assert!(walker.walk(&root, order).unwrap().is_none());
walker.into_state()
} else {
@@ -598,7 +600,8 @@ fn
walk_policy_preserves_reflected_pattern_before_default_descent() {
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));
+ let mut walker =
+ WalkWithContextPolicy::new(PolicyRegionTrace::default(),
Reenter(requested));
assert!(walker.walk(&root, order).unwrap().is_none());
let expected = match order {
WalkOrder::PreOrder => vec