This is an automated email from the ASF dual-hosted git repository. tlopex pushed a commit to branch rust-mutator-owned-state in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
commit b8af1b1229eb1cd5532c7e29f1588ab53b76f514 Author: tlopex <[email protected]> AuthorDate: Tue Sep 1 21:03:09 2026 -0400 [REFACTOR][RUST] Keep mutate dispatch state on the dispatcher --- docs/guides/rust_lang_guide.md | 50 +++--- rust/tvm-ffi-macros/src/dispatch.rs | 203 ++++------------------ rust/tvm-ffi/src/extra/structural_mutate.rs | 243 +++++++++++++++++++-------- rust/tvm-ffi/src/lib.rs | 6 +- rust/tvm-ffi/tests/test_structural_mutate.rs | 205 ++++++++++------------ 5 files changed, 314 insertions(+), 393 deletions(-) diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md index 3edfd737..3985939a 100644 --- a/docs/guides/rust_lang_guide.md +++ b/docs/guides/rust_lang_guide.md @@ -495,13 +495,13 @@ example, an integer handler can return `Result<i64>` to report failures and use completed before a later error are not rolled back, and the consumed root is not returned on error. -`structural_mutate` accepts typed callbacks in addition to a -`StructuralMutator`. Callbacks receive a `Mutator`; `MutateCallbacks` adds -state shared by the callback chain: +`structural_mutate` accepts typed callback chains in addition to a +`StructuralMutator`. Closure callbacks receive a `CallbackMutator`; +`MutateCallbacks` adds state shared by that callback chain: ```rust use tvm_ffi::{ - structural_mutate, Array, MapValue, MutateCallbacks, Mutator, + structural_mutate, Array, CallbackMutator, MapValue, MutateCallbacks, }; #[derive(Default)] @@ -512,11 +512,11 @@ struct Stats { let mut mutator = MutateCallbacks::new( Stats::default(), ( - |value: i64, mutator: &mut Mutator<Stats>| { + |value: i64, mutator: &mut CallbackMutator<Stats>| { mutator.state_mut().integers += 1; value + 1 }, - |_value: &MapValue, mutator: &mut Mutator<Stats>| { + |_value: &MapValue, mutator: &mut CallbackMutator<Stats>| { mutator.default_mutate() }, ), @@ -527,53 +527,49 @@ assert_eq!(mutated.iter().collect::<Vec<_>>(), vec![2, 3]); assert_eq!(mutator.state().integers, 2); ``` -`Mutator::mutate` uses the copy path for a borrowed value, while +`CallbackMutator::mutate` uses the copy path for a borrowed value, while `maybe_inplace_mutate` preserves the reuse opportunity of an owned value. -Callbacks are `Fn`; mutable data belongs in the mutator state. +Closure callbacks are `Fn`; mutable data belongs in the callback state. -`#[dispatch(mutate)]` groups typed `mutate_*` callbacks. Dispatch only selects -the first matching callback; `Mutator` supplies recursion, the current -definition region, and mutable state. `mutator.mutate(child)` inherits the -current region, while `mutate_with` is available for an explicit override. An -unmatched value follows default mutation with its current in-place permit: +`#[dispatch(mutate)]` groups typed `mutate_*` callbacks. The dispatch object +owns its mutable pass state, while `Mutator` supplies recursion and the current +definition region. `mutator.mutate(self, child)` safely reborrows that dispatch +object and inherits the current region; `mutate_with` is available for an +explicit override. An unmatched value follows default mutation with its +current in-place permit: ```rust use tvm_ffi::{ - dispatch, structural_mutate, Any, Array, MutateCallbacks, Mutator, + dispatch, structural_mutate, Array, Mutator, }; #[derive(Default)] -struct IncrementState { +struct Increment { integers: usize, } -struct Increment; - #[dispatch(mutate)] impl Increment { fn mutate_integer( - &self, + &mut self, value: i64, - mutator: &mut Mutator<IncrementState>, - ) -> Any { - mutator.state_mut().integers += 1; - Any::from(value + 1) + _mutator: &mut Mutator, + ) -> i64 { + self.integers += 1; + value + 1 } } -let mut increment = MutateCallbacks::new(IncrementState::default(), Increment); +let mut increment = Increment::default(); let mutated = structural_mutate( Array::new(vec![1_i64, 2]), &mut increment, )?; let mutated = Array::<i64>::try_from(mutated)?; assert_eq!(mutated.iter().collect::<Vec<_>>(), vec![2, 3]); -assert_eq!(increment.state().integers, 2); +assert_eq!(increment.integers, 2); ``` -A generated dispatch whose context state is `()` can be passed directly to -`structural_mutate`, without `MutateCallbacks`. - For a named custom recursion policy, implement `StructuralMutator` and pass `&mut` it to `structural_mutate`. `InplaceValue` is an engine-issued capability: callers cannot construct it from a read-only `MapValue`. Override diff --git a/rust/tvm-ffi-macros/src/dispatch.rs b/rust/tvm-ffi-macros/src/dispatch.rs index 243e4916..68cf56d4 100644 --- a/rust/tvm-ffi-macros/src/dispatch.rs +++ b/rust/tvm-ffi-macros/src/dispatch.rs @@ -19,28 +19,20 @@ use proc_macro::TokenStream; use proc_macro2::TokenStream as TokenStream2; -use quote::{format_ident, quote, quote_spanned}; +use quote::{quote, quote_spanned}; use syn::{ - parse_macro_input, FnArg, GenericArgument, ImplItem, ImplItemMethod, ItemImpl, Meta, - NestedMeta, PathArguments, Type, + parse_macro_input, FnArg, ImplItem, ImplItemMethod, ItemImpl, Meta, NestedMeta, PathArguments, + Type, }; use crate::utils::get_tvm_ffi_crate; pub(crate) fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { let args = parse_macro_input!(attr as DispatchArgs); - let mut item_impl = parse_macro_input!(item as ItemImpl); + let item_impl = parse_macro_input!(item as ItemImpl); match expand(&item_impl, args.mode) { - Ok(generated) => { - if matches!(args.mode, DispatchMode::Mutate) { - if let Err(error) = specialize_mutate_handlers(&mut item_impl) { - let error = error.to_compile_error(); - return quote!(#item_impl #error).into(); - } - } - quote!(#item_impl #generated).into() - } + Ok(generated) => quote!(#item_impl #generated).into(), Err(error) => { let error = error.to_compile_error(); quote!(#item_impl #error).into() @@ -48,81 +40,6 @@ pub(crate) fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { } } -fn specialize_mutate_handlers(item_impl: &mut ItemImpl) -> syn::Result<()> { - let tvm_ffi = get_tvm_ffi_crate(); - for item in &mut item_impl.items { - let ImplItem::Method(method) = item else { - continue; - }; - if !method.sig.ident.to_string().starts_with("mutate_") { - continue; - } - - let handler = parse_handler(method, DispatchMode::Mutate)?; - let state = handler - .mutate_state - .expect("mutate handlers always record their context state"); - let driver = format_ident!("__TvmFfiMutateDriver"); - if method - .sig - .generics - .type_params() - .any(|param| param.ident == driver) - { - return Err(syn::Error::new_spanned( - &method.sig.generics, - "reserved mutate-handler generic name is already in use", - )); - } - - method.sig.generics.params.push(syn::parse_quote!(#driver)); - method - .sig - .generics - .make_where_clause() - .predicates - .push(syn::parse_quote!( - #driver: #tvm_ffi::extra::structural_mutate::MutateContextDriver<#state> + ?Sized - )); - - let context = match method.sig.inputs.iter_mut().nth(2) { - Some(FnArg::Typed(context)) => context, - _ => unreachable!("the third mutate-handler argument cannot be a receiver"), - }; - let Type::Reference(reference) = context.ty.as_mut() else { - unreachable!("parse_handler already validated the mutate context"); - }; - let Type::Path(path) = reference.elem.as_mut() else { - unreachable!("parse_handler already validated the mutate context path"); - }; - let segment = path - .path - .segments - .last_mut() - .expect("a parsed Rust type path always has a segment"); - if matches!(segment.arguments, PathArguments::None) { - let arguments: syn::AngleBracketedGenericArguments = - syn::parse_quote!(<#state, #driver>); - segment.arguments = PathArguments::AngleBracketed(arguments); - continue; - } - let PathArguments::AngleBracketed(arguments) = &mut segment.arguments else { - unreachable!("parse_handler already rejected parenthesized arguments"); - }; - if !arguments - .args - .iter() - .any(|argument| matches!(argument, GenericArgument::Type(_))) - { - arguments.args.push(GenericArgument::Type(state.clone())); - } - arguments - .args - .push(GenericArgument::Type(syn::parse_quote!(#driver))); - } - Ok(()) -} - struct DispatchArgs { mode: DispatchMode, } @@ -252,16 +169,6 @@ fn expand(item_impl: &ItemImpl, mode: DispatchMode) -> syn::Result<TokenStream2> #tvm_ffi::extra::structural_mutate::IntoMutateResult::into_mutate_result }, }; - let mutate_state = if matches!(mode, DispatchMode::Mutate) { - Some( - handlers[0] - .mutate_state - .as_ref() - .expect("mutate handlers always record their context state"), - ) - } else { - None - }; let links = expand_links(&handlers, mode, &into_result, quote!(value)); let self_type = &item_impl.self_ty; let (impl_generics, _, where_clause) = item_impl.generics.split_for_impl(); @@ -338,35 +245,22 @@ fn expand(item_impl: &ItemImpl, mode: DispatchMode) -> syn::Result<TokenStream2> } } }, - DispatchMode::Mutate => { - let state = mutate_state.expect("mutate dispatch has a context state"); - quote! { - impl #impl_generics #tvm_ffi::extra::structural_mutate::MutateDispatch - for #self_type #where_clause - { - type State = #state; - - #[inline(always)] - #[allow(unreachable_code, unused_variables)] - fn dispatch_mutate<__TvmFfiMutateDriver>( - &self, - value: &#tvm_ffi::extra::structural_mutate::MapValue, - mutator: &mut #tvm_ffi::extra::structural_mutate::Mutator< - Self::State, - __TvmFfiMutateDriver, - >, - ) -> Option<#tvm_ffi::extra::structural_mutate::MutateResult> - where - __TvmFfiMutateDriver: - #tvm_ffi::extra::structural_mutate::MutateContextDriver<Self::State> - + ?Sized, - { - #(#links)* - None - } + DispatchMode::Mutate => quote! { + impl #impl_generics #tvm_ffi::extra::structural_mutate::MutateDispatch + for #self_type #where_clause + { + #[inline(always)] + #[allow(unreachable_code, unused_variables)] + fn dispatch_mutate( + &mut self, + value: &#tvm_ffi::extra::structural_mutate::MapValue, + mutator: &mut #tvm_ffi::extra::structural_mutate::Mutator, + ) -> Option<#tvm_ffi::extra::structural_mutate::MutateResult> { + #(#links)* + None } } - } + }, }; Ok(quote! { @@ -444,7 +338,6 @@ struct Handler { method: syn::Ident, argument: HandlerArgument, wants_def_region: bool, - mutate_state: Option<Type>, cfg_attrs: Vec<Meta>, } @@ -458,7 +351,7 @@ fn parse_handler(method: &ImplItemMethod, mode: DispatchMode) -> syn::Result<Han let inputs = &method.sig.inputs; let receiver_is_expected = match (mode, inputs.first()) { (DispatchMode::Mutate, Some(FnArg::Receiver(receiver))) => { - receiver.reference.is_some() && receiver.mutability.is_none() + receiver.reference.is_some() && receiver.mutability.is_some() } (_, Some(FnArg::Receiver(receiver))) => { receiver.reference.is_some() && receiver.mutability.is_some() @@ -472,7 +365,7 @@ fn parse_handler(method: &ImplItemMethod, mode: DispatchMode) -> syn::Result<Han }; if !receiver_is_expected || !arity_is_expected { let message = if matches!(mode, DispatchMode::Mutate) { - "mutate handlers must take `&self`, a node, and `&mut Mutator<State>`".to_owned() + "mutate handlers must take `&mut self`, a node, and `&mut Mutator`".to_owned() } else { format!( "{} handlers must take `&mut self`, a node, and optionally a trailing \ @@ -483,15 +376,13 @@ fn parse_handler(method: &ImplItemMethod, mode: DispatchMode) -> syn::Result<Han return Err(syn::Error::new_spanned(&method.sig, message)); } let wants_def_region = !matches!(mode, DispatchMode::Mutate) && inputs.len() == 3; - let mutate_state = if matches!(mode, DispatchMode::Mutate) { + if matches!(mode, DispatchMode::Mutate) { let context_type = match inputs.iter().nth(2) { Some(FnArg::Typed(context)) => context.ty.as_ref(), _ => unreachable!("the third argument cannot be a receiver"), }; - Some(parse_mutator_state(context_type)?) - } else { - None - }; + parse_mutator(context_type)?; + } let value_type = match inputs.iter().nth(1) { Some(FnArg::Typed(value)) => (*value.ty).clone(), @@ -521,16 +412,15 @@ fn parse_handler(method: &ImplItemMethod, mode: DispatchMode) -> syn::Result<Han method: method.sig.ident.clone(), argument, wants_def_region, - mutate_state, cfg_attrs, }) } -fn parse_mutator_state(context_type: &Type) -> syn::Result<Type> { +fn parse_mutator(context_type: &Type) -> syn::Result<()> { let Type::Reference(reference) = context_type else { return Err(syn::Error::new_spanned( context_type, - "the mutator must be `&mut Mutator<State>`", + "the mutator must be `&mut Mutator`", )); }; if reference.mutability.is_none() { @@ -542,55 +432,28 @@ fn parse_mutator_state(context_type: &Type) -> syn::Result<Type> { let Type::Path(path) = reference.elem.as_ref() else { return Err(syn::Error::new_spanned( context_type, - "expected `&mut Mutator<State>`", + "expected `&mut Mutator`", )); }; let Some(segment) = path.path.segments.last() else { return Err(syn::Error::new_spanned( context_type, - "expected `&mut Mutator<State>`", + "expected `&mut Mutator`", )); }; - let is_short_name = segment.ident == "Mutator"; - if !is_short_name && segment.ident != "MutateContext" { + if segment.ident != "Mutator" { return Err(syn::Error::new_spanned( context_type, - "expected `&mut Mutator<State>`", + "expected `&mut Mutator`", )); } - let arguments = match &segment.arguments { - PathArguments::AngleBracketed(arguments) => Some(arguments), - PathArguments::None if is_short_name => None, - _ => { - return Err(syn::Error::new_spanned( - context_type, - "`Mutator` accepts one optional state type", - )); - } - }; - let mut state_types = arguments.into_iter().flat_map(|arguments| { - arguments.args.iter().filter_map(|argument| match argument { - GenericArgument::Type(state) => Some(state.clone()), - _ => None, - }) - }); - let state = match state_types.next() { - Some(state) => state, - None if is_short_name => syn::parse_quote!(()), - None => { - return Err(syn::Error::new_spanned( - context_type, - "`MutateContext` requires a state type", - )); - } - }; - if state_types.next().is_some() { + if !matches!(segment.arguments, PathArguments::None) { return Err(syn::Error::new_spanned( context_type, - "`Mutator` accepts at most one state type", + "`Mutator` does not take a state type; store pass state on the dispatch object", )); } - Ok(state) + Ok(()) } fn presence_attrs(attrs: &[syn::Attribute]) -> syn::Result<Vec<Meta>> { diff --git a/rust/tvm-ffi/src/extra/structural_mutate.rs b/rust/tvm-ffi/src/extra/structural_mutate.rs index e69f312b..8e54dd8b 100644 --- a/rust/tvm-ffi/src/extra/structural_mutate.rs +++ b/rust/tvm-ffi/src/extra/structural_mutate.rs @@ -104,7 +104,7 @@ impl<T: Into<Any>> IntoMapResult for Result<T> { } } -/// State and recursive operations available to a mutation callback. +/// State and recursive operations available to a callback-chain mutation. /// /// A matched callback owns mutation of its value. Recursive operations /// reborrow the mutator, so mutable state cannot remain borrowed across them. @@ -116,15 +116,121 @@ pub struct MutateContext<'a, State, Driver: ?Sized = dyn MutateContextDriver<Sta _not_send_sync: PhantomData<Rc<()>>, } -/// Recursive mutation operations passed to structural-mutation callbacks. +/// Recursive mutation operations passed to closure callback chains. /// -/// This is the concise callback-facing name for [`MutateContext`]. The -/// borrow lifetime is inferred in function parameters, so stateful callbacks -/// can write `&mut Mutator<State>` and stateless callbacks can write -/// `&mut Mutator`. -pub type Mutator<'a, State = (), Driver = dyn MutateContextDriver<State> + 'a> = +/// Typed `#[dispatch(mutate)]` implementations use [`Mutator`] instead and +/// keep their mutable pass state directly on the dispatch object. +pub type CallbackMutator<'a, State = (), Driver = dyn MutateContextDriver<State> + 'a> = MutateContext<'a, State, Driver>; +/// Recursion control passed to a typed `#[dispatch(mutate)]` handler. +/// +/// The dispatch object owns all pass state. Recursive operations take that +/// object explicitly so Rust can safely reborrow the same `&mut self` for the +/// child call. +pub struct Mutator { + current: MapValue, + def_region_kind: DefRegionKind, + _not_send_sync: PhantomData<Rc<()>>, +} + +impl Mutator { + /// Complete borrowed value active at this callback. + #[inline(always)] + pub fn current(&self) -> &MapValue { + &self.current + } + + /// Definition-region state active at the callback's current value. + #[inline(always)] + pub fn def_region_kind(&self) -> DefRegionKind { + self.def_region_kind + } + + /// Definition region active at the callback's current value. + #[inline(always)] + pub fn region(&self) -> DefRegionKind { + self.def_region_kind + } + + /// Mutate a borrowed child through the same typed dispatch object. + #[inline(always)] + pub fn mutate<D, T>(&mut self, dispatch: &mut D, value: &T) -> Result<Any> + where + D: MutateDispatch, + for<'x> AnyView<'x>: From<&'x T>, + { + StructuralMutator::mutate(dispatch, value, self.def_region_kind) + } + + /// Mutate a borrowed child under an explicit definition-region state. + #[inline(always)] + pub fn mutate_with<D, T>( + &mut self, + dispatch: &mut D, + value: &T, + def_region_kind: DefRegionKind, + ) -> Result<Any> + where + D: MutateDispatch, + for<'x> AnyView<'x>: From<&'x T>, + { + StructuralMutator::mutate(dispatch, value, def_region_kind) + } + + /// Mutate an owned child and permit reuse when it remains uniquely owned. + #[inline(always)] + pub fn maybe_inplace_mutate<D, T>(&mut self, dispatch: &mut D, value: T) -> Result<Any> + where + D: MutateDispatch, + T: Into<Any>, + { + self.maybe_inplace_mutate_with(dispatch, value, self.def_region_kind) + } + + /// Mutate an owned child under an explicit definition-region state. + #[inline(always)] + pub fn maybe_inplace_mutate_with<D, T>( + &mut self, + dispatch: &mut D, + value: T, + def_region_kind: DefRegionKind, + ) -> Result<Any> + where + D: MutateDispatch, + T: Into<Any>, + { + StructuralMutator::maybe_inplace_mutate(dispatch, value, def_region_kind) + } + + /// Apply default mutation to the callback's current value. + #[inline(always)] + pub fn default_mutate<D: MutateDispatch>(&mut self, dispatch: &mut D) -> Result<Any> { + StructuralMutator::default_mutate(dispatch, &self.current, self.def_region_kind) + } + + /// Look up an invocation-local identity substitution. + #[inline(always)] + pub fn var_remap_get<D: MutateDispatch>( + &mut self, + dispatch: &mut D, + var: &MapValue, + ) -> Result<Option<Any>> { + StructuralMutator::var_remap_get(dispatch, var) + } + + /// Store an invocation-local identity substitution. + #[inline(always)] + pub fn var_remap_set<D: MutateDispatch>( + &mut self, + dispatch: &mut D, + var: &MapValue, + mutated_value: &Any, + ) -> Result<()> { + StructuralMutator::var_remap_set(dispatch, var, mutated_value) + } +} + #[doc(hidden)] /// Internal operations used by [`MutateContext`]. /// @@ -249,11 +355,12 @@ where /// Conversion into the mutator argument accepted by [`structural_mutate`]. /// /// Accepts a mutable low-level [`StructuralMutator`], a generated -/// [`MutateDispatch`], or a first-match callback chain. Use [`MutateCallbacks`] -/// when typed dispatch or a callback chain needs mutable state. +/// [`MutateDispatch`], or a first-match callback chain. Generated dispatch +/// objects keep mutable pass state directly on themselves; [`MutateCallbacks`] +/// remains available for closure callback chains with separate state. #[diagnostic::on_unimplemented( message = "`{Self}` is not a supported `structural_mutate` mutator", - note = "accepted mutators: `&mut U` where `U: StructuralMutator`; a generated `MutateDispatch<State = ()>`; an `Fn` callback over an FFI value type `T`, `&N` of an object node type, or `&MapValue`, followed by `&mut Mutator<State>`; or a tuple of up to 12 such callbacks (tuples may nest)", + note = "accepted mutators: `&mut U` where `U: StructuralMutator`; a generated `MutateDispatch`; an `Fn` callback over an FFI value type `T`, `&N` of an object node type, or `&MapValue`, followed by `&mut CallbackMutator<State>`; or a tuple of up to 12 such callbacks (tuples may nest)", note = "callback arguments need explicit type annotations; use `MutateCallbacks::new(state, callbacks)` for ordinary mutable callback state" )] pub trait IntoMutator<Marker> { @@ -297,10 +404,6 @@ pub type MutateResult = Result<Any>; /// Callback tuples use a type-erased mutation driver. pub enum DynamicMutateCallbacks {} -#[doc(hidden)] -/// Generated dispatch keeps the concrete mutation driver for inlining. -pub enum StaticMutateDispatch {} - /// One typed callback in a callback-driven structural mutator. pub trait MutateChainLink<State, Marker>: mutate_sealed::SealedLink<State, Marker> { #[doc(hidden)] @@ -318,19 +421,18 @@ pub trait MutateChainLink<State, Marker>: mutate_sealed::SealedLink<State, Marke /// /// `None` means no handler matched, so structural mutation applies its default /// behavior. A generated `#[dispatch(mutate)]` implementation tests -/// `mutate_*` methods in source order and passes the same [`Mutator`] to -/// the first match. +/// `mutate_*` methods in source order and passes the same [`Mutator`] to the +/// first match. The implementation owns its pass state and receives `&mut +/// self`, while [`Mutator`] only controls recursion and the definition region. pub trait MutateDispatch: Sized { - /// Mutable state shared by the dispatched callbacks. - type State; + fn dispatch_mutate(&mut self, value: &MapValue, mutator: &mut Mutator) -> Option<MutateResult>; +} - fn dispatch_mutate<Driver>( - &self, - value: &MapValue, - mutator: &mut Mutator<Self::State, Driver>, - ) -> Option<MutateResult> - where - Driver: MutateContextDriver<Self::State> + ?Sized; +impl<D: MutateDispatch> IntoMutator<ByMutateDispatch> for D { + #[inline] + fn mutate_root(mut self, root: Any) -> Result<Any> { + run_structural_mutator(root, &mut self) + } } mod mutate_sealed { @@ -364,29 +466,11 @@ mod mutate_sealed { O: IntoMutateResult, { } - - impl<D> SealedLink<D::State, super::ByMutateDispatch> for D where D: super::MutateDispatch {} } #[doc(hidden)] pub enum ByMutateDispatch {} -impl<D> MutateChainLink<D::State, ByMutateDispatch> for D -where - D: MutateDispatch, -{ - type Strategy = StaticMutateDispatch; - - #[inline(always)] - fn try_mutate( - &self, - value: &MapValue, - mutator: &mut MutateContext<'_, D::State>, - ) -> Option<MutateResult> { - self.dispatch_mutate(value, mutator) - } -} - #[doc(hidden)] pub struct ByMutateOwned<T>(PhantomData<T>); @@ -1094,9 +1178,48 @@ pub trait StructuralMutator: Sized { } } -// A plain closure has a fixed `&mut Mutator<State>` signature, while a -// macro-generated dispatch method can be generic over the concrete driver. -// Select the matching representation without changing the public callback API. +impl<D: MutateDispatch> StructuralMutator for D { + #[inline(always)] + fn dispatch_mutate(&mut self, value: &MapValue, def_region_kind: DefRegionKind) -> Result<Any> { + let mut mutator = Mutator { + current: MapValue::from_raw(value.raw()), + def_region_kind, + _not_send_sync: PhantomData, + }; + match MutateDispatch::dispatch_mutate(self, value, &mut mutator) { + Some(result) => result, + None => user_default_mutate(self, value.raw(), def_region_kind, Permit::Copy), + } + } + + #[inline(always)] + fn dispatch_maybe_inplace_mutate( + &mut self, + value: InplaceValue<'_>, + def_region_kind: DefRegionKind, + ) -> Result<Any> { + let mut mutator = Mutator { + current: MapValue::from_raw(value.raw()), + def_region_kind, + _not_send_sync: PhantomData, + }; + match MutateDispatch::dispatch_mutate(self, value.as_value(), &mut mutator) { + Some(result) => result, + None => { + let raw = value.raw(); + let permit = if object_is_unique(raw) { + Permit::MaybeInPlace + } else { + Permit::Copy + }; + user_default_mutate(self, raw, def_region_kind, permit) + } + } + } +} + +// Closure callback chains use a type-erased context driver so one concrete +// function signature can recurse through the complete chain. trait MutateCallbackStrategy<State, Link, Marker> { fn try_mutate<Driver>( driver: &mut Driver, @@ -1135,34 +1258,6 @@ where } } -impl<State, Dispatch> MutateCallbackStrategy<State, Dispatch, ByMutateDispatch> - for StaticMutateDispatch -where - Dispatch: MutateDispatch<State = State>, -{ - #[inline(always)] - fn try_mutate<Driver>( - driver: &mut Driver, - callback_ptr: *const Dispatch, - value: &MapValue, - def_region_kind: DefRegionKind, - ) -> Option<MutateResult> - where - Driver: MutateContextDriver<State>, - { - let mut mutator = MutateContext::<State, Driver> { - driver, - current: MapValue::from_raw(value.raw()), - def_region_kind, - _state: PhantomData, - _not_send_sync: PhantomData, - }; - // SAFETY: The dispatch value is held by the owning `Rc` or by the - // direct callback's stack slot and is only borrowed immutably. - unsafe { (&*callback_ptr).dispatch_mutate(value, &mut mutator) } - } -} - #[inline(always)] fn try_mutate_callbacks<State, Link, Marker, Driver>( driver: &mut Driver, diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index 67f96dcc..149f3103 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -49,9 +49,9 @@ pub use crate::error::{ }; pub use crate::extra::module::Module; pub use crate::extra::structural_mutate::{ - structural_map, structural_mutate, InplaceValue, IntoMapResult, IntoMapper, IntoMutator, - MapChainLink, MapDispatch, MapValue, MutateCallbacks, MutateChainLink, MutateContext, - MutateDispatch, Mutator, StructuralMutator, StructuralVarRemap, + structural_map, structural_mutate, CallbackMutator, InplaceValue, IntoMapResult, IntoMapper, + IntoMutator, MapChainLink, MapDispatch, MapValue, MutateCallbacks, MutateChainLink, + MutateContext, MutateDispatch, Mutator, StructuralMutator, StructuralVarRemap, }; pub use crate::extra::structural_visit::{ structural_visit, structural_walk, DefRegionKind, IntoVisitor, IntoWalkResult, IntoWalker, diff --git a/rust/tvm-ffi/tests/test_structural_mutate.rs b/rust/tvm-ffi/tests/test_structural_mutate.rs index b38c29ce..7527eb2a 100644 --- a/rust/tvm-ffi/tests/test_structural_mutate.rs +++ b/rust/tvm-ffi/tests/test_structural_mutate.rs @@ -28,10 +28,10 @@ use tvm_ffi::tvm_ffi_sys::{ TVMFFISEqHashKind, TVMFFITypeMetadata, TVMFFITypeRegisterAttr, }; use tvm_ffi::{ - dispatch, structural_map, structural_mutate, Any, AnyView, Array, DefRegionKind, Error, - Function, InplaceValue, Map, MapDispatch, MapValue, MutateCallbacks, Mutator, Object, - ObjectArc, ObjectCore, ObjectRefCore, Result, String as FfiString, StructuralMutator, - StructuralVarRemap, TypeIndex, WalkOrder, RUNTIME_ERROR, + dispatch, structural_map, structural_mutate, Any, AnyView, Array, CallbackMutator, + DefRegionKind, Error, Function, InplaceValue, Map, MapDispatch, MapValue, MutateCallbacks, + Mutator, Object, ObjectArc, ObjectCore, ObjectRefCore, Result, String as FfiString, + StructuralMutator, StructuralVarRemap, TypeIndex, WalkOrder, RUNTIME_ERROR, }; // These registration entry points are needed only to build reflected test @@ -1087,7 +1087,7 @@ fn callback_errors_preserve_message_and_add_object_context() { let error = match structural_mutate( Array::new(vec![1i64]), - |_integer: i64, _mutator: &mut Mutator| -> Result<i64> { + |_integer: i64, _mutator: &mut CallbackMutator| -> Result<i64> { Err(Error::new( RUNTIME_ERROR, "callback mutator failed", @@ -1133,9 +1133,10 @@ fn registered_mutation_hooks_receive_the_rust_mutator() { assert!(error.message().contains("retained after its active call")); let mutate_calls_before = REGISTERED_MUTATE_CALLS.load(Ordering::Relaxed); - let mutated = structural_mutate(source.clone(), |value: i64, _mutator: &mut Mutator| { - Any::from(value + 1) - }) + let mutated = structural_mutate( + source.clone(), + |value: i64, _mutator: &mut CallbackMutator| Any::from(value + 1), + ) .and_then(i64::try_from) .unwrap(); assert_eq!(mutated, 2); @@ -1376,17 +1377,15 @@ fn generated_map_dispatch_supports_kind_and_ordered_catch_all() { } #[derive(Default)] -struct GeneratedLeafState { +struct GeneratedLeafDispatch { integers: Vec<(i64, DefRegionKind)>, } -struct GeneratedLeafDispatch; - #[dispatch(mutate)] impl GeneratedLeafDispatch { - fn mutate_integer(&self, value: i64, mutator: &mut Mutator<GeneratedLeafState>) -> Any { + fn mutate_integer(&mut self, value: i64, mutator: &mut Mutator) -> Any { let region = mutator.region(); - mutator.state_mut().integers.push((value, region)); + self.integers.push((value, region)); Any::from(value + 1) } } @@ -1395,7 +1394,7 @@ struct GeneratedStatelessDispatch; #[dispatch(mutate)] impl GeneratedStatelessDispatch { - fn mutate_integer(&self, value: i64, _mutator: &mut Mutator) -> i64 { + fn mutate_integer(&mut self, value: i64, _mutator: &mut Mutator) -> i64 { value + 1 } } @@ -1414,7 +1413,7 @@ fn generated_stateless_mutate_dispatch_is_a_direct_callback() { fn generated_mutate_dispatch_defaults_unmatched_values_and_preserves_inplace_permit() { let root = Array::new(vec![1i64, 2]); let root_pointer = array_pointer(&root); - let mut mutator = MutateCallbacks::new(GeneratedLeafState::default(), GeneratedLeafDispatch); + let mut mutator = GeneratedLeafDispatch::default(); let mutated = structural_mutate(root, &mut mutator) .and_then(Array::<i64>::try_from) .unwrap(); @@ -1422,7 +1421,7 @@ fn generated_mutate_dispatch_defaults_unmatched_values_and_preserves_inplace_per assert_eq!(array_pointer(&mutated), root_pointer); assert_eq!(mutated.iter().collect::<Vec<_>>(), vec![2, 3]); assert_eq!( - mutator.state().integers, + mutator.integers, vec![(1, DefRegionKind::None), (2, DefRegionKind::None)] ); } @@ -1435,62 +1434,53 @@ fn generated_mutate_dispatch_default_remap_crosses_registered_hooks() { retained.take(); }); - let mut mutator = MutateCallbacks::new(GeneratedLeafState::default(), GeneratedLeafDispatch); + let mut mutator = GeneratedLeafDispatch::default(); let mutated = structural_mutate(rust_hook_node(), &mut mutator) .and_then(i64::try_from) .unwrap(); assert_eq!(mutated, 2); - assert_eq!(mutator.state().integers, vec![(1, DefRegionKind::None)]); + assert_eq!(mutator.integers, vec![(1, DefRegionKind::None)]); RETAINED_MUTATOR.with(|retained| { retained.take(); }); } #[derive(Default)] -struct GeneratedRecursiveState { +struct GeneratedRecursiveDispatch { arrays: Vec<DefRegionKind>, integers: Vec<(i64, DefRegionKind)>, } -struct GeneratedRecursiveDispatch; - #[dispatch(mutate)] impl GeneratedRecursiveDispatch { - fn mutate_array( - &self, - array: Array<i64>, - mutator: &mut Mutator<GeneratedRecursiveState>, - ) -> Result<Array<i64>> { + fn mutate_array(&mut self, array: Array<i64>, mutator: &mut Mutator) -> Result<Array<i64>> { let region = mutator.region(); - mutator.state_mut().arrays.push(region); + self.arrays.push(region); let mut mutated = Vec::with_capacity(array.len()); for value in array.iter() { - mutated.push(i64::try_from(mutator.mutate(&value)?)?); + mutated.push(i64::try_from(mutator.mutate(self, &value)?)?); } Ok(Array::new(mutated)) } - fn mutate_integer(&self, value: i64, mutator: &mut Mutator<GeneratedRecursiveState>) -> Any { + fn mutate_integer(&mut self, value: i64, mutator: &mut Mutator) -> Any { let region = mutator.region(); - mutator.state_mut().integers.push((value, region)); + self.integers.push((value, region)); Any::from(value + 10) } } #[test] fn generated_mutate_dispatch_recurses_through_context() { - let mut mutator = MutateCallbacks::new( - GeneratedRecursiveState::default(), - GeneratedRecursiveDispatch, - ); + let mut mutator = GeneratedRecursiveDispatch::default(); let mutated = structural_mutate(Array::new(vec![1i64, 2]), &mut mutator) .and_then(Array::<i64>::try_from) .unwrap(); assert_eq!(mutated.iter().collect::<Vec<_>>(), vec![11, 12]); - assert_eq!(mutator.state().arrays, vec![DefRegionKind::None]); + assert_eq!(mutator.arrays, vec![DefRegionKind::None]); assert_eq!( - mutator.state().integers, + mutator.integers, vec![(1, DefRegionKind::None), (2, DefRegionKind::None)] ); } @@ -1500,10 +1490,7 @@ fn generated_mutate_dispatch_inherits_region_during_explicit_recursion() { ensure_test_types_registered(); let _guard = REFLECTED_TEST_LOCK.lock().unwrap(); let root = rust_pair(Array::new(vec![1i64]), Any::new()); - let mut mutator = MutateCallbacks::new( - GeneratedRecursiveState::default(), - GeneratedRecursiveDispatch, - ); + let mut mutator = GeneratedRecursiveDispatch::default(); let mutated = structural_mutate(root, &mut mutator) .and_then(RustPair::try_from) @@ -1511,84 +1498,62 @@ fn generated_mutate_dispatch_inherits_region_during_explicit_recursion() { let first = Array::<i64>::try_from(mutated.data.first.clone()).unwrap(); assert_eq!(first.iter().collect::<Vec<_>>(), vec![11]); - assert_eq!(mutator.state().arrays, vec![DefRegionKind::Recursive]); - assert_eq!( - mutator.state().integers, - vec![(1, DefRegionKind::Recursive)] - ); + assert_eq!(mutator.arrays, vec![DefRegionKind::Recursive]); + assert_eq!(mutator.integers, vec![(1, DefRegionKind::Recursive)]); } #[derive(Default)] -struct GeneratedDefaultingState { +struct GeneratedDefaultingDispatch { arrays: usize, integers: Vec<i64>, } -struct GeneratedDefaultingDispatch; - #[dispatch(mutate)] impl GeneratedDefaultingDispatch { - fn mutate_array( - &self, - _array: Array<i64>, - mutator: &mut Mutator<GeneratedDefaultingState>, - ) -> Result<Any> { - mutator.state_mut().arrays += 1; - mutator.default_mutate() + fn mutate_array(&mut self, _array: Array<i64>, mutator: &mut Mutator) -> Result<Any> { + self.arrays += 1; + mutator.default_mutate(self) } - fn mutate_integer(&self, value: i64, mutator: &mut Mutator<GeneratedDefaultingState>) -> Any { - mutator.state_mut().integers.push(value); + fn mutate_integer(&mut self, value: i64, _mutator: &mut Mutator) -> Any { + self.integers.push(value); Any::from(value + 1) } } #[test] fn generated_mutate_dispatch_can_default_recurse_from_a_typed_handler() { - let mut mutator = MutateCallbacks::new( - GeneratedDefaultingState::default(), - GeneratedDefaultingDispatch, - ); + let mut mutator = GeneratedDefaultingDispatch::default(); let mutated = structural_mutate(Array::new(vec![1i64, 2]), &mut mutator) .and_then(Array::<i64>::try_from) .unwrap(); assert_eq!(mutated.iter().collect::<Vec<_>>(), vec![2, 3]); - assert_eq!(mutator.state().arrays, 1); - assert_eq!(mutator.state().integers, vec![1, 2]); + assert_eq!(mutator.arrays, 1); + assert_eq!(mutator.integers, vec![1, 2]); } -struct GeneratedRemappingState { +struct GeneratedRemappingDispatch { type_index: i32, calls: usize, } -struct GeneratedRemappingDispatch; - #[dispatch(mutate)] impl GeneratedRemappingDispatch { - fn mutate_dag_node( - &self, - _value: &RustDagNodeObj, - _mutator: &mut Mutator<GeneratedRemappingState>, - ) -> Any { + fn mutate_dag_node(&mut self, _value: &RustDagNodeObj, _mutator: &mut Mutator) -> Any { Any::from(42i64) } - fn mutate_any( - &self, - value: &MapValue, - mutator: &mut Mutator<GeneratedRemappingState>, - ) -> Result<Any> { - if value.type_index() != mutator.state().type_index { - return mutator.default_mutate(); + fn mutate_any(&mut self, value: &MapValue, mutator: &mut Mutator) -> Result<Any> { + if value.type_index() != self.type_index { + return mutator.default_mutate(self); } - if let Some(mutated) = mutator.var_remap_get(value)? { + if let Some(mutated) = mutator.var_remap_get(self, value)? { return Ok(mutated); } - mutator.state_mut().calls += 1; + self.calls += 1; let mutated = Any::from(41i64); - mutator.var_remap_set(value, &mutated)?; + mutator.var_remap_set(self, value, &mutated)?; Ok(mutated) } } @@ -1597,13 +1562,10 @@ impl GeneratedRemappingDispatch { fn generated_mutate_dispatch_uses_fresh_invocation_local_var_remap() { ensure_test_types_registered(); let var = rust_free_var(); - let mut mutator = MutateCallbacks::new( - GeneratedRemappingState { - type_index: RustFreeVarObj::type_index(), - calls: 0, - }, - GeneratedRemappingDispatch, - ); + let mut mutator = GeneratedRemappingDispatch { + type_index: RustFreeVarObj::type_index(), + calls: 0, + }; for expected_calls in [1, 2] { let root = call_global( @@ -1611,7 +1573,7 @@ fn generated_mutate_dispatch_uses_fresh_invocation_local_var_remap() { &[Any::from(var.clone()), Any::from(var.clone())], ); let mutated = structural_mutate(root, &mut mutator).unwrap(); - assert_eq!(mutator.state().calls, expected_calls); + assert_eq!(mutator.calls, expected_calls); assert_eq!(i64::try_from(array_item(&mutated, 0)).unwrap(), 41); assert_eq!(i64::try_from(array_item(&mutated, 1)).unwrap(), 41); } @@ -1705,7 +1667,7 @@ fn callbacks_return_values_convertible_into_any() { let mutated = structural_mutate( Array::new(vec![1i64, 2]), - |integer: i64, _mutator: &mut Mutator| integer * 2, + |integer: i64, _mutator: &mut CallbackMutator| integer * 2, ) .and_then(Array::<i64>::try_from) .unwrap(); @@ -1864,7 +1826,7 @@ fn callback_mutate_defaults_unmatched_values_and_preserves_root_permit() { ensure_test_types_registered(); let root = Array::new(vec![1i64, 2]); let root_pointer = array_pointer(&root); - let mutated = structural_mutate(root, |value: i64, _mutator: &mut Mutator| { + let mutated = structural_mutate(root, |value: i64, _mutator: &mut CallbackMutator| { Any::from(value + 1) }) .and_then(Array::<i64>::try_from) @@ -1879,14 +1841,14 @@ struct CallbackMutateStats { defaults: usize, } -fn stateful_mutate_integer(value: i64, mutator: &mut Mutator<CallbackMutateStats>) -> Any { +fn stateful_mutate_integer(value: i64, mutator: &mut CallbackMutator<CallbackMutateStats>) -> Any { mutator.state_mut().integers.push(value); Any::from(value + 1) } fn stateful_mutate_default( _value: &MapValue, - mutator: &mut Mutator<CallbackMutateStats>, + mutator: &mut CallbackMutator<CallbackMutateStats>, ) -> Result<Any> { mutator.state_mut().defaults += 1; mutator.default_mutate() @@ -1926,7 +1888,7 @@ struct CallbackMutateDepth { fn stateful_mutate_recursive( value: &MapValue, - mutator: &mut Mutator<CallbackMutateDepth>, + mutator: &mut CallbackMutator<CallbackMutateDepth>, ) -> Result<Any> { assert_eq!(mutator.current().type_index(), value.type_index()); { @@ -1971,8 +1933,8 @@ fn callback_mutate_current_default_is_repeatable_copy_path() { let mutated = structural_mutate( root, ( - |value: i64, _mutator: &mut Mutator| Any::from(value + 1), - |_value: &MapValue, mutator: &mut Mutator| -> Result<Any> { + |value: i64, _mutator: &mut CallbackMutator| Any::from(value + 1), + |_value: &MapValue, mutator: &mut CallbackMutator| -> Result<Any> { defaults.set(defaults.get() + 1); let first = mutator.default_mutate()?; let second = mutator.default_mutate()?; @@ -1995,8 +1957,8 @@ fn callback_mutate_match_is_final_and_same_fn_can_reenter() { let mutated = structural_mutate( Array::new(vec![1i64]), ( - |_array: Array<i64>, _mutator: &mut Mutator| Any::from(Array::new(vec![10i64])), - |value: i64, _mutator: &mut Mutator| { + |_array: Array<i64>, _mutator: &mut CallbackMutator| Any::from(Array::new(vec![10i64])), + |value: i64, _mutator: &mut CallbackMutator| { integer_calls.set(integer_calls.get() + 1); Any::from(value + 1) }, @@ -2010,7 +1972,7 @@ fn callback_mutate_match_is_final_and_same_fn_can_reenter() { let calls = Cell::new(0); let mutated = structural_mutate( Array::new(vec![1i64, 2]), - |_value: &MapValue, mutator: &mut Mutator| { + |_value: &MapValue, mutator: &mut CallbackMutator| { calls.set(calls.get() + 1); mutator.default_mutate() }, @@ -2034,10 +1996,10 @@ fn callback_mutate_supports_node_links_nested_tuples_and_reflection() { root, ( ( - |_value: f64, _mutator: &mut Mutator| Any::new(), - |_node: &RustDagNodeObj, _mutator: &mut Mutator| Any::from(7i64), + |_value: f64, _mutator: &mut CallbackMutator| Any::new(), + |_node: &RustDagNodeObj, _mutator: &mut CallbackMutator| Any::from(7i64), ), - |value: i64, mutator: &mut Mutator| { + |value: i64, mutator: &mut CallbackMutator| { regions.borrow_mut().push(mutator.def_region_kind()); Any::from(value + 1) }, @@ -2060,8 +2022,8 @@ fn callback_mutate_distinguishes_borrowed_and_owned_children() { let mutated = structural_mutate( true, ( - |_value: bool, mutator: &mut Mutator| mutator.mutate(&borrowed_child), - |value: i64, _mutator: &mut Mutator| Any::from(value + 1), + |_value: bool, mutator: &mut CallbackMutator| mutator.mutate(&borrowed_child), + |value: i64, _mutator: &mut CallbackMutator| Any::from(value + 1), ), ) .and_then(Array::<i64>::try_from) @@ -2074,12 +2036,12 @@ fn callback_mutate_distinguishes_borrowed_and_owned_children() { let mutated = structural_mutate( true, ( - |_value: bool, mutator: &mut Mutator| { + |_value: bool, mutator: &mut CallbackMutator| { let child = Array::new(vec![1i64]); owned_pointer.set(array_pointer(&child) as usize); mutator.maybe_inplace_mutate(child) }, - |value: i64, _mutator: &mut Mutator| Any::from(value + 1), + |value: i64, _mutator: &mut CallbackMutator| Any::from(value + 1), ), ) .and_then(Array::<i64>::try_from) @@ -2096,7 +2058,7 @@ fn callback_mutate_can_use_its_invocation_local_var_remap() { let type_index = RustFreeVarObj::type_index(); let mut mutator = MutateCallbacks::new( (), - |value: &MapValue, mutator: &mut Mutator| -> Result<Any> { + |value: &MapValue, mutator: &mut CallbackMutator| -> Result<Any> { if value.type_index() != type_index { return mutator.default_mutate(); } @@ -2124,16 +2086,19 @@ fn callback_mutate_can_use_its_invocation_local_var_remap() { #[test] fn nested_callback_mutate_restores_the_outer_active_mutator() { - let mutated = structural_mutate(1i64, |value: i64, mutator: &mut Mutator| -> Result<Any> { - if value != 1 { - return Ok(Any::from(value + 1)); - } - let inner = structural_mutate(2i64, |value: i64, _mutator: &mut Mutator| { - Any::from(value + 10) - })?; - assert_eq!(i64::try_from(inner).unwrap(), 12); - mutator.mutate(&3i64) - }) + let mutated = structural_mutate( + 1i64, + |value: i64, mutator: &mut CallbackMutator| -> Result<Any> { + if value != 1 { + return Ok(Any::from(value + 1)); + } + let inner = structural_mutate(2i64, |value: i64, _mutator: &mut CallbackMutator| { + Any::from(value + 10) + })?; + assert_eq!(i64::try_from(inner).unwrap(), 12); + mutator.mutate(&3i64) + }, + ) .and_then(i64::try_from) .unwrap(); assert_eq!(mutated, 4); @@ -2144,7 +2109,9 @@ fn callback_mutate_panics_resume_and_leave_the_next_run_usable() { let panic = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { structural_mutate( Array::new(vec![1i64]), - |_value: i64, _mutator: &mut Mutator| -> Any { panic!("callback mutator panic") }, + |_value: i64, _mutator: &mut CallbackMutator| -> Any { + panic!("callback mutator panic") + }, ) })) { Err(panic) => panic, @@ -2157,7 +2124,7 @@ fn callback_mutate_panics_resume_and_leave_the_next_run_usable() { let mutated = structural_mutate( Array::new(vec![1i64]), - |value: i64, _mutator: &mut Mutator| Any::from(value + 1), + |value: i64, _mutator: &mut CallbackMutator| Any::from(value + 1), ) .and_then(Array::<i64>::try_from) .unwrap();
