From 36ceafa9c233d781869eeeca451dd7390ab135d3 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Wed, 29 Jul 2026 19:35:38 -0400 Subject: [PATCH 01/20] support structural_visit. --- rust/README.md | 41 ++ rust/tvm-ffi-macros/Cargo.toml | 1 + rust/tvm-ffi-macros/src/lib.rs | 9 + rust/tvm-ffi-macros/src/visit.rs | 288 ++++++++ rust/tvm-ffi-sys/src/c_api.rs | 46 ++ rust/tvm-ffi/src/any.rs | 23 + rust/tvm-ffi/src/extra/mod.rs | 1 + rust/tvm-ffi/src/extra/structural.rs | 968 ++++++++++++++++++++++++++ rust/tvm-ffi/src/lib.rs | 6 +- rust/tvm-ffi/tests/test_dispatch.rs | 96 +++ rust/tvm-ffi/tests/test_structural.rs | 353 ++++++++++ 11 files changed, 1831 insertions(+), 1 deletion(-) create mode 100644 rust/tvm-ffi-macros/src/visit.rs create mode 100644 rust/tvm-ffi/src/extra/structural.rs create mode 100644 rust/tvm-ffi/tests/test_dispatch.rs create mode 100644 rust/tvm-ffi/tests/test_structural.rs diff --git a/rust/README.md b/rust/README.md index 31f3fb01b..9ae194ddd 100644 --- a/rust/README.md +++ b/rust/README.md @@ -28,6 +28,47 @@ This workspace contains three crates: The overall project focuses on low-level, direct access to the ABI when possible for maximum efficiency while maintaining interoperability. +## Structural Visitors and Walkers + +The `tvm-ffi` crate provides a native Rust structural walker over FFI values, +built-in containers, and reflected object fields. `#[dispatch(visit)]` turns +the `visit_*` methods in an inherent implementation into a typed, stateful +visitor: + +```rust +use tvm_ffi::{dispatch, structural_visit, Array, VisitCtx, WalkResult}; + +#[derive(Default)] +struct Sum { + value: i64, +} + +#[dispatch(visit)] +impl Sum { + fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.value += value; + WalkResult::Advance + } +} + +let root = Array::new(vec![1_i64, 2, 3]); +let mut sum = Sum::default(); +structural_visit(&root, &mut sum).unwrap(); +assert_eq!(sum.value, 6); +``` + +Typed handlers are tested in source order. Borrowed `ObjectCore` node types use +runtime subtype checks, owned arguments use `AnyCompatible` casts, and a final +`&VisitValue` handler acts as a catch-all. Use `structural_walk` to select +pre-order or post-order dispatch, or `walk`/`walk_with_context` for raw +callbacks. + +`WalkResult::Advance` visits reflected or container children, `Skip` suppresses +the current value's default recursion, and `Interrupt` stops the complete +traversal. A handler can visit selected children through `VisitCtx` before +returning `Skip`; definition-region state is available through the same +context. + ## Installation The Rust support depends on `libtvm_ffi`. diff --git a/rust/tvm-ffi-macros/Cargo.toml b/rust/tvm-ffi-macros/Cargo.toml index f8d29d406..c46a78e53 100644 --- a/rust/tvm-ffi-macros/Cargo.toml +++ b/rust/tvm-ffi-macros/Cargo.toml @@ -28,6 +28,7 @@ license = "Apache-2.0" proc-macro = true [dependencies] +proc-macro-crate = "3" proc-macro2 = "^1.0" quote = "^1.0" syn = { version = "1.0.48", features = ["full", "parsing", "extra-traits"] } diff --git a/rust/tvm-ffi-macros/src/lib.rs b/rust/tvm-ffi-macros/src/lib.rs index 03ecb050e..ada75d419 100644 --- a/rust/tvm-ffi-macros/src/lib.rs +++ b/rust/tvm-ffi-macros/src/lib.rs @@ -23,6 +23,15 @@ use proc_macro_error::proc_macro_error; mod match_any; mod object_macros; mod utils; +mod visit; + +/// Generate typed structural-visit dispatch from the `visit_*` methods in an +/// inherent implementation. +#[proc_macro_error] +#[proc_macro_attribute] +pub fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { + visit::dispatch(attr, item) +} /// Match object-backed values carried by an Any-compatible scrutinee. /// diff --git a/rust/tvm-ffi-macros/src/visit.rs b/rust/tvm-ffi-macros/src/visit.rs new file mode 100644 index 000000000..72e27a996 --- /dev/null +++ b/rust/tvm-ffi-macros/src/visit.rs @@ -0,0 +1,288 @@ +/* + * 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. + */ + +use proc_macro::TokenStream; +use proc_macro2::{Span, TokenStream as TokenStream2}; +use proc_macro_crate::{crate_name, FoundCrate}; +use quote::{quote, quote_spanned}; +use syn::{parse_macro_input, FnArg, ImplItem, ImplItemMethod, ItemImpl, Meta, NestedMeta, Type}; + +pub(crate) fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { + let mode = parse_macro_input!(attr as syn::Ident); + let item_impl = parse_macro_input!(item as ItemImpl); + + match expand(&mode, &item_impl) { + Ok(generated) => quote!(#item_impl #generated).into(), + Err(error) => { + let error = error.to_compile_error(); + quote!(#item_impl #error).into() + } + } +} + +fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result { + if mode != "visit" { + return Err(syn::Error::new(mode.span(), "expected `dispatch(visit)`")); + } + if item_impl.trait_.is_some() { + return Err(syn::Error::new_spanned( + item_impl, + "`dispatch(visit)` requires an inherent impl", + )); + } + + let handlers = item_impl + .items + .iter() + .filter_map(|item| match item { + ImplItem::Method(method) if method.sig.ident.to_string().starts_with("visit_") => { + Some(parse_handler(method)) + } + _ => None, + }) + .collect::>>()?; + + if handlers.is_empty() { + return Err(syn::Error::new_spanned( + item_impl, + "`dispatch(visit)` found no `visit_*` methods", + )); + } + let tvm_ffi = resolve_tvm_ffi_crate()?; + + let links = handlers.iter().map(|handler| { + let method = &handler.method; + let attrs = &handler.cfg_attrs; + let invoke = match &handler.argument { + HandlerArgument::Value => quote! { + return Some( + #tvm_ffi::extra::structural::IntoVisitResult::into_visit_result( + self.#method(value, ctx) + ) + ); + }, + HandlerArgument::BorrowedNode(node_type) => quote! { + if let Some(node) = value.as_node::<#node_type>() { + return Some( + #tvm_ffi::extra::structural::IntoVisitResult::into_visit_result( + self.#method(node, ctx) + ) + ); + } + }, + HandlerArgument::Owned(value_type) => quote! { + if let Some(node) = value.cast::<#value_type>() { + return Some( + #tvm_ffi::extra::structural::IntoVisitResult::into_visit_result( + self.#method(node, ctx) + ) + ); + } + }, + }; + quote! { + #(#[#attrs])* + { + #invoke + } + } + }); + 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 ordering_errors = handlers + .iter() + .enumerate() + .filter(|(_, handler)| matches!(&handler.argument, HandlerArgument::Value)) + .flat_map(|(index, handler)| { + handlers[index + 1..].iter().map(|later| { + let span = handler.method.span(); + let handler_attrs = &handler.cfg_attrs; + let later_attrs = &later.cfg_attrs; + quote_spanned! {span=> + #(#[#impl_cfg_attrs])* + #(#[#handler_attrs])* + #(#[#later_attrs])* + compile_error!( + "the `&VisitValue` catch-all handler must be last among enabled handlers" + ); + } + }) + }); + + Ok(quote! { + #(#ordering_errors)* + + #(#[#impl_cfg_attrs])* + impl #impl_generics #tvm_ffi::extra::structural::VisitDispatch + for #self_type #where_clause + { + #[allow(unreachable_code)] + fn dispatch_visit( + &mut self, + value: &#tvm_ffi::extra::structural::VisitValue, + ctx: &mut #tvm_ffi::extra::structural::VisitCtx<'_>, + ) -> Option<#tvm_ffi::extra::structural::VisitResult> { + #(#links)* + None + } + } + }) +} + +fn resolve_tvm_ffi_crate() -> syn::Result { + crate_name("tvm-ffi") + .map(crate_path) + .map_err(|error| syn::Error::new(Span::call_site(), error)) +} + +fn crate_path(found: FoundCrate) -> TokenStream2 { + match found { + FoundCrate::Itself => quote!(crate), + FoundCrate::Name(name) => { + let name = syn::parse_str::(&name) + .unwrap_or_else(|_| syn::Ident::new_raw(&name, Span::call_site())); + quote!(::#name) + } + } +} + +struct Handler { + method: syn::Ident, + argument: HandlerArgument, + cfg_attrs: Vec, +} + +enum HandlerArgument { + Value, + BorrowedNode(Type), + Owned(Type), +} + +fn parse_handler(method: &ImplItemMethod) -> syn::Result { + let inputs = &method.sig.inputs; + let receiver_is_mut = matches!( + inputs.first(), + Some(FnArg::Receiver(receiver)) + if receiver.reference.is_some() && receiver.mutability.is_some() + ); + if !receiver_is_mut || inputs.len() != 3 { + return Err(syn::Error::new_spanned( + &method.sig, + "visit handlers must take `&mut self`, a node, and a context", + )); + } + + let value_type = match inputs.iter().nth(1) { + Some(FnArg::Typed(value)) => (*value.ty).clone(), + _ => unreachable!("the second argument cannot be a receiver"), + }; + let argument = match &value_type { + Type::Reference(reference) if reference.mutability.is_none() => { + if is_visit_value(reference.elem.as_ref()) { + HandlerArgument::Value + } else { + HandlerArgument::BorrowedNode((*reference.elem).clone()) + } + } + Type::Reference(_) => { + return Err(syn::Error::new_spanned( + value_type, + "visit handler values cannot be mutable references", + )); + } + _ => HandlerArgument::Owned(value_type), + }; + let cfg_attrs = presence_attrs(&method.attrs)?; + Ok(Handler { + method: method.sig.ident.clone(), + argument, + cfg_attrs, + }) +} + +fn presence_attrs(attrs: &[syn::Attribute]) -> syn::Result> { + attrs + .iter() + .filter(|attr| attr.path.is_ident("cfg") || attr.path.is_ident("cfg_attr")) + .map(|attr| attr.parse_meta().map(presence_meta)) + .filter_map(Result::transpose) + .collect() +} + +fn presence_meta(meta: Meta) -> Option { + if meta.path().is_ident("cfg") { + return Some(meta); + } + let Meta::List(mut list) = meta else { + return None; + }; + if !list.path.is_ident("cfg_attr") { + return None; + } + + let mut items = list.nested.into_iter(); + let condition = items.next()?; + let mut retained = syn::punctuated::Punctuated::new(); + retained.push(condition); + for item in items { + if let NestedMeta::Meta(meta) = item { + if let Some(meta) = presence_meta(meta) { + retained.push(NestedMeta::Meta(meta)); + } + } + } + if retained.len() == 1 { + None + } else { + list.nested = retained; + Some(Meta::List(list)) + } +} + +fn is_visit_value(value_type: &Type) -> bool { + let Type::Path(path) = value_type else { + return false; + }; + path.path + .segments + .last() + .is_some_and(|segment| segment.ident == "VisitValue") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn renamed_dependency_uses_its_imported_name() { + assert_eq!( + crate_path(FoundCrate::Name("renamed_tvm_ffi".to_string())).to_string(), + ":: renamed_tvm_ffi" + ); + } + + #[test] + fn keyword_dependency_uses_a_raw_identifier() { + assert_eq!( + crate_path(FoundCrate::Name("type".to_string())).to_string(), + ":: r#type" + ); + } +} diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs index 910a0505a..7d2a235df 100644 --- a/rust/tvm-ffi-sys/src/c_api.rs +++ b/rust/tvm-ffi-sys/src/c_api.rs @@ -81,6 +81,16 @@ pub enum TVMFFITypeIndex { kTVMFFIModule = 73, /// Opaque python object. kTVMFFIOpaquePyObject = 74, + /// Mutable list object. + kTVMFFIList = 75, + /// Mutable dict object. + kTVMFFIDict = 76, + /// Structural visit interrupt object. + kTVMFFIVisitInterrupt = 77, + /// End of the statically allocated object type-index range. + kTVMFFIStaticObjectEnd = 78, + /// Start of dynamically allocated object type indices. + kTVMFFIDynObjectBegin = 128, } #[repr(i32)] @@ -91,6 +101,34 @@ pub enum TVMFFIObjectDeleterFlagBitMask { kTVMFFIObjectDeleterFlagBitMaskBoth = (1 << 0) | (1 << 1), } +/// Bit flags attached to reflected fields. +#[repr(i32)] +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum TVMFFIFieldFlagBitMask { + kTVMFFIFieldFlagBitMaskWritable = 1 << 0, + kTVMFFIFieldFlagBitMaskHasDefault = 1 << 1, + kTVMFFIFieldFlagBitMaskIsStaticMethod = 1 << 2, + kTVMFFIFieldFlagBitMaskSEqHashIgnore = 1 << 3, + kTVMFFIFieldFlagBitMaskSEqHashDefRecursive = 1 << 4, + kTVMFFIFieldFlagBitMaskDefaultFromFactory = 1 << 5, + kTVMFFIFieldFlagBitMaskReprOff = 1 << 6, + kTVMFFIFieldFlagBitMaskCompareOff = 1 << 7, + kTVMFFIFieldFlagBitMaskHashOff = 1 << 8, + kTVMFFIFieldFlagBitMaskInitOff = 1 << 9, + kTVMFFIFieldFlagBitMaskKwOnly = 1 << 10, + kTVMFFIFieldFlagBitSetterIsFunctionObj = 1 << 11, + kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive = 1 << 12, +} + +/// Definition-region mode used by structural traversal. +#[repr(i32)] +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum TVMFFIDefRegionKind { + kTVMFFIDefRegionKindNone = 0, + kTVMFFIDefRegionKindRecursive = 1, + kTVMFFIDefRegionKindNonRecursive = 2, +} + /// Handle to Object from C API's pov pub type TVMFFIObjectHandle = *mut c_void; pub type TVMFFIObjectDeleter = unsafe extern "C" fn(self_ptr: *mut c_void, flags: i32); @@ -410,6 +448,14 @@ unsafe extern "C" { pub fn TVMFFISetCustomAllocator(allocator: *mut TVMFFICustomAllocator) -> i32; pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex: *mut i32) -> i32; + pub fn TVMFFITypeRegisterAttr( + type_index: i32, + attr_name: *const TVMFFIByteArray, + attr_value: *const TVMFFIAny, + ) -> i32; + pub fn TVMFFIGetTypeAttrColumn( + attr_name: *const TVMFFIByteArray, + ) -> *const TVMFFITypeAttrColumn; pub fn TVMFFIFunctionGetGlobal( name: *const TVMFFIByteArray, out: *mut TVMFFIObjectHandle, diff --git a/rust/tvm-ffi/src/any.rs b/rust/tvm-ffi/src/any.rs index 47379f47a..1ebf017c7 100644 --- a/rust/tvm-ffi/src/any.rs +++ b/rust/tvm-ffi/src/any.rs @@ -53,6 +53,25 @@ impl<'a> AnyView<'a> { self.data.type_index } + #[inline] + pub(crate) fn as_raw_ffi_any(&self) -> &TVMFFIAny { + &self.data + } + + /// Construct a borrowed view from its ABI representation. + /// + /// # Safety + /// + /// The caller must keep every resource referenced by `data` alive for the + /// returned view's complete lifetime. + #[inline] + pub(crate) unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self { + Self { + data, + _phantom: std::marker::PhantomData, + } + } + /// More strict version than try_from/try_into /// /// This function will not try to cast the type @@ -151,6 +170,10 @@ impl Any { pub fn type_index(&self) -> i32 { self.data.type_index } + #[inline] + pub(crate) fn as_raw_ffi_any(&self) -> &TVMFFIAny { + &self.data + } /// Try to query if stored typed in Any exactly matches the type T /// /// This function is fast in the case of failure and can be used to check diff --git a/rust/tvm-ffi/src/extra/mod.rs b/rust/tvm-ffi/src/extra/mod.rs index 2a4c01d0b..cbf12bb2e 100644 --- a/rust/tvm-ffi/src/extra/mod.rs +++ b/rust/tvm-ffi/src/extra/mod.rs @@ -17,3 +17,4 @@ * under the License. */ pub mod module; +pub mod structural; diff --git a/rust/tvm-ffi/src/extra/structural.rs b/rust/tvm-ffi/src/extra/structural.rs new file mode 100644 index 000000000..5c99be379 --- /dev/null +++ b/rust/tvm-ffi/src/extra/structural.rs @@ -0,0 +1,968 @@ +/* + * 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. + */ + +//! Native Rust structural visiting. +//! +//! This module separates the two jobs involved in a visit: +//! +//! * [`VisitValue`] provides borrowed matching for generated Rust dispatch. +//! * `NativeWalker` owns recursion through containers and reflected fields. +//! +//! The runtime object registry is open, so the walker uses the stable tvm-ffi +//! reflection ABI for arbitrary registered object types. That ABI is only the +//! object-description boundary: traversal, control flow, typed dispatch, +//! visitor state, and definition-region propagation remain in Rust. +//! +//! A Rust handler may override a type's children by visiting them through +//! [`VisitCtx`] and returning [`WalkResult::Skip`]. No C++ +//! `ffi.StructuralVisitor` is constructed and no C++ default-visit function is +//! called. A non-container type with a foreign `__s_visit__` hook must be +//! handled this way; advancing into its default children is rejected instead +//! of silently substituting reflection with potentially different semantics. + +use std::ops::ControlFlow; +use std::os::raw::c_void; +use std::ptr::NonNull; + +use crate::any::{Any, AnyView}; +use crate::error::{Error, Result, RUNTIME_ERROR, TYPE_ERROR}; +use crate::function::Function; +use crate::object::ObjectCore; +use crate::tvm_ffi_sys::TVMFFIFieldFlagBitMask::{ + kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive, kTVMFFIFieldFlagBitMaskSEqHashDefRecursive, + kTVMFFIFieldFlagBitMaskSEqHashIgnore, +}; +use crate::tvm_ffi_sys::{ + TVMFFIAny, TVMFFIByteArray, TVMFFIDefRegionKind, TVMFFIFieldInfo, TVMFFIGetTypeAttrColumn, + TVMFFIGetTypeInfo, TVMFFIObject, TVMFFITypeAttrColumn, TVMFFITypeIndex, +}; + +const STRUCTURAL_VISIT_ATTR: &str = "__s_visit__"; +const FLAG_SEQ_HASH_IGNORE: i64 = kTVMFFIFieldFlagBitMaskSEqHashIgnore as i64; +const FLAG_SEQ_HASH_DEF_RECURSIVE: i64 = kTVMFFIFieldFlagBitMaskSEqHashDefRecursive as i64; +const FLAG_SEQ_HASH_DEF_NON_RECURSIVE: i64 = kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive as i64; + +/// What a callback asks the Rust walker to do with the current value. +pub enum WalkResult { + /// Continue and visit this value's children. + Advance, + /// Continue without visiting this value's children or firing its exit hook. + Skip, + /// Halt the entire traversal. + Interrupt, + /// Halt the entire traversal and return a payload to the caller. + InterruptWith(Any), +} + +impl WalkResult { + /// Halt traversal with an FFI-compatible payload. + pub fn interrupt_with>(payload: T) -> Self { + Self::InterruptWith(payload.into()) + } +} + +/// Convert either an infallible or fallible typed handler result. +/// +/// This keeps simple handlers terse while allowing a handler to return +/// `tvm_ffi::Result` and use `?`. +pub trait IntoVisitResult { + fn into_visit_result(self) -> Result; +} + +impl IntoVisitResult for WalkResult { + fn into_visit_result(self) -> Result { + Ok(self) + } +} + +impl IntoVisitResult for Result { + fn into_visit_result(self) -> Result { + self + } +} + +/// Whether a callback runs before or after a value's children. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Phase { + /// Before the value's children. + Enter, + /// After the value's children. + Exit, +} + +/// Callback order for [`structural_walk`]. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum WalkOrder { + /// Run the typed handler before the current value's children. + #[default] + PreOrder, + /// Run the typed handler after the current value's children. + PostOrder, +} + +/// Definition-region state active at the current value. +/// +/// Reflected fields marked `SEqHashDefRecursive` or +/// `SEqHashDefNonRecursive` override the inherited state for that field's +/// complete recursive visit. +#[repr(i32)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum DefRegionKind { + /// The value is outside a definition region. + #[default] + None = 0, + /// Definitions apply recursively through the visited value. + Recursive = 1, + /// Definitions apply to the visited value using non-recursive semantics. + NonRecursive = 2, +} + +const _: () = { + assert!(DefRegionKind::None as i32 == TVMFFIDefRegionKind::kTVMFFIDefRegionKindNone as i32); + assert!( + DefRegionKind::Recursive as i32 + == TVMFFIDefRegionKind::kTVMFFIDefRegionKindRecursive as i32 + ); + assert!( + DefRegionKind::NonRecursive as i32 + == TVMFFIDefRegionKind::kTVMFFIDefRegionKindNonRecursive as i32 + ); +}; + +/// Result of a completed Rust walk. +/// +/// `Continue(())` means the whole graph was visited. `Break(payload)` means a +/// handler interrupted it; a payload-less interrupt carries `ffi::None`. +pub type VisitOutcome = ControlFlow; + +/// Fallible result returned by generated typed dispatch. +#[doc(hidden)] +pub type VisitResult = Result; + +/// A borrowed view of a raw tvm-ffi value. +/// +/// Generated visitors match this value without taking ownership: borrowed +/// object-node handlers use [`VisitValue::as_node`], while POD or object-ref +/// value handlers use [`VisitValue::cast`]. +#[repr(transparent)] +pub struct VisitValue(TVMFFIAny); + +impl VisitValue { + #[inline] + fn from_raw(raw: TVMFFIAny) -> Self { + VisitValue(raw) + } + + /// Convert the value into an owned typed handle. + #[inline] + pub fn cast(&self) -> Option { + unsafe { + if R::check_any_strict(&self.0) { + Some(R::copy_from_any_view_after_check(&self.0)) + } else { + None + } + } + } + + /// Runtime type index stored in this value. + #[inline] + pub fn type_index(&self) -> i32 { + self.0.type_index + } + + /// Borrow the value as node type `N` if it is an instance of that type. + #[inline] + pub fn as_node(&self) -> Option<&N> { + if self.0.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + return None; + } + if !is_instance(self.0.type_index, N::type_index()) { + return None; + } + Some(unsafe { &*(self.0.data_union.v_obj as *const N) }) + } +} + +enum NativeHalt { + Interrupt(Any), + Error(Error), +} + +impl From for NativeHalt { + fn from(error: Error) -> Self { + NativeHalt::Error(error) + } +} + +type NativeResult = std::result::Result<(), NativeHalt>; + +/// Typed dispatch implemented by the visitor object itself. +/// +/// [`crate::dispatch`] tests the implementation's `visit_*` methods in source +/// order. Borrowed node arguments use refcount-free subtype checks, owned +/// FFI-compatible arguments use exact value casts, and `&VisitValue` is a +/// catch-all. `None` asks the Rust walker to continue normally. +pub trait VisitDispatch: Sized { + fn dispatch_visit(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) + -> Option; +} + +/// Recursive traversal access passed to a typed handler. +/// +/// The context contains the walker, not the visitor. A handler lends its +/// current `&mut self` back to [`VisitCtx::visit`], so nested traversal is an +/// ordinary checked Rust reborrow and needs no raw visitor pointer. +pub struct VisitCtx<'a> { + walker: &'a NativeWalker, + order: WalkOrder, + def_region_kind: DefRegionKind, + halted: Option, +} + +impl VisitCtx<'_> { + /// Return the definition-region state active at the current node. + pub fn def_region_kind(&self) -> DefRegionKind { + self.def_region_kind + } + + /// Visit `child` immediately with the same typed dispatcher. + pub fn visit(&mut self, visitor: &mut V, child: &T) -> bool + where + V: VisitDispatch, + for<'x> AnyView<'x>: From<&'x T>, + { + self.visit_with_def_region(visitor, child, self.def_region_kind) + } + + /// Visit `child` under an explicitly selected definition-region state. + /// + /// The override is scoped to this recursive call. The current context is + /// unchanged after success, error, or interruption. + pub fn visit_with_def_region( + &mut self, + visitor: &mut V, + child: &T, + def_region_kind: DefRegionKind, + ) -> bool + where + V: VisitDispatch, + for<'x> AnyView<'x>: From<&'x T>, + { + if self.halted.is_some() { + return false; + } + let mut dispatch = DispatchVisitor { + visitor, + order: self.order, + }; + let result = + self.walker + .visit_raw(raw_of(AnyView::from(child)), &mut dispatch, def_region_kind); + self.absorb(result) + } + + fn absorb(&mut self, result: NativeResult) -> bool { + match result { + Ok(()) => true, + Err(halt) => { + self.halted = Some(halt); + false + } + } + } +} + +trait NativeVisit { + fn order(&self) -> WalkOrder { + WalkOrder::PreOrder + } + + fn enter(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result; + + fn exit(&mut self, _value: &VisitValue, _ctx: &mut VisitCtx<'_>) -> Result { + Ok(WalkResult::Advance) + } +} + +struct DispatchVisitor<'a, V> { + visitor: &'a mut V, + order: WalkOrder, +} + +impl NativeVisit for DispatchVisitor<'_, V> { + fn order(&self) -> WalkOrder { + self.order + } + + fn enter(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + match self.order { + WalkOrder::PreOrder => self + .visitor + .dispatch_visit(value, ctx) + .unwrap_or(Ok(WalkResult::Advance)), + WalkOrder::PostOrder => Ok(WalkResult::Advance), + } + } + + fn exit(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + match self.order { + WalkOrder::PreOrder => Ok(WalkResult::Advance), + WalkOrder::PostOrder => self + .visitor + .dispatch_visit(value, ctx) + .unwrap_or(Ok(WalkResult::Advance)), + } + } +} + +struct CallbackVisitor(F); + +impl NativeVisit for CallbackVisitor +where + F: FnMut(&VisitValue, Phase, DefRegionKind) -> O, + O: IntoVisitResult, +{ + fn enter(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + (self.0)(value, Phase::Enter, ctx.def_region_kind()).into_visit_result() + } + + fn exit(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + (self.0)(value, Phase::Exit, ctx.def_region_kind()).into_visit_result() + } +} + +/// Stateless Rust recursion engine. +struct NativeWalker { + structural_visit: Option, +} + +impl NativeWalker { + fn new() -> Self { + Self { + structural_visit: type_attr_column(STRUCTURAL_VISIT_ATTR), + } + } + + fn visit_raw( + &self, + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, + ) -> NativeResult { + if value.type_index == TVMFFITypeIndex::kTVMFFINone as i32 { + return Ok(()); + } + + let visit_value = VisitValue::from_raw(value); + let mut ctx = VisitCtx { + walker: self, + order: visitor.order(), + def_region_kind, + halted: None, + }; + let enter = match visitor.enter(&visit_value, &mut ctx) { + Ok(flow) => flow, + Err(error) => return Err(Self::with_value_context(error.into(), value)), + }; + if let Some(halt) = ctx.halted.take() { + return Err(Self::with_value_context(halt, value)); + } + match enter { + WalkResult::Advance => {} + WalkResult::Skip => return Ok(()), + WalkResult::Interrupt => return Err(NativeHalt::Interrupt(Any::new())), + WalkResult::InterruptWith(payload) => return Err(NativeHalt::Interrupt(payload)), + } + + if let Err(halt) = self.visit_children_raw(value, visitor, def_region_kind) { + return Err(Self::with_value_context(halt, value)); + } + + let exit = match visitor.exit(&visit_value, &mut ctx) { + Ok(flow) => flow, + Err(error) => return Err(Self::with_value_context(error.into(), value)), + }; + if let Some(halt) = ctx.halted.take() { + return Err(Self::with_value_context(halt, value)); + } + match exit { + WalkResult::Interrupt => Err(NativeHalt::Interrupt(Any::new())), + WalkResult::InterruptWith(payload) => Err(NativeHalt::Interrupt(payload)), + WalkResult::Advance | WalkResult::Skip => Ok(()), + } + } + + fn visit_children_raw( + &self, + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, + ) -> NativeResult { + match value.type_index { + x if x == TVMFFITypeIndex::kTVMFFIArray as i32 + || x == TVMFFITypeIndex::kTVMFFIList as i32 => + { + return self.visit_sequence(value, visitor, def_region_kind); + } + x if x == TVMFFITypeIndex::kTVMFFIMap as i32 + || x == TVMFFITypeIndex::kTVMFFIDict as i32 => + { + return self.visit_map(value, visitor, def_region_kind); + } + _ => {} + } + + self.reject_foreign_structural_visit(value.type_index)?; + if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + Ok(()) + } else { + self.visit_reflected_fields(value, visitor, def_region_kind) + } + } + + fn visit_sequence( + &self, + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, + ) -> NativeResult { + let seq = unsafe { &*(value.data_union.v_obj as *const SeqPrefix) }; + if seq.size < 0 { + return Err(runtime_error("native visitor: sequence reports a negative size").into()); + } + if seq.data.is_null() && seq.size != 0 { + return Err(runtime_error( + "native visitor: non-empty sequence has a null data pointer", + ) + .into()); + } + let size = usize::try_from(seq.size) + .map_err(|_| runtime_error("native visitor: sequence size does not fit usize"))?; + if size == 0 { + return Ok(()); + } + + if value.type_index == TVMFFITypeIndex::kTVMFFIList as i32 { + // List storage may be invalidated by a re-entrant callback. Own a + // snapshot before running the first callback. + let children: Vec = { + let cells = unsafe { std::slice::from_raw_parts(seq.data, size) }; + cells + .iter() + .map(|cell| Any::from(unsafe { view_of(cell) })) + .collect() + }; + for (index, mut child) in children.into_iter().enumerate() { + let raw = raw_of_owned(&mut child); + self.visit_raw(raw, visitor, def_region_kind) + .map_err(|halt| { + with_error_context(halt, &format!("sequence item [{index}]")) + })?; + } + return Ok(()); + } + + // Array is immutable, so its element cells remain stable throughout + // recursive callbacks and need no refcounted snapshot. + let cells = unsafe { std::slice::from_raw_parts(seq.data, size) }; + for (index, child) in cells.iter().enumerate() { + self.visit_raw(*child, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("sequence item [{index}]")))?; + } + Ok(()) + } + + fn visit_map( + &self, + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, + ) -> NativeResult { + // Map storage is private C++. The Rust binding itself uses these public + // iterator functors; using them here does not invoke structural + // visiting or transfer traversal control out of Rust. + let is_dict = value.type_index == TVMFFITypeIndex::kTVMFFIDict as i32; + let (size_name, iter_name) = if is_dict { + ("ffi.DictSize", "ffi.DictForwardIterFunctor") + } else { + ("ffi.MapSize", "ffi.MapForwardIterFunctor") + }; + let size = Function::get_global(size_name)? + .call_packed(&[unsafe { view_of(&value) }]) + .and_then(i64::try_from)?; + if size < 0 { + return Err(runtime_error("native visitor: map reports a negative size").into()); + } + let size = usize::try_from(size) + .map_err(|_| runtime_error("native visitor: map size does not fit usize"))?; + if size == 0 { + return Ok(()); + } + + let iter_any = + Function::get_global(iter_name)?.call_packed(&[unsafe { view_of(&value) }])?; + let iter = Function::try_from(iter_any)?; + + if is_dict { + // Dict mutation invalidates its iterator, so snapshot all entries + // before dispatching to user code. + let mut entries = Vec::with_capacity(size); + for index in 0..size { + let key = iter.call_packed(&[AnyView::from(&0i64)])?; + let map_value = iter.call_packed(&[AnyView::from(&1i64)])?; + entries.push((key, map_value)); + if index + 1 != size { + iter.call_packed(&[AnyView::from(&2i64)])?; + } + } + + for (index, (mut key, mut map_value)) in entries.into_iter().enumerate() { + let key_raw = raw_of_owned(&mut key); + self.visit_raw(key_raw, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("dict key [{index}]")))?; + let value_raw = raw_of_owned(&mut map_value); + self.visit_raw(value_raw, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("dict value [{index}]")))?; + } + return Ok(()); + } + + // Map is immutable. Retain only the current owned key/value pair. + for index in 0..size { + let mut key = iter.call_packed(&[AnyView::from(&0i64)])?; + let mut map_value = iter.call_packed(&[AnyView::from(&1i64)])?; + let key_raw = raw_of_owned(&mut key); + self.visit_raw(key_raw, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("map key [{index}]")))?; + let value_raw = raw_of_owned(&mut map_value); + self.visit_raw(value_raw, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("map value [{index}]")))?; + if index + 1 != size { + iter.call_packed(&[AnyView::from(&2i64)])?; + } + } + Ok(()) + } + + fn visit_reflected_fields( + &self, + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, + ) -> NativeResult { + if unsafe { TVMFFIGetTypeInfo(value.type_index) }.is_null() { + return Err(runtime_error(&format!( + "native visitor: unregistered type index {}", + value.type_index + )) + .into()); + } + let object = unsafe { value.data_union.v_obj } as *mut u8; + let halted = unsafe { + for_each_field(value.type_index, |field| { + match self.visit_reflected_field(object, field, visitor, def_region_kind) { + Ok(()) => ControlFlow::Continue(()), + Err(halt) => ControlFlow::Break(halt), + } + }) + }; + halted.map_or(Ok(()), Err) + } + + unsafe fn visit_reflected_field( + &self, + object: *mut u8, + field: &TVMFFIFieldInfo, + visitor: &mut V, + inherited_region: DefRegionKind, + ) -> NativeResult { + if field.flags & FLAG_SEQ_HASH_IGNORE != 0 { + return Ok(()); + } + + let Some(getter) = field.getter else { + return Err(NativeHalt::Error(runtime_error(&format!( + "native visitor: reflected field `{}` has no getter", + field.name.as_str() + )))); + }; + let address = object.offset(field.offset as isize) as *mut c_void; + let mut child_raw = TVMFFIAny::new(); + if getter(address, &mut child_raw) != 0 { + return Err(with_error_context( + NativeHalt::Error(Error::from_raised()), + &format!("field `{}`", field.name.as_str()), + )); + } + + // A reflection getter returns an owned Any. Keep it alive while the + // recursive walk borrows its raw cell. + let mut child = Any::from_raw_ffi_any(child_raw); + let borrowed = raw_of_owned(&mut child); + let child_region = field_def_region(field, inherited_region); + self.visit_raw(borrowed, visitor, child_region) + .map_err(|halt| with_error_context(halt, &format!("field `{}`", field.name.as_str()))) + } + + fn reject_foreign_structural_visit(&self, type_index: i32) -> Result<()> { + let Some(attr) = self + .structural_visit + .and_then(|column| column.get(type_index)) + else { + return Ok(()); + }; + match attr.type_index { + x if x == TVMFFITypeIndex::kTVMFFINone as i32 => Ok(()), + x if x == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 + || x == TVMFFITypeIndex::kTVMFFIFunction as i32 => + { + let value_type = if type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + format!("type index {type_index}") + } else { + format!("type `{}`", type_key_of(type_index)) + }; + Err(runtime_error(&format!( + "native visitor: {value_type} registers foreign `{STRUCTURAL_VISIT_ATTR}`; \ + use a matching pre-order Rust handler, visit its children through \ + `VisitCtx`, and return `WalkResult::Skip`" + ))) + } + _ => Err(Error::new( + TYPE_ERROR, + &format!( + "{STRUCTURAL_VISIT_ATTR} must be an opaque function pointer or ffi.Function" + ), + "", + )), + } + } + + fn with_value_context(halt: NativeHalt, value: TVMFFIAny) -> NativeHalt { + if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + halt + } else { + with_error_context(halt, &format!("object `{}`", type_key_of(value.type_index))) + } + } +} + +/// Visit `root` in pre-order with typed handlers stored in `visitor`. +pub fn structural_visit(root: &R, visitor: &mut V) -> Result +where + V: VisitDispatch, + for<'x> AnyView<'x>: From<&'x R>, +{ + structural_walk(root, visitor, WalkOrder::PreOrder) +} + +/// Walk `root` with typed handlers and state stored in `walker`. +/// +/// `walker` may use [`crate::dispatch`] exactly like a visitor. Each matching +/// handler runs once, before or after the value's children according to +/// `order`. +pub fn structural_walk(root: &R, walker: &mut W, order: WalkOrder) -> Result +where + W: VisitDispatch, + for<'x> AnyView<'x>: From<&'x R>, +{ + let native_walker = NativeWalker::new(); + let mut dispatch = DispatchVisitor { + visitor: walker, + order, + }; + finish(native_walker.visit_raw( + raw_of(AnyView::from(root)), + &mut dispatch, + DefRegionKind::None, + )) +} + +/// Native pre/post walk used by analyses that need to observe every raw value. +pub fn walk(root: &R, mut callback: F) -> Result +where + for<'x> AnyView<'x>: From<&'x R>, + F: FnMut(&VisitValue, Phase) -> O, + O: IntoVisitResult, +{ + walk_with_context(root, move |value, phase, _def_region_kind| { + callback(value, phase) + }) +} + +/// Native pre/post walk whose callback also receives definition-region state. +pub fn walk_with_context(root: &R, callback: F) -> Result +where + for<'x> AnyView<'x>: From<&'x R>, + F: FnMut(&VisitValue, Phase, DefRegionKind) -> O, + O: IntoVisitResult, +{ + let walker = NativeWalker::new(); + let mut callback = CallbackVisitor(callback); + finish(walker.visit_raw( + raw_of(AnyView::from(root)), + &mut callback, + DefRegionKind::None, + )) +} + +fn finish(result: NativeResult) -> Result { + match result { + Ok(()) => Ok(ControlFlow::Continue(())), + Err(NativeHalt::Error(error)) => Err(error), + Err(NativeHalt::Interrupt(payload)) => Ok(ControlFlow::Break(payload)), + } +} + +fn field_def_region(field: &TVMFFIFieldInfo, inherited: DefRegionKind) -> DefRegionKind { + if field.flags & FLAG_SEQ_HASH_DEF_NON_RECURSIVE != 0 { + DefRegionKind::NonRecursive + } else if field.flags & FLAG_SEQ_HASH_DEF_RECURSIVE != 0 { + DefRegionKind::Recursive + } else { + inherited + } +} + +fn with_error_context(halt: NativeHalt, frame: &str) -> NativeHalt { + match halt { + NativeHalt::Error(error) => NativeHalt::Error(Error::with_appended_backtrace( + error, + &format!("[native structural visit] {frame}\n"), + )), + interrupt => interrupt, + } +} + +fn runtime_error(message: &str) -> Error { + Error::new(RUNTIME_ERROR, message, "") +} + +/// Layout prefix shared by the C++ `ArrayObj` and `ListObj`. +#[repr(C)] +struct SeqPrefix { + _header: TVMFFIObject, + data: *const TVMFFIAny, + size: i64, +} + +const _: () = { + assert!(std::mem::offset_of!(SeqPrefix, data) == 24); + assert!(std::mem::offset_of!(SeqPrefix, size) == 32); +}; + +#[derive(Clone, Copy)] +struct TypeAttrColumn(NonNull); + +impl TypeAttrColumn { + /// Copy one borrowed cell; ownership remains with the registry. + fn get(self, type_index: i32) -> Option { + unsafe { + let column = self.0.as_ref(); + let index = type_index - column.begin_index; + if index < 0 || index >= column.size || column.data.is_null() { + None + } else { + Some(*column.data.offset(index as isize)) + } + } + } +} + +fn type_attr_column(attr_name: &str) -> Option { + unsafe { + let attr_name = TVMFFIByteArray::from_str(attr_name); + NonNull::new(TVMFFIGetTypeAttrColumn(&attr_name).cast_mut()).map(TypeAttrColumn) + } +} + +fn type_key_of(type_index: i32) -> String { + unsafe { + let info = TVMFFIGetTypeInfo(type_index); + if info.is_null() { + format!("") + } else { + (*info).type_key.as_str().to_string() + } + } +} + +fn is_instance(object_type_index: i32, base_type_index: i32) -> bool { + if object_type_index == base_type_index { + return true; + } + unsafe { + let info = TVMFFIGetTypeInfo(object_type_index); + let base_info = TVMFFIGetTypeInfo(base_type_index); + if info.is_null() || base_info.is_null() { + return false; + } + let base_depth = (*base_info).type_depth; + if (*info).type_depth <= base_depth { + return false; + } + let ancestors = (*info).type_acenstors; + if ancestors.is_null() { + return false; + } + let ancestor = *ancestors.offset(base_depth as isize); + !ancestor.is_null() && (*ancestor).type_index == base_type_index + } +} + +/// Visit every reflected field of `type_index` and its ancestors in the same +/// parent-to-child order as C++ `ForEachFieldInfoWithEarlyStop`. +/// +/// # Safety +/// +/// `type_index` must be a registered type index. +unsafe fn for_each_field( + type_index: i32, + mut callback: impl FnMut(&'static TVMFFIFieldInfo) -> ControlFlow, +) -> Option { + let info = TVMFFIGetTypeInfo(type_index); + if info.is_null() { + return None; + } + + // Ancestor slot 0 is the root Object. C++ starts at slot 1, walks toward + // the immediate parent, then visits the concrete type's own fields. + for depth in 1..(*info).type_depth { + let ancestor = *(*info).type_acenstors.offset(depth as isize); + if let Some(value) = visit_field_level(ancestor, &mut callback) { + return Some(value); + } + } + visit_field_level(info, &mut callback) +} + +unsafe fn visit_field_level( + info: *const crate::tvm_ffi_sys::TVMFFITypeInfo, + callback: &mut impl FnMut(&'static TVMFFIFieldInfo) -> ControlFlow, +) -> Option { + if info.is_null() || (*info).fields.is_null() { + return None; + } + let fields = std::slice::from_raw_parts((*info).fields, (*info).num_fields as usize); + for field in fields { + // C reflection tables are immortal once registered. + let field: &'static TVMFFIFieldInfo = &*(field as *const TVMFFIFieldInfo); + if let ControlFlow::Break(value) = callback(field) { + return Some(value); + } + } + None +} + +fn raw_of(view: AnyView<'_>) -> TVMFFIAny { + *view.as_raw_ffi_any() +} + +fn raw_of_owned(any: &mut Any) -> TVMFFIAny { + *any.as_raw_ffi_any() +} + +unsafe fn view_of(raw: &TVMFFIAny) -> AnyView<'_> { + unsafe { AnyView::from_raw_ffi_any(*raw) } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Array; + + struct RegionProbe(Vec); + + impl NativeVisit for RegionProbe { + fn enter(&mut self, _value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + self.0.push(ctx.def_region_kind()); + Ok(WalkResult::Advance) + } + } + + #[derive(Default)] + struct TypedRegionProbe(Vec); + + #[crate::dispatch(visit)] + impl TypedRegionProbe { + fn visit_integer(&mut self, _value: i64, ctx: &mut VisitCtx<'_>) -> WalkResult { + self.0.push(ctx.def_region_kind()); + WalkResult::Advance + } + } + + unsafe extern "C" fn clone_any_field(field: *mut c_void, result: *mut TVMFFIAny) -> i32 { + let value = &*(field as *const Any); + *result = Any::into_raw_ffi_any(value.clone()); + 0 + } + + #[test] + fn def_region_is_inherited_through_containers() { + let root = Array::new(vec![1i64, 2]); + let walker = NativeWalker::new(); + let mut probe = RegionProbe(Vec::new()); + assert!(walker + .visit_raw( + raw_of(AnyView::from(&root)), + &mut probe, + DefRegionKind::Recursive, + ) + .is_ok()); + assert_eq!(probe.0, vec![DefRegionKind::Recursive; 3]); + } + + #[test] + fn reflected_field_def_region_reaches_typed_handler_and_restores() { + let walker = NativeWalker::new(); + let mut probe = TypedRegionProbe::default(); + let mut dispatch = DispatchVisitor { + visitor: &mut probe, + order: WalkOrder::PreOrder, + }; + let mut value = Any::from(7i64); + let mut field: TVMFFIFieldInfo = unsafe { std::mem::zeroed() }; + field.name = unsafe { TVMFFIByteArray::from_str("value") }; + field.getter = Some(clone_any_field); + let object = (&mut value as *mut Any).cast::(); + + for flags in [ + FLAG_SEQ_HASH_DEF_RECURSIVE, + 0, + FLAG_SEQ_HASH_DEF_NON_RECURSIVE, + FLAG_SEQ_HASH_DEF_NON_RECURSIVE | FLAG_SEQ_HASH_DEF_RECURSIVE, + FLAG_SEQ_HASH_IGNORE, + ] { + field.flags = flags; + assert!(unsafe { + walker.visit_reflected_field(object, &field, &mut dispatch, DefRegionKind::None) + } + .is_ok()); + } + assert_eq!( + probe.0, + vec![ + DefRegionKind::Recursive, + DefRegionKind::None, + DefRegionKind::NonRecursive, + DefRegionKind::NonRecursive, + ] + ); + } +} diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index 124923080..5cbe7d764 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -46,13 +46,17 @@ pub use crate::error::{ ATTRIBUTE_ERROR, INDEX_ERROR, KEY_ERROR, RUNTIME_ERROR, TYPE_ERROR, VALUE_ERROR, }; pub use crate::extra::module::Module; +pub use crate::extra::structural::{ + structural_visit, structural_walk, walk, walk_with_context, DefRegionKind, Phase, VisitCtx, + VisitDispatch, VisitOutcome, VisitValue, WalkOrder, WalkResult, +}; pub use crate::function::Function; pub use crate::object::ObjectRefCast; pub use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems, ObjectRefCore}; pub use crate::optional::Optional; pub use crate::string::{Bytes, String}; pub use crate::type_traits::AnyCompatible; -pub use tvm_ffi_macros::match_any; +pub use tvm_ffi_macros::{dispatch, match_any}; pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex; pub use tvm_ffi_sys::{ diff --git a/rust/tvm-ffi/tests/test_dispatch.rs b/rust/tvm-ffi/tests/test_dispatch.rs new file mode 100644 index 000000000..2b5869abf --- /dev/null +++ b/rust/tvm-ffi/tests/test_dispatch.rs @@ -0,0 +1,96 @@ +/* + * 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. + */ + +//! This integration test compiles as a downstream crate, checking every public +//! path emitted by the dispatch macro outside `tvm_ffi` itself. + +use tvm_ffi::{dispatch, structural_visit, Array, Object, VisitCtx, VisitDispatch, WalkResult}; + +struct ExternalCounter { + objects: usize, +} + +#[dispatch(visit)] +impl ExternalCounter { + #[cfg(any(unix, windows))] + #[cfg_attr(all(), inline)] + fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.objects += 1; + WalkResult::Advance + } +} + +fn assert_visit_dispatch() {} + +const _: fn() = assert_visit_dispatch::; + +struct CfgAttrCounter; + +#[dispatch(visit)] +impl CfgAttrCounter { + #[cfg(any())] + fn visit_disabled_catch_all( + &mut self, + _value: &tvm_ffi::VisitValue, + _ctx: &mut VisitCtx<'_>, + ) -> WalkResult { + WalkResult::Advance + } + + #[cfg_attr(all(), cfg(any()))] + fn visit_disabled(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + WalkResult::Advance + } + + fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + WalkResult::Advance + } +} + +const _: fn() = assert_visit_dispatch::; + +struct DisabledCounter; +const _: usize = std::mem::size_of::(); + +#[dispatch(visit)] +#[cfg(any())] +impl DisabledCounter { + fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + WalkResult::Advance + } +} + +struct CfgAttrDisabledCounter; +const _: usize = std::mem::size_of::(); + +#[dispatch(visit)] +#[cfg_attr(all(), cfg(any()))] +impl CfgAttrDisabledCounter { + fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + WalkResult::Advance + } +} + +#[test] +fn generated_dispatch_uses_public_downstream_paths() { + let root = Array::new(vec![1i64, 2]); + let mut visitor = ExternalCounter { objects: 0 }; + assert!(structural_visit(&root, &mut visitor).unwrap().is_continue()); + assert_eq!(visitor.objects, 1); +} diff --git a/rust/tvm-ffi/tests/test_structural.rs b/rust/tvm-ffi/tests/test_structural.rs new file mode 100644 index 000000000..6dbf90ee9 --- /dev/null +++ b/rust/tvm-ffi/tests/test_structural.rs @@ -0,0 +1,353 @@ +/* + * 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. + */ + +use std::ops::ControlFlow; + +use tvm_ffi::tvm_ffi_sys::{TVMFFIByteArray, TVMFFITypeIndex, TVMFFITypeRegisterAttr}; +use tvm_ffi::{ + dispatch, structural_visit, structural_walk, walk, walk_with_context, Any, AnyView, Array, + DefRegionKind, Error, Function, Map, Phase, Shape, String as FfiString, VisitCtx, VisitValue, + WalkOrder, WalkResult, RUNTIME_ERROR, +}; + +fn runtime_error(message: &str) -> Error { + Error::new(RUNTIME_ERROR, message, "") +} + +#[test] +fn plain_walk_uses_native_sequence_fallback() { + let root = Array::new(vec![1i64, 2, 3]); + let mut integers = 0; + assert!(walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }) + .unwrap() + .is_continue()); + assert_eq!(integers, 3); +} + +#[test] +fn plain_walk_uses_native_map_fallback() { + let root: Map = [(FfiString::from("a"), 1i64), (FfiString::from("b"), 2i64)] + .into_iter() + .collect(); + let mut integers = 0; + assert!(walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }) + .unwrap() + .is_continue()); + assert_eq!(integers, 2); +} + +struct SkipForeignShape; + +#[dispatch(visit)] +impl SkipForeignShape { + fn visit_shape(&mut self, _shape: Shape, _ctx: &mut VisitCtx<'_>) -> WalkResult { + WalkResult::Skip + } +} + +#[test] +fn foreign_structural_visit_requires_explicit_rust_override() { + let hook = Function::get_global("ffi.ArraySize").unwrap(); + let attr_name = unsafe { TVMFFIByteArray::from_str("__s_visit__") }; + let mut attr_value = Any::from(hook); + assert_eq!( + unsafe { + TVMFFITypeRegisterAttr( + TVMFFITypeIndex::kTVMFFIShape as i32, + &attr_name, + Any::as_data_ptr(&mut attr_value), + ) + }, + 0 + ); + + let root = Shape::from([2i64, 3]); + let error = match walk(&root, |_value, _phase| WalkResult::Advance) { + Err(error) => error, + Ok(_) => panic!("foreign structural visit unexpectedly used reflection"), + }; + assert!(error.message().contains("registers foreign `__s_visit__`")); + assert!(error.message().contains("return `WalkResult::Skip`")); + + assert!(structural_visit(&root, &mut SkipForeignShape) + .unwrap() + .is_continue()); +} + +#[test] +fn mutable_list_is_snapshotted_before_callbacks() { + let root = Function::get_global("ffi.List") + .unwrap() + .call_packed(&[AnyView::from(&1i64), AnyView::from(&2i64)]) + .unwrap(); + let captured = root.clone(); + let append = Function::get_global("ffi.ListAppend").unwrap(); + let mut appended = false; + let mut integers = Vec::new(); + + assert!(walk(&root, |value, phase| { + if phase == Phase::Enter { + if let Some(integer) = value.cast::() { + integers.push(integer); + if !appended { + append + .call_packed(&[AnyView::from(&captured), AnyView::from(&3i64)]) + .unwrap(); + appended = true; + } + } + } + WalkResult::Advance + }) + .unwrap() + .is_continue()); + + assert_eq!(integers, vec![1, 2]); + let size = Function::get_global("ffi.ListSize") + .unwrap() + .call_packed(&[AnyView::from(&root)]) + .and_then(i64::try_from) + .unwrap(); + assert_eq!(size, 3); +} + +#[test] +fn mutable_dict_is_snapshotted_before_callbacks() { + let root = Function::get_global("ffi.Dict") + .unwrap() + .call_packed(&[ + AnyView::from(&FfiString::from("a")), + AnyView::from(&1i64), + AnyView::from(&FfiString::from("b")), + AnyView::from(&2i64), + ]) + .unwrap(); + let captured = root.clone(); + let set_item = Function::get_global("ffi.DictSetItem").unwrap(); + let mut inserted = false; + let mut integers = Vec::new(); + + assert!(walk(&root, |value, phase| { + if phase == Phase::Enter { + if let Some(integer) = value.cast::() { + integers.push(integer); + if !inserted { + set_item + .call_packed(&[ + AnyView::from(&captured), + AnyView::from(&FfiString::from("c")), + AnyView::from(&3i64), + ]) + .unwrap(); + inserted = true; + } + } + } + WalkResult::Advance + }) + .unwrap() + .is_continue()); + + integers.sort_unstable(); + assert_eq!(integers, vec![1, 2]); + let size = Function::get_global("ffi.DictSize") + .unwrap() + .call_packed(&[AnyView::from(&root)]) + .and_then(i64::try_from) + .unwrap(); + assert_eq!(size, 3); +} + +#[test] +fn interrupt_stops_without_running_remaining_callbacks() { + let root = Array::new(vec![1i64, 2, 3]); + let mut integers = 0; + let outcome = walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::().is_some() { + integers += 1; + return WalkResult::Interrupt; + } + WalkResult::Advance + }) + .unwrap(); + assert!(outcome.is_break()); + assert_eq!(integers, 1); +} + +#[derive(Default)] +struct ManualRegionProbe { + seen: Vec, +} + +#[dispatch(visit)] +impl ManualRegionProbe { + fn visit_array(&mut self, array: Array, ctx: &mut VisitCtx<'_>) -> WalkResult { + let overridden = array.get(0).unwrap(); + if !ctx.visit_with_def_region(self, &overridden, DefRegionKind::NonRecursive) { + return WalkResult::Interrupt; + } + let inherited = array.get(1).unwrap(); + if !ctx.visit(self, &inherited) { + return WalkResult::Interrupt; + } + WalkResult::Skip + } + + fn visit_integer(&mut self, _value: i64, ctx: &mut VisitCtx<'_>) -> WalkResult { + self.seen.push(ctx.def_region_kind()); + WalkResult::Advance + } +} + +#[test] +fn manual_child_visit_can_override_def_region() { + let root = Array::new(vec![7i64, 8]); + let mut probe = ManualRegionProbe::default(); + assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + assert_eq!( + probe.seen, + vec![DefRegionKind::NonRecursive, DefRegionKind::None] + ); +} + +#[derive(Default)] +struct GenericDispatchProbe { + integers: Vec, + objects: usize, + catch_all: usize, +} + +#[dispatch(visit)] +impl GenericDispatchProbe { + fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.integers.push(value); + WalkResult::Advance + } + + fn visit_object(&mut self, _value: &tvm_ffi::Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.objects += 1; + WalkResult::Advance + } + + fn visit_any(&mut self, _value: &VisitValue, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.catch_all += 1; + WalkResult::Advance + } +} + +#[test] +fn generated_dispatch_supports_pod_and_ordered_catch_all() { + let root = Array::new(vec![1i64, 2]); + let mut probe = GenericDispatchProbe::default(); + assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + assert_eq!(probe.integers, vec![1, 2]); + assert_eq!(probe.objects, 1); + + let floats = Array::new(vec![1.0f64, 2.0]); + assert!(structural_visit(&floats, &mut probe).unwrap().is_continue()); + assert_eq!(probe.objects, 2); + assert_eq!(probe.catch_all, 2); +} + +#[derive(Default)] +struct OrderProbe { + events: Vec, +} + +#[dispatch(visit)] +impl OrderProbe { + fn visit_array(&mut self, _array: Array, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.events.push("array".to_string()); + WalkResult::Advance + } + + fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.events.push(format!("int:{value}")); + WalkResult::Advance + } +} + +#[test] +fn stateful_structural_walk_supports_post_order() { + let root = Array::new(vec![1i64, 2]); + let mut probe = OrderProbe::default(); + assert!(structural_walk(&root, &mut probe, WalkOrder::PostOrder) + .unwrap() + .is_continue()); + assert_eq!(probe.events, vec!["int:1", "int:2", "array"]); +} + +#[test] +fn interrupt_payload_is_returned_to_the_caller() { + let root = Array::new(vec![1i64, 2]); + let outcome = walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::() == Some(1) { + return WalkResult::interrupt_with(42i64); + } + WalkResult::Advance + }) + .unwrap(); + let ControlFlow::Break(payload) = outcome else { + panic!("walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(payload).unwrap(), 42); +} + +#[test] +fn handler_errors_include_native_visit_path() { + let root = Array::new(vec![1i64]); + let error = match walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::().is_some() { + Err(runtime_error("handler failed")) + } else { + Ok(WalkResult::Advance) + } + }) { + Err(error) => error, + Ok(_) => panic!("handler unexpectedly succeeded"), + }; + assert_eq!(error.message(), "handler failed"); + assert!(error.backtrace().contains("sequence item [0]")); + assert!(error.backtrace().contains("object `ffi.Array`")); +} + +#[test] +fn raw_walk_context_receives_def_region() { + let root = Array::new(vec![1i64]); + let mut regions = Vec::new(); + assert!(walk_with_context(&root, |_value, phase, region| { + if phase == Phase::Enter { + regions.push(region); + } + WalkResult::Advance + }) + .unwrap() + .is_continue()); + assert_eq!(regions, vec![DefRegionKind::None; 2]); +} From 932213a0c57eda7391dd8ac336a809aa78e3923b Mon Sep 17 00:00:00 2001 From: yuchuan Date: Wed, 29 Jul 2026 20:08:12 -0400 Subject: [PATCH 02/20] finish. --- rust/tvm-ffi-macros/src/visit.rs | 14 +++++++------- rust/tvm-ffi/src/extra/mod.rs | 2 +- .../extra/{structural.rs => structural_visit.rs} | 0 rust/tvm-ffi/src/lib.rs | 2 +- ...test_structural.rs => test_structural_visit.rs} | 0 5 files changed, 9 insertions(+), 9 deletions(-) rename rust/tvm-ffi/src/extra/{structural.rs => structural_visit.rs} (100%) rename rust/tvm-ffi/tests/{test_structural.rs => test_structural_visit.rs} (100%) diff --git a/rust/tvm-ffi-macros/src/visit.rs b/rust/tvm-ffi-macros/src/visit.rs index 72e27a996..13047d7ef 100644 --- a/rust/tvm-ffi-macros/src/visit.rs +++ b/rust/tvm-ffi-macros/src/visit.rs @@ -72,7 +72,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result let invoke = match &handler.argument { HandlerArgument::Value => quote! { return Some( - #tvm_ffi::extra::structural::IntoVisitResult::into_visit_result( + #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( self.#method(value, ctx) ) ); @@ -80,7 +80,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result HandlerArgument::BorrowedNode(node_type) => quote! { if let Some(node) = value.as_node::<#node_type>() { return Some( - #tvm_ffi::extra::structural::IntoVisitResult::into_visit_result( + #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( self.#method(node, ctx) ) ); @@ -89,7 +89,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result HandlerArgument::Owned(value_type) => quote! { if let Some(node) = value.cast::<#value_type>() { return Some( - #tvm_ffi::extra::structural::IntoVisitResult::into_visit_result( + #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( self.#method(node, ctx) ) ); @@ -130,15 +130,15 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result #(#ordering_errors)* #(#[#impl_cfg_attrs])* - impl #impl_generics #tvm_ffi::extra::structural::VisitDispatch + impl #impl_generics #tvm_ffi::extra::structural_visit::VisitDispatch for #self_type #where_clause { #[allow(unreachable_code)] fn dispatch_visit( &mut self, - value: &#tvm_ffi::extra::structural::VisitValue, - ctx: &mut #tvm_ffi::extra::structural::VisitCtx<'_>, - ) -> Option<#tvm_ffi::extra::structural::VisitResult> { + value: &#tvm_ffi::extra::structural_visit::VisitValue, + ctx: &mut #tvm_ffi::extra::structural_visit::VisitCtx<'_>, + ) -> Option<#tvm_ffi::extra::structural_visit::VisitResult> { #(#links)* None } diff --git a/rust/tvm-ffi/src/extra/mod.rs b/rust/tvm-ffi/src/extra/mod.rs index cbf12bb2e..32fed9fa3 100644 --- a/rust/tvm-ffi/src/extra/mod.rs +++ b/rust/tvm-ffi/src/extra/mod.rs @@ -17,4 +17,4 @@ * under the License. */ pub mod module; -pub mod structural; +pub mod structural_visit; diff --git a/rust/tvm-ffi/src/extra/structural.rs b/rust/tvm-ffi/src/extra/structural_visit.rs similarity index 100% rename from rust/tvm-ffi/src/extra/structural.rs rename to rust/tvm-ffi/src/extra/structural_visit.rs diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index 5cbe7d764..7990bf891 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -46,7 +46,7 @@ pub use crate::error::{ ATTRIBUTE_ERROR, INDEX_ERROR, KEY_ERROR, RUNTIME_ERROR, TYPE_ERROR, VALUE_ERROR, }; pub use crate::extra::module::Module; -pub use crate::extra::structural::{ +pub use crate::extra::structural_visit::{ structural_visit, structural_walk, walk, walk_with_context, DefRegionKind, Phase, VisitCtx, VisitDispatch, VisitOutcome, VisitValue, WalkOrder, WalkResult, }; diff --git a/rust/tvm-ffi/tests/test_structural.rs b/rust/tvm-ffi/tests/test_structural_visit.rs similarity index 100% rename from rust/tvm-ffi/tests/test_structural.rs rename to rust/tvm-ffi/tests/test_structural_visit.rs From 9adb020d4450fa094c9aa0814cd8dcd2cbc4ad7e Mon Sep 17 00:00:00 2001 From: yuchuan Date: Wed, 29 Jul 2026 20:52:17 -0400 Subject: [PATCH 03/20] update readme. --- rust/README.md | 33 ++++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/rust/README.md b/rust/README.md index 9ae194ddd..e65e9d1a6 100644 --- a/rust/README.md +++ b/rust/README.md @@ -33,31 +33,42 @@ efficiency while maintaining interoperability. The `tvm-ffi` crate provides a native Rust structural walker over FFI values, built-in containers, and reflected object fields. `#[dispatch(visit)]` turns the `visit_*` methods in an inherent implementation into a typed, stateful -visitor: +visitor. Each value is automatically dispatched to the handler matching its +runtime type: ```rust -use tvm_ffi::{dispatch, structural_visit, Array, VisitCtx, WalkResult}; +use tvm_ffi::{dispatch, structural_visit, Function, VisitCtx, WalkResult}; #[derive(Default)] -struct Sum { - value: i64, +struct Calculator { + value: f64, } #[dispatch(visit)] -impl Sum { +impl Calculator { fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { - self.value += value; + self.value += value as f64; + WalkResult::Advance + } + + fn visit_float(&mut self, value: f64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + self.value -= value; WalkResult::Advance } } -let root = Array::new(vec![1_i64, 2, 3]); -let mut sum = Sum::default(); -structural_visit(&root, &mut sum).unwrap(); -assert_eq!(sum.value, 6); +let values = Function::get_global("ffi.Array") + .unwrap() + .call_tuple((10_i64, 2.5_f64)) + .unwrap(); +let mut calculator = Calculator::default(); +assert!(structural_visit(&values, &mut calculator) + .unwrap() + .is_continue()); +assert_eq!(calculator.value, 7.5); ``` -Typed handlers are tested in source order. Borrowed `ObjectCore` node types use +Typed handlers are tested in source order: borrowed `ObjectCore` node types use runtime subtype checks, owned arguments use `AnyCompatible` casts, and a final `&VisitValue` handler acts as a catch-all. Use `structural_walk` to select pre-order or post-order dispatch, or `walk`/`walk_with_context` for raw From 12fcdb0c35a05e833c72b41cb9935a5d5ce331d3 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 30 Jul 2026 11:37:52 -0400 Subject: [PATCH 04/20] remove ctx. --- rust/README.md | 38 +++- rust/tvm-ffi-macros/src/visit.rs | 47 ++++- rust/tvm-ffi-sys/src/c_api.rs | 12 ++ rust/tvm-ffi/src/extra/structural_visit.rs | 192 ++++++++++---------- rust/tvm-ffi/src/lib.rs | 2 +- rust/tvm-ffi/tests/test_dispatch.rs | 20 +- rust/tvm-ffi/tests/test_structural_visit.rs | 34 ++-- 7 files changed, 205 insertions(+), 140 deletions(-) diff --git a/rust/README.md b/rust/README.md index e65e9d1a6..bf669edca 100644 --- a/rust/README.md +++ b/rust/README.md @@ -37,7 +37,7 @@ visitor. Each value is automatically dispatched to the handler matching its runtime type: ```rust -use tvm_ffi::{dispatch, structural_visit, Function, VisitCtx, WalkResult}; +use tvm_ffi::{dispatch, structural_visit, Function, WalkResult}; #[derive(Default)] struct Calculator { @@ -46,12 +46,12 @@ struct Calculator { #[dispatch(visit)] impl Calculator { - fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_integer(&mut self, value: i64) -> WalkResult { self.value += value as f64; WalkResult::Advance } - fn visit_float(&mut self, value: f64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_float(&mut self, value: f64) -> WalkResult { self.value -= value; WalkResult::Advance } @@ -76,9 +76,35 @@ callbacks. `WalkResult::Advance` visits reflected or container children, `Skip` suppresses the current value's default recursion, and `Interrupt` stops the complete -traversal. A handler can visit selected children through `VisitCtx` before -returning `Skip`; definition-region state is available through the same -context. +traversal. + +A handler may declare an optional third `DefRegionKind` parameter to observe +the definition-region state at the current value (maintained automatically +from reflected field flags): + +```rust,ignore +fn visit_var(&mut self, var: &VarObj, kind: DefRegionKind) -> WalkResult { + // kind tells a definition position apart from a use position. + WalkResult::Advance +} +``` + +To take over a value's children, a pre-order handler visits them through +`VisitDispatch::subvisit` and returns `Skip`. `subvisit` reports a nested +interrupt as `ControlFlow::Break`; propagate it (and errors, via `?`) instead +of dropping the result: + +```rust,ignore +fn visit_func(&mut self, func: &FuncObj, kind: DefRegionKind) -> Result { + if let ControlFlow::Break(payload) = self.subvisit(&func.params, DefRegionKind::Recursive)? { + return Ok(WalkResult::InterruptWith(payload)); + } + if let ControlFlow::Break(payload) = self.subvisit(&func.body, kind)? { + return Ok(WalkResult::InterruptWith(payload)); + } + Ok(WalkResult::Skip) +} +``` ## Installation diff --git a/rust/tvm-ffi-macros/src/visit.rs b/rust/tvm-ffi-macros/src/visit.rs index 13047d7ef..54bec9dff 100644 --- a/rust/tvm-ffi-macros/src/visit.rs +++ b/rust/tvm-ffi-macros/src/visit.rs @@ -69,11 +69,16 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result let links = handlers.iter().map(|handler| { let method = &handler.method; let attrs = &handler.cfg_attrs; + let kind_arg = if handler.takes_def_region_kind { + quote!(, def_region_kind) + } else { + quote!() + }; let invoke = match &handler.argument { HandlerArgument::Value => quote! { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(value, ctx) + self.#method(value #kind_arg) ) ); }, @@ -81,7 +86,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result if let Some(node) = value.as_node::<#node_type>() { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(node, ctx) + self.#method(node #kind_arg) ) ); } @@ -90,7 +95,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result if let Some(node) = value.cast::<#value_type>() { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(node, ctx) + self.#method(node #kind_arg) ) ); } @@ -133,11 +138,11 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result impl #impl_generics #tvm_ffi::extra::structural_visit::VisitDispatch for #self_type #where_clause { - #[allow(unreachable_code)] + #[allow(unreachable_code, unused_variables)] fn dispatch_visit( &mut self, value: &#tvm_ffi::extra::structural_visit::VisitValue, - ctx: &mut #tvm_ffi::extra::structural_visit::VisitCtx<'_>, + def_region_kind: #tvm_ffi::extra::structural_visit::DefRegionKind, ) -> Option<#tvm_ffi::extra::structural_visit::VisitResult> { #(#links)* None @@ -166,6 +171,7 @@ fn crate_path(found: FoundCrate) -> TokenStream2 { struct Handler { method: syn::Ident, argument: HandlerArgument, + takes_def_region_kind: bool, cfg_attrs: Vec, } @@ -182,13 +188,29 @@ fn parse_handler(method: &ImplItemMethod) -> syn::Result { Some(FnArg::Receiver(receiver)) if receiver.reference.is_some() && receiver.mutability.is_some() ); - if !receiver_is_mut || inputs.len() != 3 { + if !receiver_is_mut || (inputs.len() != 2 && inputs.len() != 3) { return Err(syn::Error::new_spanned( &method.sig, - "visit handlers must take `&mut self`, a node, and a context", + "visit handlers must take `&mut self`, a node, and optionally a `DefRegionKind`", )); } + let takes_def_region_kind = if inputs.len() == 3 { + let kind_type = match inputs.iter().nth(2) { + Some(FnArg::Typed(kind)) => &kind.ty, + _ => unreachable!("the third argument cannot be a receiver"), + }; + if !is_def_region_kind(kind_type) { + return Err(syn::Error::new_spanned( + kind_type, + "the third visit handler argument must be `DefRegionKind` (by value)", + )); + } + true + } else { + false + }; + let value_type = match inputs.iter().nth(1) { Some(FnArg::Typed(value)) => (*value.ty).clone(), _ => unreachable!("the second argument cannot be a receiver"), @@ -213,6 +235,7 @@ fn parse_handler(method: &ImplItemMethod) -> syn::Result { Ok(Handler { method: method.sig.ident.clone(), argument, + takes_def_region_kind, cfg_attrs, }) } @@ -266,6 +289,16 @@ fn is_visit_value(value_type: &Type) -> bool { .is_some_and(|segment| segment.ident == "VisitValue") } +fn is_def_region_kind(kind_type: &Type) -> bool { + let Type::Path(path) = kind_type else { + return false; + }; + path.path + .segments + .last() + .is_some_and(|segment| segment.ident == "DefRegionKind" && segment.arguments.is_empty()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs index 7d2a235df..203446c89 100644 --- a/rust/tvm-ffi-sys/src/c_api.rs +++ b/rust/tvm-ffi-sys/src/c_api.rs @@ -129,6 +129,18 @@ pub enum TVMFFIDefRegionKind { kTVMFFIDefRegionKindNonRecursive = 2, } +/// Structural equality/hash participation kind stored in type metadata. +#[repr(i32)] +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum TVMFFISEqHashKind { + kTVMFFISEqHashKindUnsupported = 0, + kTVMFFISEqHashKindTreeNode = 1, + kTVMFFISEqHashKindFreeVar = 2, + kTVMFFISEqHashKindDAGNode = 3, + kTVMFFISEqHashKindConstTreeNode = 4, + kTVMFFISEqHashKindUniqueInstance = 5, +} + /// Handle to Object from C API's pov pub type TVMFFIObjectHandle = *mut c_void; pub type TVMFFIObjectDeleter = unsafe extern "C" fn(self_ptr: *mut c_void, flags: i32); diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index 5c99be379..a54a4388d 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -30,7 +30,7 @@ //! visitor state, and definition-region propagation remain in Rust. //! //! A Rust handler may override a type's children by visiting them through -//! [`VisitCtx`] and returning [`WalkResult::Skip`]. No C++ +//! [`VisitDispatch::subvisit`] and returning [`WalkResult::Skip`]. No C++ //! `ffi.StructuralVisitor` is constructed and no C++ default-visit function is //! called. A non-container type with a foreign `__s_visit__` hook must be //! handled this way; advancing into its default children is rejected instead @@ -50,7 +50,7 @@ use crate::tvm_ffi_sys::TVMFFIFieldFlagBitMask::{ }; use crate::tvm_ffi_sys::{ TVMFFIAny, TVMFFIByteArray, TVMFFIDefRegionKind, TVMFFIFieldInfo, TVMFFIGetTypeAttrColumn, - TVMFFIGetTypeInfo, TVMFFIObject, TVMFFITypeAttrColumn, TVMFFITypeIndex, + TVMFFIGetTypeInfo, TVMFFIObject, TVMFFISEqHashKind, TVMFFITypeAttrColumn, TVMFFITypeIndex, }; const STRUCTURAL_VISIT_ATTR: &str = "__s_visit__"; @@ -220,83 +220,41 @@ type NativeResult = std::result::Result<(), NativeHalt>; /// FFI-compatible arguments use exact value casts, and `&VisitValue` is a /// catch-all. `None` asks the Rust walker to continue normally. pub trait VisitDispatch: Sized { - fn dispatch_visit(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) - -> Option; -} - -/// Recursive traversal access passed to a typed handler. -/// -/// The context contains the walker, not the visitor. A handler lends its -/// current `&mut self` back to [`VisitCtx::visit`], so nested traversal is an -/// ordinary checked Rust reborrow and needs no raw visitor pointer. -pub struct VisitCtx<'a> { - walker: &'a NativeWalker, - order: WalkOrder, - def_region_kind: DefRegionKind, - halted: Option, -} - -impl VisitCtx<'_> { - /// Return the definition-region state active at the current node. - pub fn def_region_kind(&self) -> DefRegionKind { - self.def_region_kind - } - - /// Visit `child` immediately with the same typed dispatcher. - pub fn visit(&mut self, visitor: &mut V, child: &T) -> bool - where - V: VisitDispatch, - for<'x> AnyView<'x>: From<&'x T>, - { - self.visit_with_def_region(visitor, child, self.def_region_kind) - } - - /// Visit `child` under an explicitly selected definition-region state. - /// - /// The override is scoped to this recursive call. The current context is - /// unchanged after success, error, or interruption. - pub fn visit_with_def_region( + fn dispatch_visit( &mut self, - visitor: &mut V, - child: &T, + value: &VisitValue, def_region_kind: DefRegionKind, - ) -> bool + ) -> Option; + + /// Visit `child` immediately with this dispatcher under `def_region_kind`. + /// + /// This is how a pre-order handler takes over a value's children: visit + /// each selected child, then return [`WalkResult::Skip`]. Pass the + /// handler's received `def_region_kind` to keep the inherited + /// definition-region state, or another kind to override it for exactly + /// this subtree. The nested traversal always dispatches pre-order. + /// + /// `Err` carries a nested handler failure and should be propagated with + /// `?`. `Ok(ControlFlow::Break(payload))` reports a nested interrupt; + /// return [`WalkResult::InterruptWith`] with the payload to keep + /// halting. Dropping the result silently swallows both. + fn subvisit(&mut self, child: &T, def_region_kind: DefRegionKind) -> Result where - V: VisitDispatch, for<'x> AnyView<'x>: From<&'x T>, { - if self.halted.is_some() { - return false; - } + let walker = NativeWalker::new(); let mut dispatch = DispatchVisitor { - visitor, - order: self.order, + visitor: self, + order: WalkOrder::PreOrder, }; - let result = - self.walker - .visit_raw(raw_of(AnyView::from(child)), &mut dispatch, def_region_kind); - self.absorb(result) - } - - fn absorb(&mut self, result: NativeResult) -> bool { - match result { - Ok(()) => true, - Err(halt) => { - self.halted = Some(halt); - false - } - } + finish(walker.visit_raw(raw_of(AnyView::from(child)), &mut dispatch, def_region_kind)) } } trait NativeVisit { - fn order(&self) -> WalkOrder { - WalkOrder::PreOrder - } - - fn enter(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result; + fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result; - fn exit(&mut self, _value: &VisitValue, _ctx: &mut VisitCtx<'_>) -> Result { + fn exit(&mut self, _value: &VisitValue, _def_region_kind: DefRegionKind) -> Result { Ok(WalkResult::Advance) } } @@ -307,26 +265,22 @@ struct DispatchVisitor<'a, V> { } impl NativeVisit for DispatchVisitor<'_, V> { - fn order(&self) -> WalkOrder { - self.order - } - - fn enter(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { match self.order { WalkOrder::PreOrder => self .visitor - .dispatch_visit(value, ctx) + .dispatch_visit(value, def_region_kind) .unwrap_or(Ok(WalkResult::Advance)), WalkOrder::PostOrder => Ok(WalkResult::Advance), } } - fn exit(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { + fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { match self.order { WalkOrder::PreOrder => Ok(WalkResult::Advance), WalkOrder::PostOrder => self .visitor - .dispatch_visit(value, ctx) + .dispatch_visit(value, def_region_kind) .unwrap_or(Ok(WalkResult::Advance)), } } @@ -339,12 +293,12 @@ where F: FnMut(&VisitValue, Phase, DefRegionKind) -> O, O: IntoVisitResult, { - fn enter(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { - (self.0)(value, Phase::Enter, ctx.def_region_kind()).into_visit_result() + fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { + (self.0)(value, Phase::Enter, def_region_kind).into_visit_result() } - fn exit(&mut self, value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { - (self.0)(value, Phase::Exit, ctx.def_region_kind()).into_visit_result() + fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { + (self.0)(value, Phase::Exit, def_region_kind).into_visit_result() } } @@ -371,19 +325,10 @@ impl NativeWalker { } let visit_value = VisitValue::from_raw(value); - let mut ctx = VisitCtx { - walker: self, - order: visitor.order(), - def_region_kind, - halted: None, - }; - let enter = match visitor.enter(&visit_value, &mut ctx) { + let enter = match visitor.enter(&visit_value, def_region_kind) { Ok(flow) => flow, Err(error) => return Err(Self::with_value_context(error.into(), value)), }; - if let Some(halt) = ctx.halted.take() { - return Err(Self::with_value_context(halt, value)); - } match enter { WalkResult::Advance => {} WalkResult::Skip => return Ok(()), @@ -395,13 +340,10 @@ impl NativeWalker { return Err(Self::with_value_context(halt, value)); } - let exit = match visitor.exit(&visit_value, &mut ctx) { + let exit = match visitor.exit(&visit_value, def_region_kind) { Ok(flow) => flow, Err(error) => return Err(Self::with_value_context(error.into(), value)), }; - if let Some(halt) = ctx.halted.take() { - return Err(Self::with_value_context(halt, value)); - } match exit { WalkResult::Interrupt => Err(NativeHalt::Interrupt(Any::new())), WalkResult::InterruptWith(payload) => Err(NativeHalt::Interrupt(payload)), @@ -567,13 +509,23 @@ impl NativeWalker { visitor: &mut V, def_region_kind: DefRegionKind, ) -> NativeResult { - if unsafe { TVMFFIGetTypeInfo(value.type_index) }.is_null() { + let type_info = unsafe { TVMFFIGetTypeInfo(value.type_index) }; + if type_info.is_null() { return Err(runtime_error(&format!( "native visitor: unregistered type index {}", value.type_index )) .into()); } + let seq_hash_kind = unsafe { + let metadata = (*type_info).metadata; + if metadata.is_null() { + TVMFFISEqHashKind::kTVMFFISEqHashKindUnsupported as i32 + } else { + (*metadata).structural_eq_hash_kind + } + }; + let def_region_kind = free_var_child_region(def_region_kind, seq_hash_kind); let object = unsafe { value.data_union.v_obj } as *mut u8; let halted = unsafe { for_each_field(value.type_index, |field| { @@ -641,7 +593,7 @@ impl NativeWalker { Err(runtime_error(&format!( "native visitor: {value_type} registers foreign `{STRUCTURAL_VISIT_ATTR}`; \ use a matching pre-order Rust handler, visit its children through \ - `VisitCtx`, and return `WalkResult::Skip`" + `VisitDispatch::subvisit`, and return `WalkResult::Skip`" ))) } _ => Err(Error::new( @@ -740,6 +692,20 @@ fn field_def_region(field: &TVMFFIFieldInfo, inherited: DefRegionKind) -> DefReg } } +/// A non-recursive definition applies to a FreeVar value itself, but not to +/// the FreeVar's own reflected children: nested free vars there must resolve +/// against an outer binding instead of rebinding. Mirrors C++ +/// `VisitReflectedFieldsExpected`. +fn free_var_child_region(inherited: DefRegionKind, structural_eq_hash_kind: i32) -> DefRegionKind { + if inherited == DefRegionKind::NonRecursive + && structural_eq_hash_kind == TVMFFISEqHashKind::kTVMFFISEqHashKindFreeVar as i32 + { + DefRegionKind::None + } else { + inherited + } +} + fn with_error_context(halt: NativeHalt, frame: &str) -> NativeHalt { match halt { NativeHalt::Error(error) => NativeHalt::Error(Error::with_appended_backtrace( @@ -890,8 +856,12 @@ mod tests { struct RegionProbe(Vec); impl NativeVisit for RegionProbe { - fn enter(&mut self, _value: &VisitValue, ctx: &mut VisitCtx<'_>) -> Result { - self.0.push(ctx.def_region_kind()); + fn enter( + &mut self, + _value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result { + self.0.push(def_region_kind); Ok(WalkResult::Advance) } } @@ -901,8 +871,8 @@ mod tests { #[crate::dispatch(visit)] impl TypedRegionProbe { - fn visit_integer(&mut self, _value: i64, ctx: &mut VisitCtx<'_>) -> WalkResult { - self.0.push(ctx.def_region_kind()); + fn visit_integer(&mut self, _value: i64, def_region_kind: DefRegionKind) -> WalkResult { + self.0.push(def_region_kind); WalkResult::Advance } } @@ -965,4 +935,28 @@ mod tests { ] ); } + + #[test] + fn non_recursive_region_is_clamped_for_free_var_children_only() { + use TVMFFISEqHashKind::{kTVMFFISEqHashKindFreeVar, kTVMFFISEqHashKindTreeNode}; + + let free_var = kTVMFFISEqHashKindFreeVar as i32; + let tree_node = kTVMFFISEqHashKindTreeNode as i32; + assert_eq!( + free_var_child_region(DefRegionKind::NonRecursive, free_var), + DefRegionKind::None + ); + assert_eq!( + free_var_child_region(DefRegionKind::Recursive, free_var), + DefRegionKind::Recursive + ); + assert_eq!( + free_var_child_region(DefRegionKind::None, free_var), + DefRegionKind::None + ); + assert_eq!( + free_var_child_region(DefRegionKind::NonRecursive, tree_node), + DefRegionKind::NonRecursive + ); + } } diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index 7990bf891..987bb9480 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -47,7 +47,7 @@ pub use crate::error::{ }; pub use crate::extra::module::Module; pub use crate::extra::structural_visit::{ - structural_visit, structural_walk, walk, walk_with_context, DefRegionKind, Phase, VisitCtx, + structural_visit, structural_walk, walk, walk_with_context, DefRegionKind, Phase, VisitDispatch, VisitOutcome, VisitValue, WalkOrder, WalkResult, }; pub use crate::function::Function; diff --git a/rust/tvm-ffi/tests/test_dispatch.rs b/rust/tvm-ffi/tests/test_dispatch.rs index 2b5869abf..4370dcd42 100644 --- a/rust/tvm-ffi/tests/test_dispatch.rs +++ b/rust/tvm-ffi/tests/test_dispatch.rs @@ -20,7 +20,9 @@ //! This integration test compiles as a downstream crate, checking every public //! path emitted by the dispatch macro outside `tvm_ffi` itself. -use tvm_ffi::{dispatch, structural_visit, Array, Object, VisitCtx, VisitDispatch, WalkResult}; +use tvm_ffi::{ + dispatch, structural_visit, Array, DefRegionKind, Object, VisitDispatch, WalkResult, +}; struct ExternalCounter { objects: usize, @@ -30,7 +32,7 @@ struct ExternalCounter { impl ExternalCounter { #[cfg(any(unix, windows))] #[cfg_attr(all(), inline)] - fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_object(&mut self, _value: &Object, _kind: DefRegionKind) -> WalkResult { self.objects += 1; WalkResult::Advance } @@ -45,20 +47,16 @@ struct CfgAttrCounter; #[dispatch(visit)] impl CfgAttrCounter { #[cfg(any())] - fn visit_disabled_catch_all( - &mut self, - _value: &tvm_ffi::VisitValue, - _ctx: &mut VisitCtx<'_>, - ) -> WalkResult { + fn visit_disabled_catch_all(&mut self, _value: &tvm_ffi::VisitValue) -> WalkResult { WalkResult::Advance } #[cfg_attr(all(), cfg(any()))] - fn visit_disabled(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_disabled(&mut self, _value: &Object, _kind: DefRegionKind) -> WalkResult { WalkResult::Advance } - fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_object(&mut self, _value: &Object) -> WalkResult { WalkResult::Advance } } @@ -71,7 +69,7 @@ const _: usize = std::mem::size_of::(); #[dispatch(visit)] #[cfg(any())] impl DisabledCounter { - fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_object(&mut self, _value: &Object) -> WalkResult { WalkResult::Advance } } @@ -82,7 +80,7 @@ const _: usize = std::mem::size_of::(); #[dispatch(visit)] #[cfg_attr(all(), cfg(any()))] impl CfgAttrDisabledCounter { - fn visit_object(&mut self, _value: &Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_object(&mut self, _value: &Object, _kind: DefRegionKind) -> WalkResult { WalkResult::Advance } } diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index 6dbf90ee9..09922417e 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -22,8 +22,8 @@ use std::ops::ControlFlow; use tvm_ffi::tvm_ffi_sys::{TVMFFIByteArray, TVMFFITypeIndex, TVMFFITypeRegisterAttr}; use tvm_ffi::{ dispatch, structural_visit, structural_walk, walk, walk_with_context, Any, AnyView, Array, - DefRegionKind, Error, Function, Map, Phase, Shape, String as FfiString, VisitCtx, VisitValue, - WalkOrder, WalkResult, RUNTIME_ERROR, + DefRegionKind, Error, Function, Map, Phase, Result, Shape, String as FfiString, VisitDispatch, + VisitValue, WalkOrder, WalkResult, RUNTIME_ERROR, }; fn runtime_error(message: &str) -> Error { @@ -66,7 +66,7 @@ struct SkipForeignShape; #[dispatch(visit)] impl SkipForeignShape { - fn visit_shape(&mut self, _shape: Shape, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_shape(&mut self, _shape: Shape) -> WalkResult { WalkResult::Skip } } @@ -207,20 +207,22 @@ struct ManualRegionProbe { #[dispatch(visit)] impl ManualRegionProbe { - fn visit_array(&mut self, array: Array, ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_array(&mut self, array: Array, kind: DefRegionKind) -> Result { let overridden = array.get(0).unwrap(); - if !ctx.visit_with_def_region(self, &overridden, DefRegionKind::NonRecursive) { - return WalkResult::Interrupt; + if let ControlFlow::Break(payload) = + self.subvisit(&overridden, DefRegionKind::NonRecursive)? + { + return Ok(WalkResult::InterruptWith(payload)); } let inherited = array.get(1).unwrap(); - if !ctx.visit(self, &inherited) { - return WalkResult::Interrupt; + if let ControlFlow::Break(payload) = self.subvisit(&inherited, kind)? { + return Ok(WalkResult::InterruptWith(payload)); } - WalkResult::Skip + Ok(WalkResult::Skip) } - fn visit_integer(&mut self, _value: i64, ctx: &mut VisitCtx<'_>) -> WalkResult { - self.seen.push(ctx.def_region_kind()); + fn visit_integer(&mut self, _value: i64, kind: DefRegionKind) -> WalkResult { + self.seen.push(kind); WalkResult::Advance } } @@ -245,17 +247,17 @@ struct GenericDispatchProbe { #[dispatch(visit)] impl GenericDispatchProbe { - fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_integer(&mut self, value: i64) -> WalkResult { self.integers.push(value); WalkResult::Advance } - fn visit_object(&mut self, _value: &tvm_ffi::Object, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_object(&mut self, _value: &tvm_ffi::Object) -> WalkResult { self.objects += 1; WalkResult::Advance } - fn visit_any(&mut self, _value: &VisitValue, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_any(&mut self, _value: &VisitValue) -> WalkResult { self.catch_all += 1; WalkResult::Advance } @@ -282,12 +284,12 @@ struct OrderProbe { #[dispatch(visit)] impl OrderProbe { - fn visit_array(&mut self, _array: Array, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_array(&mut self, _array: Array) -> WalkResult { self.events.push("array".to_string()); WalkResult::Advance } - fn visit_integer(&mut self, value: i64, _ctx: &mut VisitCtx<'_>) -> WalkResult { + fn visit_integer(&mut self, value: i64) -> WalkResult { self.events.push(format!("int:{value}")); WalkResult::Advance } From e7fef47c4466b66d039fde79a99a7c71bad202f6 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 30 Jul 2026 12:12:02 -0400 Subject: [PATCH 05/20] remove the kind. --- rust/README.md | 46 ++++++++++----------- rust/tvm-ffi/src/extra/structural_visit.rs | 31 ++++++++++++++ rust/tvm-ffi/tests/test_structural_visit.rs | 39 +++++++++++++++++ 3 files changed, 93 insertions(+), 23 deletions(-) diff --git a/rust/README.md b/rust/README.md index bf669edca..7718cc348 100644 --- a/rust/README.md +++ b/rust/README.md @@ -70,29 +70,22 @@ assert_eq!(calculator.value, 7.5); Typed handlers are tested in source order: borrowed `ObjectCore` node types use runtime subtype checks, owned arguments use `AnyCompatible` casts, and a final -`&VisitValue` handler acts as a catch-all. Use `structural_walk` to select -pre-order or post-order dispatch, or `walk`/`walk_with_context` for raw -callbacks. - -`WalkResult::Advance` visits reflected or container children, `Skip` suppresses -the current value's default recursion, and `Interrupt` stops the complete -traversal. - -A handler may declare an optional third `DefRegionKind` parameter to observe -the definition-region state at the current value (maintained automatically -from reflected field flags): - -```rust,ignore -fn visit_var(&mut self, var: &VarObj, kind: DefRegionKind) -> WalkResult { - // kind tells a definition position apart from a use position. - WalkResult::Advance -} -``` - -To take over a value's children, a pre-order handler visits them through -`VisitDispatch::subvisit` and returns `Skip`. `subvisit` reports a nested -interrupt as `ControlFlow::Break`; propagate it (and errors, via `?`) instead -of dropping the result: +`&VisitValue` handler acts as a catch-all. A handler may take an optional third +`DefRegionKind` argument to observe the definition-region state at the current +value. Use `structural_walk` to select pre-order or post-order dispatch, or +`walk`/`walk_with_context` for raw callbacks that fire at both `Phase::Enter` +and `Phase::Exit` of every value. + +`WalkResult::Advance` visits container or reflected children, `Skip` suppresses +the current value's default recursion, and `Interrupt`/`InterruptWith` halt the +walk, surfacing to the caller as `ControlFlow::Break`. Handlers and callbacks +may also return `Result` to propagate errors with `?`. + +A pre-order handler can take over a value's children instead of advancing: +visit selected children with `subvisit`, or delegate the walker's default +child recursion with `subvisit_children`, then return `Skip`. Both report a +nested interrupt as `ControlFlow::Break`; propagate it (and errors, via `?`) +instead of dropping the result: ```rust,ignore fn visit_func(&mut self, func: &FuncObj, kind: DefRegionKind) -> Result { @@ -106,6 +99,13 @@ fn visit_func(&mut self, func: &FuncObj, kind: DefRegionKind) -> Result Result { + let walker = NativeWalker::new(); + let mut dispatch = DispatchVisitor { + visitor: self, + order: WalkOrder::PreOrder, + }; + let result = walker + .visit_children_raw(value.0, &mut dispatch, def_region_kind) + .map_err(|halt| NativeWalker::with_value_context(halt, value.0)); + finish(result) + } } trait NativeVisit { diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index 09922417e..07dacdc10 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -277,6 +277,45 @@ fn generated_dispatch_supports_pod_and_ordered_catch_all() { assert_eq!(probe.catch_all, 2); } +#[derive(Default)] +struct StraddleProbe { + events: Vec, +} + +#[dispatch(visit)] +impl StraddleProbe { + fn visit_any(&mut self, value: &VisitValue, kind: DefRegionKind) -> Result { + let label = match value.cast::() { + Some(integer) => format!("int:{integer}"), + None => "node".to_string(), + }; + self.events.push(format!("enter:{label}")); + if let ControlFlow::Break(payload) = self.subvisit_children(value, kind)? { + return Ok(WalkResult::InterruptWith(payload)); + } + self.events.push(format!("exit:{label}")); + Ok(WalkResult::Skip) + } +} + +#[test] +fn catch_all_handler_can_straddle_default_children() { + let root = Array::new(vec![1i64, 2]); + let mut probe = StraddleProbe::default(); + assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + assert_eq!( + probe.events, + vec![ + "enter:node", + "enter:int:1", + "exit:int:1", + "enter:int:2", + "exit:int:2", + "exit:node", + ] + ); +} + #[derive(Default)] struct OrderProbe { events: Vec, From 0964c2e1c2619b7eafb86a40705b2e25d7021ed7 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 30 Jul 2026 12:46:02 -0400 Subject: [PATCH 06/20] remove kind. --- rust/README.md | 33 +++--- rust/tvm-ffi-macros/src/visit.rs | 112 +++++++++++------- rust/tvm-ffi/src/extra/structural_visit.rs | 121 +++++++++++++++----- rust/tvm-ffi/tests/test_dispatch.rs | 23 ++-- rust/tvm-ffi/tests/test_structural_visit.rs | 67 ++++++++--- 5 files changed, 246 insertions(+), 110 deletions(-) diff --git a/rust/README.md b/rust/README.md index 7718cc348..d13a27bca 100644 --- a/rust/README.md +++ b/rust/README.md @@ -37,14 +37,15 @@ visitor. Each value is automatically dispatched to the handler matching its runtime type: ```rust -use tvm_ffi::{dispatch, structural_visit, Function, WalkResult}; +use tvm_ffi::{dispatch, structural_visit, DefRegionKind, Function, WalkResult}; #[derive(Default)] struct Calculator { + def_region: DefRegionKind, value: f64, } -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl Calculator { fn visit_integer(&mut self, value: i64) -> WalkResult { self.value += value as f64; @@ -70,11 +71,13 @@ assert_eq!(calculator.value, 7.5); Typed handlers are tested in source order: borrowed `ObjectCore` node types use runtime subtype checks, owned arguments use `AnyCompatible` casts, and a final -`&VisitValue` handler acts as a catch-all. A handler may take an optional third -`DefRegionKind` argument to observe the definition-region state at the current -value. Use `structural_walk` to select pre-order or post-order dispatch, or -`walk`/`walk_with_context` for raw callbacks that fire at both `Phase::Enter` -and `Phase::Exit` of every value. +`&VisitValue` handler acts as a catch-all. Every visitor names a +`DefRegionKind` mirror field through `def_region = `; the walker keeps +that field equal to the definition-region state of the value being dispatched +(and rewinds it after nested traversal), so a handler reads the state through +`self.def_region_kind()` and never writes the field. Use `structural_walk` to +select pre-order or post-order dispatch, or `walk`/`walk_with_context` for raw +callbacks that fire at both `Phase::Enter` and `Phase::Exit` of every value. `WalkResult::Advance` visits container or reflected children, `Skip` suppresses the current value's default recursion, and `Interrupt`/`InterruptWith` halt the @@ -83,16 +86,20 @@ may also return `Result` to propagate errors with `?`. A pre-order handler can take over a value's children instead of advancing: visit selected children with `subvisit`, or delegate the walker's default -child recursion with `subvisit_children`, then return `Skip`. Both report a -nested interrupt as `ControlFlow::Break`; propagate it (and errors, via `?`) -instead of dropping the result: +child recursion with `subvisit_children`, then return `Skip`. Both inherit +the current definition-region state; the `_with_def_region` variants override +it for exactly that subtree. All of them report a nested interrupt as +`ControlFlow::Break`; propagate it (and errors, via `?`) instead of dropping +the result: ```rust,ignore -fn visit_func(&mut self, func: &FuncObj, kind: DefRegionKind) -> Result { - if let ControlFlow::Break(payload) = self.subvisit(&func.params, DefRegionKind::Recursive)? { +fn visit_func(&mut self, func: &FuncObj) -> Result { + if let ControlFlow::Break(payload) = + self.subvisit_with_def_region(&func.params, DefRegionKind::Recursive)? + { return Ok(WalkResult::InterruptWith(payload)); } - if let ControlFlow::Break(payload) = self.subvisit(&func.body, kind)? { + if let ControlFlow::Break(payload) = self.subvisit(&func.body)? { return Ok(WalkResult::InterruptWith(payload)); } Ok(WalkResult::Skip) diff --git a/rust/tvm-ffi-macros/src/visit.rs b/rust/tvm-ffi-macros/src/visit.rs index 54bec9dff..4a0d44a4c 100644 --- a/rust/tvm-ffi-macros/src/visit.rs +++ b/rust/tvm-ffi-macros/src/visit.rs @@ -24,10 +24,10 @@ use quote::{quote, quote_spanned}; use syn::{parse_macro_input, FnArg, ImplItem, ImplItemMethod, ItemImpl, Meta, NestedMeta, Type}; pub(crate) fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { - let mode = parse_macro_input!(attr as syn::Ident); + let args = parse_macro_input!(attr as DispatchArgs); let item_impl = parse_macro_input!(item as ItemImpl); - match expand(&mode, &item_impl) { + match expand(&args, &item_impl) { Ok(generated) => quote!(#item_impl #generated).into(), Err(error) => { let error = error.to_compile_error(); @@ -36,10 +36,43 @@ pub(crate) fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { } } -fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result { - if mode != "visit" { - return Err(syn::Error::new(mode.span(), "expected `dispatch(visit)`")); +struct DispatchArgs { + def_region_field: syn::Ident, +} + +impl syn::parse::Parse for DispatchArgs { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let mode: syn::Ident = input.parse()?; + if mode != "visit" { + return Err(syn::Error::new( + mode.span(), + "expected `dispatch(visit, def_region = )`", + )); + } + if input.parse::().is_err() { + return Err(syn::Error::new( + mode.span(), + "`dispatch(visit)` requires `def_region = ` naming the visitor's \ + `DefRegionKind` mirror field", + )); + } + let key: syn::Ident = input.parse()?; + if key != "def_region" { + return Err(syn::Error::new( + key.span(), + "expected `def_region = `", + )); + } + input.parse::()?; + let def_region_field: syn::Ident = input.parse()?; + if !input.is_empty() { + return Err(input.error("unexpected tokens after `def_region = `")); + } + Ok(DispatchArgs { def_region_field }) } +} + +fn expand(args: &DispatchArgs, item_impl: &ItemImpl) -> syn::Result { if item_impl.trait_.is_some() { return Err(syn::Error::new_spanned( item_impl, @@ -65,20 +98,16 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result )); } let tvm_ffi = resolve_tvm_ffi_crate()?; + let def_region_field = &args.def_region_field; let links = handlers.iter().map(|handler| { let method = &handler.method; let attrs = &handler.cfg_attrs; - let kind_arg = if handler.takes_def_region_kind { - quote!(, def_region_kind) - } else { - quote!() - }; let invoke = match &handler.argument { HandlerArgument::Value => quote! { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(value #kind_arg) + self.#method(value) ) ); }, @@ -86,7 +115,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result if let Some(node) = value.as_node::<#node_type>() { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(node #kind_arg) + self.#method(node) ) ); } @@ -95,7 +124,7 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result if let Some(node) = value.cast::<#value_type>() { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(node #kind_arg) + self.#method(node) ) ); } @@ -142,11 +171,33 @@ fn expand(mode: &syn::Ident, item_impl: &ItemImpl) -> syn::Result fn dispatch_visit( &mut self, value: &#tvm_ffi::extra::structural_visit::VisitValue, - def_region_kind: #tvm_ffi::extra::structural_visit::DefRegionKind, ) -> Option<#tvm_ffi::extra::structural_visit::VisitResult> { #(#links)* None } + + fn def_region_kind(&self) -> #tvm_ffi::extra::structural_visit::DefRegionKind { + self.#def_region_field + } + + #[doc(hidden)] + fn def_region_slot( + &mut self, + ) -> &mut #tvm_ffi::extra::structural_visit::DefRegionKind { + &mut self.#def_region_field + } + } + + #(#[#impl_cfg_attrs])* + impl #impl_generics #self_type #where_clause { + /// Definition-region state at the value currently being dispatched. + /// + /// Inherent mirror of `VisitDispatch::def_region_kind`, callable + /// without importing the trait. + #[allow(dead_code)] + fn def_region_kind(&self) -> #tvm_ffi::extra::structural_visit::DefRegionKind { + self.#def_region_field + } } }) } @@ -171,7 +222,6 @@ fn crate_path(found: FoundCrate) -> TokenStream2 { struct Handler { method: syn::Ident, argument: HandlerArgument, - takes_def_region_kind: bool, cfg_attrs: Vec, } @@ -188,29 +238,14 @@ fn parse_handler(method: &ImplItemMethod) -> syn::Result { Some(FnArg::Receiver(receiver)) if receiver.reference.is_some() && receiver.mutability.is_some() ); - if !receiver_is_mut || (inputs.len() != 2 && inputs.len() != 3) { + if !receiver_is_mut || inputs.len() != 2 { return Err(syn::Error::new_spanned( &method.sig, - "visit handlers must take `&mut self`, a node, and optionally a `DefRegionKind`", + "visit handlers must take `&mut self` and a node; read definition-region state \ + through `self.def_region_kind()`", )); } - let takes_def_region_kind = if inputs.len() == 3 { - let kind_type = match inputs.iter().nth(2) { - Some(FnArg::Typed(kind)) => &kind.ty, - _ => unreachable!("the third argument cannot be a receiver"), - }; - if !is_def_region_kind(kind_type) { - return Err(syn::Error::new_spanned( - kind_type, - "the third visit handler argument must be `DefRegionKind` (by value)", - )); - } - true - } else { - false - }; - let value_type = match inputs.iter().nth(1) { Some(FnArg::Typed(value)) => (*value.ty).clone(), _ => unreachable!("the second argument cannot be a receiver"), @@ -235,7 +270,6 @@ fn parse_handler(method: &ImplItemMethod) -> syn::Result { Ok(Handler { method: method.sig.ident.clone(), argument, - takes_def_region_kind, cfg_attrs, }) } @@ -289,16 +323,6 @@ fn is_visit_value(value_type: &Type) -> bool { .is_some_and(|segment| segment.ident == "VisitValue") } -fn is_def_region_kind(kind_type: &Type) -> bool { - let Type::Path(path) = kind_type else { - return false; - }; - path.path - .segments - .last() - .is_some_and(|segment| segment.ident == "DefRegionKind" && segment.arguments.is_empty()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index 4952aa15c..d3ad344f6 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -219,26 +219,59 @@ type NativeResult = std::result::Result<(), NativeHalt>; /// order. Borrowed node arguments use refcount-free subtype checks, owned /// FFI-compatible arguments use exact value casts, and `&VisitValue` is a /// catch-all. `None` asks the Rust walker to continue normally. +/// +/// Every visitor carries a definition-region mirror field, named through +/// `#[dispatch(visit, def_region = )]`. The walker writes the field +/// before each dispatched handler and restores it afterwards, so during a +/// handler [`VisitDispatch::def_region_kind`] always reports the state at +/// the current value — including after nested [`VisitDispatch::subvisit`] +/// calls. Treat the field as read-only; the walker's own propagation never +/// reads it, so overwriting it only misleads your own reads and the +/// kind-inheriting `subvisit` forms. pub trait VisitDispatch: Sized { - fn dispatch_visit( - &mut self, - value: &VisitValue, - def_region_kind: DefRegionKind, - ) -> Option; + fn dispatch_visit(&mut self, value: &VisitValue) -> Option; + + /// Return the definition-region state at the value being dispatched. + /// + /// Meaningful only while a handler is running; outside a walk the mirror + /// holds its initial or last-restored value. + fn def_region_kind(&self) -> DefRegionKind; + + /// Mirror storage refreshed by the walker around each dispatch. + #[doc(hidden)] + fn def_region_slot(&mut self) -> &mut DefRegionKind; - /// Visit `child` immediately with this dispatcher under `def_region_kind`. + /// Visit `child` immediately, inheriting the current definition-region + /// state. /// /// This is how a pre-order handler takes over a value's children: visit - /// each selected child, then return [`WalkResult::Skip`]. Pass the - /// handler's received `def_region_kind` to keep the inherited - /// definition-region state, or another kind to override it for exactly - /// this subtree. The nested traversal always dispatches pre-order. + /// each selected child, then return [`WalkResult::Skip`]. Use + /// [`VisitDispatch::subvisit_with_def_region`] to override the state for + /// exactly this subtree. The nested traversal always dispatches + /// pre-order. /// /// `Err` carries a nested handler failure and should be propagated with /// `?`. `Ok(ControlFlow::Break(payload))` reports a nested interrupt; /// return [`WalkResult::InterruptWith`] with the payload to keep /// halting. Dropping the result silently swallows both. - fn subvisit(&mut self, child: &T, def_region_kind: DefRegionKind) -> Result + fn subvisit(&mut self, child: &T) -> Result + where + for<'x> AnyView<'x>: From<&'x T>, + { + let def_region_kind = self.def_region_kind(); + self.subvisit_with_def_region(child, def_region_kind) + } + + /// Visit `child` immediately under an explicitly selected + /// definition-region state. + /// + /// The override is scoped to this recursive call; see + /// [`VisitDispatch::subvisit`] for the result contract. + fn subvisit_with_def_region( + &mut self, + child: &T, + def_region_kind: DefRegionKind, + ) -> Result where for<'x> AnyView<'x>: From<&'x T>, { @@ -251,10 +284,10 @@ pub trait VisitDispatch: Sized { } /// Visit `value`'s children — not `value` itself — with the walker's - /// default rules: container contents for `Array`/`List`/`Map`/`Dict`, - /// reflected structural fields otherwise. + /// default rules, inheriting the current definition-region state. /// - /// This is the Rust analog of C++ + /// Children are container contents for `Array`/`List`/`Map`/`Dict` and + /// reflected structural fields otherwise. This is the Rust analog of C++ /// `StructuralVisitorObj::DefaultVisitExpected`: a handler may run its /// enter logic, delegate the default child recursion explicitly, run its /// exit logic with the same locals in scope, and return @@ -262,10 +295,17 @@ pub trait VisitDispatch: Sized { /// knowledge of the value's concrete type, so it also works from a /// `&VisitValue` catch-all handler. /// - /// The result contract matches `subvisit`: propagate `Err` with `?` and - /// map `Ok(ControlFlow::Break(payload))` to - /// [`WalkResult::InterruptWith`]. - fn subvisit_children( + /// The result contract matches [`VisitDispatch::subvisit`]. + fn subvisit_children(&mut self, value: &VisitValue) -> Result { + let def_region_kind = self.def_region_kind(); + self.subvisit_children_with_def_region(value, def_region_kind) + } + + /// Visit `value`'s children with the walker's default rules under an + /// explicitly selected definition-region state. + /// + /// See [`VisitDispatch::subvisit_children`]. + fn subvisit_children_with_def_region( &mut self, value: &VisitValue, def_region_kind: DefRegionKind, @@ -295,13 +335,31 @@ struct DispatchVisitor<'a, V> { order: WalkOrder, } +impl DispatchVisitor<'_, V> { + /// Run one typed dispatch with the mirror scoped to `def_region_kind`. + /// + /// The save/restore pair keeps the mirror equal to the dispatched value's + /// state for the whole handler call, and transparently rewinds it after + /// nested `subvisit` recursion. + fn dispatch_scoped( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result { + let saved = std::mem::replace(self.visitor.def_region_slot(), def_region_kind); + let result = self + .visitor + .dispatch_visit(value) + .unwrap_or(Ok(WalkResult::Advance)); + *self.visitor.def_region_slot() = saved; + result + } +} + impl NativeVisit for DispatchVisitor<'_, V> { fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { match self.order { - WalkOrder::PreOrder => self - .visitor - .dispatch_visit(value, def_region_kind) - .unwrap_or(Ok(WalkResult::Advance)), + WalkOrder::PreOrder => self.dispatch_scoped(value, def_region_kind), WalkOrder::PostOrder => Ok(WalkResult::Advance), } } @@ -309,10 +367,7 @@ impl NativeVisit for DispatchVisitor<'_, V> { fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { match self.order { WalkOrder::PreOrder => Ok(WalkResult::Advance), - WalkOrder::PostOrder => self - .visitor - .dispatch_visit(value, def_region_kind) - .unwrap_or(Ok(WalkResult::Advance)), + WalkOrder::PostOrder => self.dispatch_scoped(value, def_region_kind), } } } @@ -898,12 +953,16 @@ mod tests { } #[derive(Default)] - struct TypedRegionProbe(Vec); + struct TypedRegionProbe { + def_region: DefRegionKind, + seen: Vec, + } - #[crate::dispatch(visit)] + #[crate::dispatch(visit, def_region = def_region)] impl TypedRegionProbe { - fn visit_integer(&mut self, _value: i64, def_region_kind: DefRegionKind) -> WalkResult { - self.0.push(def_region_kind); + fn visit_integer(&mut self, _value: i64) -> WalkResult { + let kind = self.def_region_kind(); + self.seen.push(kind); WalkResult::Advance } } @@ -957,7 +1016,7 @@ mod tests { .is_ok()); } assert_eq!( - probe.0, + probe.seen, vec![ DefRegionKind::Recursive, DefRegionKind::None, diff --git a/rust/tvm-ffi/tests/test_dispatch.rs b/rust/tvm-ffi/tests/test_dispatch.rs index 4370dcd42..b02da3650 100644 --- a/rust/tvm-ffi/tests/test_dispatch.rs +++ b/rust/tvm-ffi/tests/test_dispatch.rs @@ -24,15 +24,17 @@ use tvm_ffi::{ dispatch, structural_visit, Array, DefRegionKind, Object, VisitDispatch, WalkResult, }; +#[derive(Default)] struct ExternalCounter { + def_region: DefRegionKind, objects: usize, } -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl ExternalCounter { #[cfg(any(unix, windows))] #[cfg_attr(all(), inline)] - fn visit_object(&mut self, _value: &Object, _kind: DefRegionKind) -> WalkResult { + fn visit_object(&mut self, _value: &Object) -> WalkResult { self.objects += 1; WalkResult::Advance } @@ -42,9 +44,12 @@ fn assert_visit_dispatch() {} const _: fn() = assert_visit_dispatch::; -struct CfgAttrCounter; +#[derive(Default)] +struct CfgAttrCounter { + def_region: DefRegionKind, +} -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl CfgAttrCounter { #[cfg(any())] fn visit_disabled_catch_all(&mut self, _value: &tvm_ffi::VisitValue) -> WalkResult { @@ -52,7 +57,7 @@ impl CfgAttrCounter { } #[cfg_attr(all(), cfg(any()))] - fn visit_disabled(&mut self, _value: &Object, _kind: DefRegionKind) -> WalkResult { + fn visit_disabled(&mut self, _value: &Object) -> WalkResult { WalkResult::Advance } @@ -66,7 +71,7 @@ const _: fn() = assert_visit_dispatch::; struct DisabledCounter; const _: usize = std::mem::size_of::(); -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] #[cfg(any())] impl DisabledCounter { fn visit_object(&mut self, _value: &Object) -> WalkResult { @@ -77,10 +82,10 @@ impl DisabledCounter { struct CfgAttrDisabledCounter; const _: usize = std::mem::size_of::(); -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] #[cfg_attr(all(), cfg(any()))] impl CfgAttrDisabledCounter { - fn visit_object(&mut self, _value: &Object, _kind: DefRegionKind) -> WalkResult { + fn visit_object(&mut self, _value: &Object) -> WalkResult { WalkResult::Advance } } @@ -88,7 +93,7 @@ impl CfgAttrDisabledCounter { #[test] fn generated_dispatch_uses_public_downstream_paths() { let root = Array::new(vec![1i64, 2]); - let mut visitor = ExternalCounter { objects: 0 }; + let mut visitor = ExternalCounter::default(); assert!(structural_visit(&root, &mut visitor).unwrap().is_continue()); assert_eq!(visitor.objects, 1); } diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index 07dacdc10..6b98e5cac 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -62,9 +62,12 @@ fn plain_walk_uses_native_map_fallback() { assert_eq!(integers, 2); } -struct SkipForeignShape; +#[derive(Default)] +struct SkipForeignShape { + def_region: DefRegionKind, +} -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl SkipForeignShape { fn visit_shape(&mut self, _shape: Shape) -> WalkResult { WalkResult::Skip @@ -95,7 +98,7 @@ fn foreign_structural_visit_requires_explicit_rust_override() { assert!(error.message().contains("registers foreign `__s_visit__`")); assert!(error.message().contains("return `WalkResult::Skip`")); - assert!(structural_visit(&root, &mut SkipForeignShape) + assert!(structural_visit(&root, &mut SkipForeignShape::default()) .unwrap() .is_continue()); } @@ -202,31 +205,66 @@ fn interrupt_stops_without_running_remaining_callbacks() { #[derive(Default)] struct ManualRegionProbe { + def_region: DefRegionKind, seen: Vec, } -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl ManualRegionProbe { - fn visit_array(&mut self, array: Array, kind: DefRegionKind) -> Result { + fn visit_array(&mut self, array: Array) -> Result { + let before = self.def_region_kind(); let overridden = array.get(0).unwrap(); if let ControlFlow::Break(payload) = - self.subvisit(&overridden, DefRegionKind::NonRecursive)? + self.subvisit_with_def_region(&overridden, DefRegionKind::NonRecursive)? { return Ok(WalkResult::InterruptWith(payload)); } + // The walker rewinds the mirror after nested dispatches. + assert_eq!(self.def_region_kind(), before); let inherited = array.get(1).unwrap(); - if let ControlFlow::Break(payload) = self.subvisit(&inherited, kind)? { + if let ControlFlow::Break(payload) = self.subvisit(&inherited)? { return Ok(WalkResult::InterruptWith(payload)); } Ok(WalkResult::Skip) } - fn visit_integer(&mut self, _value: i64, kind: DefRegionKind) -> WalkResult { + fn visit_integer(&mut self, _value: i64) -> WalkResult { + let kind = self.def_region_kind(); self.seen.push(kind); WalkResult::Advance } } +#[derive(Default)] +struct MirrorCorruptionProbe { + def_region: DefRegionKind, + seen: Vec, +} + +#[dispatch(visit, def_region = def_region)] +impl MirrorCorruptionProbe { + fn visit_array(&mut self, _array: Array) -> WalkResult { + // Overwriting the mirror must not leak into the walker's own + // by-value propagation. + self.def_region = DefRegionKind::Recursive; + WalkResult::Advance + } + + fn visit_integer(&mut self, _value: i64) -> WalkResult { + let kind = self.def_region_kind(); + self.seen.push(kind); + WalkResult::Advance + } +} + +#[test] +fn mirror_corruption_does_not_affect_walker_propagation() { + let root = Array::new(vec![7i64]); + let mut probe = MirrorCorruptionProbe::default(); + assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + assert_eq!(probe.seen, vec![DefRegionKind::None]); +} + #[test] fn manual_child_visit_can_override_def_region() { let root = Array::new(vec![7i64, 8]); @@ -240,12 +278,13 @@ fn manual_child_visit_can_override_def_region() { #[derive(Default)] struct GenericDispatchProbe { + def_region: DefRegionKind, integers: Vec, objects: usize, catch_all: usize, } -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl GenericDispatchProbe { fn visit_integer(&mut self, value: i64) -> WalkResult { self.integers.push(value); @@ -279,18 +318,19 @@ fn generated_dispatch_supports_pod_and_ordered_catch_all() { #[derive(Default)] struct StraddleProbe { + def_region: DefRegionKind, events: Vec, } -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl StraddleProbe { - fn visit_any(&mut self, value: &VisitValue, kind: DefRegionKind) -> Result { + fn visit_any(&mut self, value: &VisitValue) -> Result { let label = match value.cast::() { Some(integer) => format!("int:{integer}"), None => "node".to_string(), }; self.events.push(format!("enter:{label}")); - if let ControlFlow::Break(payload) = self.subvisit_children(value, kind)? { + if let ControlFlow::Break(payload) = self.subvisit_children(value)? { return Ok(WalkResult::InterruptWith(payload)); } self.events.push(format!("exit:{label}")); @@ -318,10 +358,11 @@ fn catch_all_handler_can_straddle_default_children() { #[derive(Default)] struct OrderProbe { + def_region: DefRegionKind, events: Vec, } -#[dispatch(visit)] +#[dispatch(visit, def_region = def_region)] impl OrderProbe { fn visit_array(&mut self, _array: Array) -> WalkResult { self.events.push("array".to_string()); From 9e86a84902ee6ed2529ae843425246838880c66f Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 30 Jul 2026 13:02:57 -0400 Subject: [PATCH 07/20] add tests. Signed-off-by: yuchuan --- rust/tvm-ffi/tests/test_visitor_alignment.rs | 139 +++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 rust/tvm-ffi/tests/test_visitor_alignment.rs diff --git a/rust/tvm-ffi/tests/test_visitor_alignment.rs b/rust/tvm-ffi/tests/test_visitor_alignment.rs new file mode 100644 index 000000000..62c74d54b --- /dev/null +++ b/rust/tvm-ffi/tests/test_visitor_alignment.rs @@ -0,0 +1,139 @@ +/* + * 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. + */ + +//! Rust mirror of the C++ visitor example. +//! +//! `RecordingVisitor` lines up member-for-member with the C++ +//! `TestVisitorObj` (tests/cpp/extra/test_structural_visit.cc), and its +//! `visit_array` handler plays the role of the C++ `TFuncObj::StructuralVisit` +//! hook (tests/cpp/testing_object.h): the first element is visited as a +//! recursive definition region, the rest inherit the surrounding state. + +use std::ops::ControlFlow; + +use tvm_ffi::{ + dispatch, structural_visit, Array, DefRegionKind, Result, String as FfiString, VisitDispatch, + VisitValue, WalkResult, +}; + +/// C++: class TestVisitorObj : public StructuralVisitorObj +#[derive(Default)] +struct RecordingVisitor { + /// C++: `def_region_mode_` (base-class member maintained by the walker). + def_region: DefRegionKind, + /// C++: `std::vector visited;` + visited: Vec, + /// C++: `std::vector modes;` + modes: Vec, + /// C++: `ObjectRef interrupt_on;` + interrupt_on: Option, +} + +#[dispatch(visit, def_region = def_region)] +impl RecordingVisitor { + /// C++ analog: `TFuncObj::StructuralVisit` — visit "params" (element 0) + /// under a recursive definition region, then the "body" (element 1) + /// under the inherited state, and skip the default recursion. + fn visit_array(&mut self, array: Array) -> Result { + self.visited.push("array".to_string()); + self.modes.push(self.def_region_kind()); + + // C++: visitor->WithDefRegionKind(kTVMFFIDefRegionKindRecursive, + // [&] { return visitor->VisitExpected(self->params); }) + let params = array.get(0).unwrap(); + if let ControlFlow::Break(payload) = + self.subvisit_with_def_region(¶ms, DefRegionKind::Recursive)? + { + return Ok(WalkResult::InterruptWith(payload)); + } + + // C++: visitor->VisitExpected(self->body) (inherits the current state) + let body = array.get(1).unwrap(); + if let ControlFlow::Break(payload) = self.subvisit(&body)? { + return Ok(WalkResult::InterruptWith(payload)); + } + Ok(WalkResult::Skip) + } + + /// C++ analog: `TestVisitorObj::VisitImpl` — record every value together + /// with the active def-region state, optionally interrupt with a payload, + /// otherwise delegate the default recursion. + fn visit_any(&mut self, value: &VisitValue) -> Result { + let label = match value.cast::() { + Some(integer) => integer.to_string(), + None => "obj".to_string(), + }; + // C++: visited.push_back(value_ref); + // modes.push_back(def_region_mode_); + self.visited.push(label); + self.modes.push(self.def_region_kind()); + + // C++: if (value_ref.same_as(interrupt_on)) + // return VisitInterrupt(String("stop")); + if self.interrupt_on.is_some() && value.cast::() == self.interrupt_on { + return Ok(WalkResult::interrupt_with(FfiString::from("stop"))); + } + + // C++: return DefaultVisitExpected(value); + if let ControlFlow::Break(payload) = self.subvisit_children(value)? { + return Ok(WalkResult::InterruptWith(payload)); + } + Ok(WalkResult::Skip) + } +} + +/// C++ analog: TEST(StructuralVisitor, TraversesFunction) — the def-region +/// state flips to Recursive under "params" and back to None for the "body". +#[test] +fn records_values_and_def_region_modes() { + let root = Array::new(vec![10i64, 20]); + let mut visitor = RecordingVisitor::default(); + + let outcome = structural_visit(&root, &mut visitor).unwrap(); + + assert!(outcome.is_continue()); + assert_eq!(visitor.visited, vec!["array", "10", "20"]); + assert_eq!( + visitor.modes, + vec![ + DefRegionKind::None, // the array itself + DefRegionKind::Recursive, // element 0: the "params" position + DefRegionKind::None, // element 1: the "body" position + ] + ); +} + +/// C++ analog: TEST(StructuralVisitor, StopsOnInterrupt) — the traversal +/// halts at the marked value and the payload reaches the caller. +#[test] +fn stops_on_interrupt_with_payload() { + let root = Array::new(vec![10i64, 20]); + let mut visitor = RecordingVisitor { + interrupt_on: Some(20), + ..RecordingVisitor::default() + }; + + let outcome = structural_visit(&root, &mut visitor).unwrap(); + + let ControlFlow::Break(payload) = outcome else { + panic!("traversal unexpectedly completed"); + }; + assert_eq!(FfiString::try_from(payload).unwrap().as_str(), "stop"); + assert_eq!(visitor.visited, vec!["array", "10", "20"]); +} From 003d60a638e582cefce38ba6ce4cc2da4b367644 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 30 Jul 2026 17:31:26 -0400 Subject: [PATCH 08/20] perf. Signed-off-by: yuchuan optim the effciciency. Signed-off-by: yuchuan --- rust/tvm-ffi/src/extra/structural_visit.rs | 232 +++++++++++++++++++- rust/tvm-ffi/tests/test_structural_visit.rs | 59 +++++ 2 files changed, 285 insertions(+), 6 deletions(-) diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index d3ad344f6..55fc261ce 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -39,6 +39,7 @@ use std::ops::ControlFlow; use std::os::raw::c_void; use std::ptr::NonNull; +use std::sync::atomic::{AtomicU8, Ordering}; use crate::any::{Any, AnyView}; use crate::error::{Error, Result, RUNTIME_ERROR, TYPE_ERROR}; @@ -193,8 +194,17 @@ impl VisitValue { if self.0.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { return None; } - if !is_instance(self.0.type_index, N::type_index()) { - return None; + let base_type_index = N::type_index(); + if self.0.type_index != base_type_index { + // A final type has no registered subtype, so a differing index can + // never match: reject with the integer compare alone, mirroring the + // `_type_final` fast path of C++ `IsObjectInstance`. + if N::TYPE_FINAL { + return None; + } + if !is_instance_at_depth(self.0.type_index, base_type_index, N::TYPE_DEPTH) { + return None; + } } Some(unsafe { &*(self.0.data_union.v_obj as *const N) }) } @@ -437,6 +447,7 @@ impl NativeWalker { } } + #[inline] fn visit_children_raw( &self, value: TVMFFIAny, @@ -452,6 +463,16 @@ impl NativeWalker { x if x == TVMFFITypeIndex::kTVMFFIMap as i32 || x == TVMFFITypeIndex::kTVMFFIDict as i32 => { + // Fast path: read the MapBaseObj storage layout directly, like + // the SeqPrefix path for arrays — zero FFI calls per entry. + // Dict entries are snapshotted first to keep the re-entrant + // mutation guard. If the one-time layout validation fails + // (e.g. an ABI-debug build), fall back to the packed-functor + // iteration protocol. + if map_layout_usable(value) { + let snapshot = x == TVMFFITypeIndex::kTVMFFIDict as i32; + return self.visit_map_layout(value, visitor, def_region_kind, snapshot); + } return self.visit_map(value, visitor, def_region_kind); } _ => {} @@ -465,6 +486,7 @@ impl NativeWalker { } } + #[inline(never)] fn visit_sequence( &self, value: TVMFFIAny, @@ -517,6 +539,61 @@ impl NativeWalker { Ok(()) } + /// Walk map/dict entries by reading the `MapBaseObj` storage directly — + /// the map analog of the `SeqPrefix` array fast path. `snapshot` first + /// takes owned copies of all entries (Dict re-entrant mutation guard). + #[inline(never)] + fn visit_map_layout( + &self, + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, + snapshot: bool, + ) -> NativeResult { + let map = unsafe { &*(value.data_union.v_obj as *const MapPrefix) }; + let size = map.size as usize; + if size == 0 { + return Ok(()); + } + let mut cursor = unsafe { MapCursor::new(map) }; + + if snapshot { + let mut entries: Vec<(Any, Any)> = Vec::with_capacity(size); + for _ in 0..size { + let Some((key, val)) = (unsafe { cursor.next() }) else { + return Err(runtime_error("native visitor: map iteration ended early").into()); + }; + entries.push(( + Any::from(unsafe { view_of(&key) }), + Any::from(unsafe { view_of(&val) }), + )); + } + for (index, (mut key, mut val)) in entries.into_iter().enumerate() { + let key_raw = raw_of_owned(&mut key); + self.visit_raw(key_raw, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("dict key [{index}]")))?; + let val_raw = raw_of_owned(&mut val); + self.visit_raw(val_raw, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("dict value [{index}]")))?; + } + return Ok(()); + } + + // Immutable map: entry cells stay stable throughout recursive + // callbacks, so visit them in place. The `size` bound also guards the + // dense iteration list against corruption-induced cycles. + for index in 0..size { + let Some((key, val)) = (unsafe { cursor.next() }) else { + return Err(runtime_error("native visitor: map iteration ended early").into()); + }; + self.visit_raw(key, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("map key [{index}]")))?; + self.visit_raw(val, visitor, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("map value [{index}]")))?; + } + Ok(()) + } + fn visit_map( &self, value: TVMFFIAny, @@ -589,6 +666,7 @@ impl NativeWalker { Ok(()) } + #[inline] fn visit_reflected_fields( &self, value: TVMFFIAny, @@ -768,6 +846,7 @@ fn finish(result: NativeResult) -> Result { } } +#[inline] fn field_def_region(field: &TVMFFIFieldInfo, inherited: DefRegionKind) -> DefRegionKind { if field.flags & FLAG_SEQ_HASH_DEF_NON_RECURSIVE != 0 { DefRegionKind::NonRecursive @@ -782,6 +861,7 @@ fn field_def_region(field: &TVMFFIFieldInfo, inherited: DefRegionKind) -> DefReg /// the FreeVar's own reflected children: nested free vars there must resolve /// against an outer binding instead of rebinding. Mirrors C++ /// `VisitReflectedFieldsExpected`. +#[inline] fn free_var_child_region(inherited: DefRegionKind, structural_eq_hash_kind: i32) -> DefRegionKind { if inherited == DefRegionKind::NonRecursive && structural_eq_hash_kind == TVMFFISEqHashKind::kTVMFFISEqHashKindFreeVar as i32 @@ -806,6 +886,142 @@ fn runtime_error(message: &str) -> Error { Error::new(RUNTIME_ERROR, message, "") } +/// Layout prefix shared by the C++ `MapObj` and `DictObj` (`MapBaseObj`, +/// release ABI without `TVM_FFI_DEBUG_WITH_ABI_CHANGE`). +#[repr(C)] +struct MapPrefix { + _header: TVMFFIObject, + data: *mut u8, + size: u64, + slots: u64, + _data_deleter: Option, +} + +/// Dense-layout extension of the prefix (`DenseMapBaseObj`). +#[repr(C)] +struct DenseMapPrefix { + base: MapPrefix, + fib_shift: u32, + iter_list_head: u64, + iter_list_tail: u64, +} + +const _: () = { + assert!(std::mem::offset_of!(MapPrefix, data) == 24); + assert!(std::mem::offset_of!(MapPrefix, size) == 32); + assert!(std::mem::offset_of!(MapPrefix, slots) == 40); + assert!(std::mem::offset_of!(MapPrefix, _data_deleter) == 48); + assert!(std::mem::offset_of!(DenseMapPrefix, fib_shift) == 56); + assert!(std::mem::offset_of!(DenseMapPrefix, iter_list_head) == 64); +}; + +/// MSB tag on `slots_` marking the small (inline KV array) layout. +const MAP_SMALL_TAG: u64 = 1 << 63; +/// `kInvalidIndex`: terminator of the dense iteration list. +const MAP_INVALID_INDEX: u64 = u64::MAX; +/// `kBlockCap`: entries per dense block. +const MAP_BLOCK_CAP: u64 = 16; +/// `sizeof(ItemType)`: KV pair (32 bytes) + prev/next indices (16 bytes). +const MAP_ITEM_SIZE: usize = 48; +/// `sizeof(Block)`: `kBlockCap` metadata bytes + `kBlockCap` items. +const MAP_BLOCK_SIZE: usize = 16 + 16 * MAP_ITEM_SIZE; +/// Byte offset of `ItemType::next` (after the 32-byte KV pair and `prev`). +const MAP_ITEM_NEXT_OFFSET: usize = 40; + +/// Borrowed traversal cursor over either map storage layout, yielding entries +/// in the same order as the C++ iterator. +enum MapCursor { + Small { + kv: *const TVMFFIAny, + index: usize, + size: usize, + }, + Dense { + data: *const u8, + index: u64, + }, +} + +impl MapCursor { + #[inline] + unsafe fn new(map: &MapPrefix) -> MapCursor { + if map.slots & MAP_SMALL_TAG != 0 { + MapCursor::Small { + kv: map.data as *const TVMFFIAny, + index: 0, + size: map.size as usize, + } + } else { + let dense = &*(map as *const MapPrefix as *const DenseMapPrefix); + MapCursor::Dense { + data: map.data, + index: dense.iter_list_head, + } + } + } + + #[inline] + unsafe fn next(&mut self) -> Option<(TVMFFIAny, TVMFFIAny)> { + match self { + MapCursor::Small { kv, index, size } => { + if *index >= *size { + return None; + } + let pair = kv.add(*index * 2); + *index += 1; + Some((*pair, *pair.add(1))) + } + MapCursor::Dense { data, index } => { + if *index == MAP_INVALID_INDEX { + return None; + } + let block = data.add((*index / MAP_BLOCK_CAP) as usize * MAP_BLOCK_SIZE); + let item = block.add(MAP_BLOCK_CAP as usize + (*index % MAP_BLOCK_CAP) as usize * MAP_ITEM_SIZE); + let key = *(item as *const TVMFFIAny); + let val = *(item.add(16) as *const TVMFFIAny); + *index = *(item.add(MAP_ITEM_NEXT_OFFSET) as *const u64); + Some((key, val)) + } + } + } +} + +/// Process-wide result of the one-time map layout validation: +/// 0 = unknown, 1 = usable, 2 = unusable. +static MAP_LAYOUT_STATE: AtomicU8 = AtomicU8::new(0); + +#[inline] +fn map_layout_usable(value: TVMFFIAny) -> bool { + match MAP_LAYOUT_STATE.load(Ordering::Relaxed) { + 1 => true, + 2 => false, + _ => { + let usable = validate_map_layout(value); + MAP_LAYOUT_STATE.store(if usable { 1 } else { 2 }, Ordering::Relaxed); + usable + } + } +} + +/// Cross-check the mirrored `MapBaseObj` layout against the public size +/// functor once per process. An ABI-debug build inserts a state marker that +/// shifts every field by 8 bytes, which this detects: offset 32 then holds a +/// pointer value that cannot equal the reported entry count. +fn validate_map_layout(value: TVMFFIAny) -> bool { + let expected = (|| -> Result { + let is_dict = value.type_index == TVMFFITypeIndex::kTVMFFIDict as i32; + let name = if is_dict { "ffi.DictSize" } else { "ffi.MapSize" }; + Function::get_global(name)? + .call_packed(&[unsafe { view_of(&value) }]) + .and_then(i64::try_from) + })(); + let Ok(expected) = expected else { + return false; + }; + let map = unsafe { &*(value.data_union.v_obj as *const MapPrefix) }; + expected >= 0 && map.size == expected as u64 +} + /// Layout prefix shared by the C++ `ArrayObj` and `ListObj`. #[repr(C)] struct SeqPrefix { @@ -855,17 +1071,18 @@ fn type_key_of(type_index: i32) -> String { } } -fn is_instance(object_type_index: i32, base_type_index: i32) -> bool { +/// Subtype check with the base's inheritance depth supplied by the caller +/// (`ObjectCore::TYPE_DEPTH`), so only the object's type info is fetched. +#[inline] +fn is_instance_at_depth(object_type_index: i32, base_type_index: i32, base_depth: i32) -> bool { if object_type_index == base_type_index { return true; } unsafe { let info = TVMFFIGetTypeInfo(object_type_index); - let base_info = TVMFFIGetTypeInfo(base_type_index); - if info.is_null() || base_info.is_null() { + if info.is_null() { return false; } - let base_depth = (*base_info).type_depth; if (*info).type_depth <= base_depth { return false; } @@ -922,14 +1139,17 @@ unsafe fn visit_field_level( None } +#[inline] fn raw_of(view: AnyView<'_>) -> TVMFFIAny { *view.as_raw_ffi_any() } +#[inline] fn raw_of_owned(any: &mut Any) -> TVMFFIAny { *any.as_raw_ffi_any() } +#[inline] unsafe fn view_of(raw: &TVMFFIAny) -> AnyView<'_> { unsafe { AnyView::from_raw_ffi_any(*raw) } } diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index 6b98e5cac..16ae9c930 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -187,6 +187,65 @@ fn mutable_dict_is_snapshotted_before_callbacks() { assert_eq!(size, 3); } +#[test] +fn dense_map_layout_is_traversed_completely() { + // More than 4 entries forces the dense (block + iteration list) layout. + let root: Map = (0..9) + .map(|i| (FfiString::from(format!("k{i}")), i as i64)) + .collect(); + let mut sum = 0; + let mut strings = 0; + assert!(walk(&root, |value, phase| { + if phase == Phase::Enter { + if let Some(integer) = value.cast::() { + sum += integer; + } else if value.cast::().is_some() { + strings += 1; + } + } + WalkResult::Advance + }) + .unwrap() + .is_continue()); + assert_eq!(sum, (0..9).sum::()); + assert_eq!(strings, 9); +} + +#[test] +fn interrupt_payload_crosses_map_traversal() { + let root: Map = [(FfiString::from("a"), 1i64), (FfiString::from("b"), 2i64)] + .into_iter() + .collect(); + let outcome = walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::().is_some() { + return WalkResult::interrupt_with(99i64); + } + WalkResult::Advance + }) + .unwrap(); + let ControlFlow::Break(payload) = outcome else { + panic!("map walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(payload).unwrap(), 99); +} + +#[test] +fn handler_error_crosses_map_traversal() { + let root: Map = [(FfiString::from("a"), 1i64)].into_iter().collect(); + let error = match walk(&root, |value, phase| { + if phase == Phase::Enter && value.cast::().is_some() { + Err(runtime_error("map handler failed")) + } else { + Ok(WalkResult::Advance) + } + }) { + Err(error) => error, + Ok(_) => panic!("map handler unexpectedly succeeded"), + }; + assert_eq!(error.message(), "map handler failed"); + assert!(error.backtrace().contains("object `ffi.Map`")); +} + #[test] fn interrupt_stops_without_running_remaining_callbacks() { let root = Array::new(vec![1i64, 2, 3]); From 87dfa895a36c083d765222db2e2c2d0edeb439ec Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Thu, 30 Jul 2026 21:33:20 -0400 Subject: [PATCH 09/20] [DOC] Use native Array constructor in Rust README Signed-off-by: tlopex <820958424@qq.com> --- rust/README.md | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/rust/README.md b/rust/README.md index d13a27bca..5b9a876d3 100644 --- a/rust/README.md +++ b/rust/README.md @@ -37,7 +37,7 @@ visitor. Each value is automatically dispatched to the handler matching its runtime type: ```rust -use tvm_ffi::{dispatch, structural_visit, DefRegionKind, Function, WalkResult}; +use tvm_ffi::{dispatch, structural_visit, Array, DefRegionKind, WalkResult}; #[derive(Default)] struct Calculator { @@ -58,15 +58,12 @@ impl Calculator { } } -let values = Function::get_global("ffi.Array") - .unwrap() - .call_tuple((10_i64, 2.5_f64)) - .unwrap(); +let values = Array::new(vec![10_i64, 2]); let mut calculator = Calculator::default(); assert!(structural_visit(&values, &mut calculator) .unwrap() .is_continue()); -assert_eq!(calculator.value, 7.5); +assert_eq!(calculator.value, 12.0); ``` Typed handlers are tested in source order: borrowed `ObjectCore` node types use From 65c5a688db61d6239276596face0da5a21eb42ac Mon Sep 17 00:00:00 2001 From: yuchuan Date: Fri, 31 Jul 2026 10:29:44 -0400 Subject: [PATCH 10/20] align with c++ side, update the macro. Signed-off-by: yuchuan --- rust/README.md | 130 +- rust/tvm-ffi-macros/Cargo.toml | 1 - rust/tvm-ffi-macros/src/visit.rs | 132 +- rust/tvm-ffi/src/extra/structural_visit.rs | 1229 ++++++++++-------- rust/tvm-ffi/src/lib.rs | 4 +- rust/tvm-ffi/tests/test_dispatch.rs | 49 +- rust/tvm-ffi/tests/test_structural_visit.rs | 518 +++++--- rust/tvm-ffi/tests/test_visitor_alignment.rs | 92 +- 8 files changed, 1251 insertions(+), 904 deletions(-) diff --git a/rust/README.md b/rust/README.md index 5b9a876d3..226e082ae 100644 --- a/rust/README.md +++ b/rust/README.md @@ -30,22 +30,33 @@ efficiency while maintaining interoperability. ## Structural Visitors and Walkers -The `tvm-ffi` crate provides a native Rust structural walker over FFI values, -built-in containers, and reflected object fields. `#[dispatch(visit)]` turns -the `visit_*` methods in an inherent implementation into a typed, stateful -visitor. Each value is automatically dispatched to the handler matching its -runtime type: +The `tvm-ffi` crate provides native Rust structural traversal over FFI values, +built-in containers, and reflected object fields, split into the same two +layers as the C++ API: + +- **Walk layer (observer)** — `structural_walk`, the analog of C++ + `StructuralWalk`: the walker owns recursion; handlers observe each value + and steer traversal by returning a `WalkResult`. +- **Visitor layer (user-driven)** — `structural_visit` with a + `StructuralVisitor`, the analog of a hand-written C++ + `StructuralVisitorObj`: your `visit` method runs for the root and then + controls all recursion itself. + +### Observer walks + +`#[dispatch(visit)]` turns the `visit_*` methods in an inherent +implementation into a typed, stateful observer. Each value is automatically +dispatched to the handler matching its runtime type: ```rust -use tvm_ffi::{dispatch, structural_visit, Array, DefRegionKind, WalkResult}; +use tvm_ffi::{dispatch, structural_walk, Array, WalkOrder, WalkResult}; #[derive(Default)] struct Calculator { - def_region: DefRegionKind, value: f64, } -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] impl Calculator { fn visit_integer(&mut self, value: i64) -> WalkResult { self.value += value as f64; @@ -60,55 +71,98 @@ impl Calculator { let values = Array::new(vec![10_i64, 2]); let mut calculator = Calculator::default(); -assert!(structural_visit(&values, &mut calculator) +assert!(structural_walk(&values, &mut calculator, WalkOrder::PreOrder) .unwrap() - .is_continue()); + .is_none()); assert_eq!(calculator.value, 12.0); ``` Typed handlers are tested in source order: borrowed `ObjectCore` node types use runtime subtype checks, owned arguments use `AnyCompatible` casts, and a final -`&VisitValue` handler acts as a catch-all. Every visitor names a -`DefRegionKind` mirror field through `def_region = `; the walker keeps -that field equal to the definition-region state of the value being dispatched -(and rewinds it after nested traversal), so a handler reads the state through -`self.def_region_kind()` and never writes the field. Use `structural_walk` to -select pre-order or post-order dispatch, or `walk`/`walk_with_context` for raw -callbacks that fire at both `Phase::Enter` and `Phase::Exit` of every value. +`&VisitValue` handler acts as a catch-all. A handler that needs the +definition-region state declares a trailing `DefRegionKind` argument and the +generated dispatch forwards it by arity — the analog of a C++ `StructuralWalk` +callback accepting `(value, def_region_kind)` instead of `(value)`. +`WalkOrder` selects pre-order or post-order dispatch. + +`structural_walk` also accepts a bare closure — a catch-all observer taking +`&VisitValue` with an optional trailing `DefRegionKind`, mirroring the C++ +callback overloads (annotate the closure arguments so the handler shape can +be inferred): + +```rust +use tvm_ffi::{structural_walk, Array, VisitValue, WalkOrder, WalkResult}; + +let values = Array::new(vec![1_i64, 2, 3]); +let mut integers = 0; +assert!(structural_walk( + &values, + |value: &VisitValue| { + if value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, +) +.unwrap() +.is_none()); +assert_eq!(integers, 3); +``` `WalkResult::Advance` visits container or reflected children, `Skip` suppresses the current value's default recursion, and `Interrupt`/`InterruptWith` halt the -walk, surfacing to the caller as `ControlFlow::Break`. Handlers and callbacks -may also return `Result` to propagate errors with `?`. - -A pre-order handler can take over a value's children instead of advancing: -visit selected children with `subvisit`, or delegate the walker's default -child recursion with `subvisit_children`, then return `Skip`. Both inherit -the current definition-region state; the `_with_def_region` variants override -it for exactly that subtree. All of them report a nested interrupt as -`ControlFlow::Break`; propagate it (and errors, via `?`) instead of dropping -the result: +walk. Every traversal returns `Result>`, matching the +C++ `Expected>`: `Ok(None)` means the whole graph was +visited, `Ok(Some(interrupt))` carries the interrupting handler's payload. +Handlers and callbacks may also return `Result` to propagate +errors with `?`. + +### User-driven visitors + +When traversal itself is part of the analysis — visiting selected children, +custom orders, definition-region overrides — implement `StructuralVisitor`. +`visit` receives each value with its definition-region state and descends only +where it chooses: `default_visit_children` delegates the default child +recursion (the analog of C++ `DefaultVisitExpected`), and `visit_child` visits +one child under an explicit state (the analog of `Visit` under +`WithDefRegionKind`): ```rust,ignore -fn visit_func(&mut self, func: &FuncObj) -> Result { - if let ControlFlow::Break(payload) = - self.subvisit_with_def_region(&func.params, DefRegionKind::Recursive)? - { - return Ok(WalkResult::InterruptWith(payload)); +use tvm_ffi::{DefRegionKind, Result, StructuralVisitor, VisitInterrupt, VisitValue}; + +struct FuncVisitor; + +impl StructuralVisitor for FuncVisitor { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + if let Some(func) = value.as_node::() { + // Parameters bind recursively; the body inherits the state. + if let Some(interrupt) = self.visit_child(&func.params, DefRegionKind::Recursive)? { + return Ok(Some(interrupt)); + } + return self.visit_child(&func.body, def_region_kind); + } + self.default_visit_children(value, def_region_kind) } - if let ControlFlow::Break(payload) = self.subvisit(&func.body)? { - return Ok(WalkResult::InterruptWith(payload)); - } - Ok(WalkResult::Skip) } ``` +Returning without descending skips a value's children; +`Ok(Some(VisitInterrupt))` halts the traversal. Nested +`visit_child`/`default_visit_children` calls report a nested interrupt through +their return value — propagate it (and errors, via `?`) instead of dropping +the result. + Recursion runs natively in Rust; no C++ visitor is constructed. Mutable `List`/`Dict` contents are snapshotted before callbacks run, so re-entrant mutation cannot invalidate a traversal. A non-container type that registers a foreign `__s_visit__` hook is rejected rather than silently walked through -reflection; handle it with a matching pre-order handler, `subvisit`, and -`Skip`. +reflection; visit its children explicitly from a `StructuralVisitor`, or skip +it in a walk with a pre-order `WalkResult::Skip` handler. ## Installation diff --git a/rust/tvm-ffi-macros/Cargo.toml b/rust/tvm-ffi-macros/Cargo.toml index c46a78e53..f8d29d406 100644 --- a/rust/tvm-ffi-macros/Cargo.toml +++ b/rust/tvm-ffi-macros/Cargo.toml @@ -28,7 +28,6 @@ license = "Apache-2.0" proc-macro = true [dependencies] -proc-macro-crate = "3" proc-macro2 = "^1.0" quote = "^1.0" syn = { version = "1.0.48", features = ["full", "parsing", "extra-traits"] } diff --git a/rust/tvm-ffi-macros/src/visit.rs b/rust/tvm-ffi-macros/src/visit.rs index 4a0d44a4c..4bad37db0 100644 --- a/rust/tvm-ffi-macros/src/visit.rs +++ b/rust/tvm-ffi-macros/src/visit.rs @@ -18,16 +18,17 @@ */ use proc_macro::TokenStream; -use proc_macro2::{Span, TokenStream as TokenStream2}; -use proc_macro_crate::{crate_name, FoundCrate}; +use proc_macro2::TokenStream as TokenStream2; use quote::{quote, quote_spanned}; use syn::{parse_macro_input, FnArg, ImplItem, ImplItemMethod, ItemImpl, Meta, NestedMeta, 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 _args = parse_macro_input!(attr as DispatchArgs); let item_impl = parse_macro_input!(item as ItemImpl); - match expand(&args, &item_impl) { + match expand(&item_impl) { Ok(generated) => quote!(#item_impl #generated).into(), Err(error) => { let error = error.to_compile_error(); @@ -36,43 +37,25 @@ pub(crate) fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream { } } -struct DispatchArgs { - def_region_field: syn::Ident, -} +struct DispatchArgs; impl syn::parse::Parse for DispatchArgs { fn parse(input: syn::parse::ParseStream) -> syn::Result { let mode: syn::Ident = input.parse()?; if mode != "visit" { - return Err(syn::Error::new( - mode.span(), - "expected `dispatch(visit, def_region = )`", - )); - } - if input.parse::().is_err() { - return Err(syn::Error::new( - mode.span(), - "`dispatch(visit)` requires `def_region = ` naming the visitor's \ - `DefRegionKind` mirror field", - )); - } - let key: syn::Ident = input.parse()?; - if key != "def_region" { - return Err(syn::Error::new( - key.span(), - "expected `def_region = `", - )); + return Err(syn::Error::new(mode.span(), "expected `dispatch(visit)`")); } - input.parse::()?; - let def_region_field: syn::Ident = input.parse()?; if !input.is_empty() { - return Err(input.error("unexpected tokens after `def_region = `")); + return Err(input.error( + "`dispatch(visit)` takes no further arguments; a handler that needs the \ + definition-region state declares a trailing `DefRegionKind` argument", + )); } - Ok(DispatchArgs { def_region_field }) + Ok(DispatchArgs) } } -fn expand(args: &DispatchArgs, item_impl: &ItemImpl) -> syn::Result { +fn expand(item_impl: &ItemImpl) -> syn::Result { if item_impl.trait_.is_some() { return Err(syn::Error::new_spanned( item_impl, @@ -97,17 +80,24 @@ fn expand(args: &DispatchArgs, item_impl: &ItemImpl) -> syn::Result quote! { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(value) + self.#method(value #kind_arg) ) ); }, @@ -115,7 +105,7 @@ fn expand(args: &DispatchArgs, item_impl: &ItemImpl) -> syn::Result() { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(node) + self.#method(node #kind_arg) ) ); } @@ -124,7 +114,7 @@ fn expand(args: &DispatchArgs, item_impl: &ItemImpl) -> syn::Result() { return Some( #tvm_ffi::extra::structural_visit::IntoVisitResult::into_visit_result( - self.#method(node) + self.#method(node #kind_arg) ) ); } @@ -171,57 +161,20 @@ fn expand(args: &DispatchArgs, item_impl: &ItemImpl) -> syn::Result Option<#tvm_ffi::extra::structural_visit::VisitResult> { #(#links)* None } - - fn def_region_kind(&self) -> #tvm_ffi::extra::structural_visit::DefRegionKind { - self.#def_region_field - } - - #[doc(hidden)] - fn def_region_slot( - &mut self, - ) -> &mut #tvm_ffi::extra::structural_visit::DefRegionKind { - &mut self.#def_region_field - } - } - - #(#[#impl_cfg_attrs])* - impl #impl_generics #self_type #where_clause { - /// Definition-region state at the value currently being dispatched. - /// - /// Inherent mirror of `VisitDispatch::def_region_kind`, callable - /// without importing the trait. - #[allow(dead_code)] - fn def_region_kind(&self) -> #tvm_ffi::extra::structural_visit::DefRegionKind { - self.#def_region_field - } } }) } -fn resolve_tvm_ffi_crate() -> syn::Result { - crate_name("tvm-ffi") - .map(crate_path) - .map_err(|error| syn::Error::new(Span::call_site(), error)) -} - -fn crate_path(found: FoundCrate) -> TokenStream2 { - match found { - FoundCrate::Itself => quote!(crate), - FoundCrate::Name(name) => { - let name = syn::parse_str::(&name) - .unwrap_or_else(|_| syn::Ident::new_raw(&name, Span::call_site())); - quote!(::#name) - } - } -} - struct Handler { method: syn::Ident, argument: HandlerArgument, + /// The handler declared a trailing `DefRegionKind` argument. + wants_def_region: bool, cfg_attrs: Vec, } @@ -238,13 +191,14 @@ fn parse_handler(method: &ImplItemMethod) -> syn::Result { Some(FnArg::Receiver(receiver)) if receiver.reference.is_some() && receiver.mutability.is_some() ); - if !receiver_is_mut || inputs.len() != 2 { + if !receiver_is_mut || !(inputs.len() == 2 || inputs.len() == 3) { return Err(syn::Error::new_spanned( &method.sig, - "visit handlers must take `&mut self` and a node; read definition-region state \ - through `self.def_region_kind()`", + "visit handlers must take `&mut self`, a node, and optionally a trailing \ + `DefRegionKind` argument", )); } + let wants_def_region = inputs.len() == 3; let value_type = match inputs.iter().nth(1) { Some(FnArg::Typed(value)) => (*value.ty).clone(), @@ -270,6 +224,7 @@ fn parse_handler(method: &ImplItemMethod) -> syn::Result { Ok(Handler { method: method.sig.ident.clone(), argument, + wants_def_region, cfg_attrs, }) } @@ -322,24 +277,3 @@ fn is_visit_value(value_type: &Type) -> bool { .last() .is_some_and(|segment| segment.ident == "VisitValue") } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn renamed_dependency_uses_its_imported_name() { - assert_eq!( - crate_path(FoundCrate::Name("renamed_tvm_ffi".to_string())).to_string(), - ":: renamed_tvm_ffi" - ); - } - - #[test] - fn keyword_dependency_uses_a_raw_identifier() { - assert_eq!( - crate_path(FoundCrate::Name("type".to_string())).to_string(), - ":: r#type" - ); - } -} diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index 55fc261ce..6fccdd4ed 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -19,27 +19,41 @@ //! Native Rust structural visiting. //! -//! This module separates the two jobs involved in a visit: +//! Two public layers mirror the C++ API split: //! -//! * [`VisitValue`] provides borrowed matching for generated Rust dispatch. -//! * `NativeWalker` owns recursion through containers and reflected fields. +//! * [`StructuralVisitor`] + [`structural_visit`] — the visitor drives +//! recursion itself, like a hand-written C++ `StructuralVisitorObj`: +//! [`StructuralVisitor::visit`] runs once per reached value and descends +//! only where it calls [`StructuralVisitor::default_visit_children`] or +//! [`StructuralVisitor::visit_child`]. +//! * [`VisitDispatch`] + [`structural_walk`] — observer callbacks, like C++ +//! `StructuralWalk`: the walker recurses on its own and callbacks steer it +//! through the returned [`WalkResult`] (advance, skip, interrupt). +//! +//! Both layers thread the definition-region state explicitly: walk handlers +//! opt in with a trailing [`DefRegionKind`] argument, and a visitor receives +//! and forwards it when descending. +//! +//! Underneath both, [`VisitValue`] provides borrowed matching for typed Rust +//! dispatch and the stateless recursion engine (`visit_raw` and the +//! `visit_*` helpers below) owns iteration over containers and reflected +//! fields. //! //! The runtime object registry is open, so the walker uses the stable tvm-ffi //! reflection ABI for arbitrary registered object types. That ABI is only the //! object-description boundary: traversal, control flow, typed dispatch, //! visitor state, and definition-region propagation remain in Rust. //! -//! A Rust handler may override a type's children by visiting them through -//! [`VisitDispatch::subvisit`] and returning [`WalkResult::Skip`]. No C++ -//! `ffi.StructuralVisitor` is constructed and no C++ default-visit function is -//! called. A non-container type with a foreign `__s_visit__` hook must be -//! handled this way; advancing into its default children is rejected instead -//! of silently substituting reflection with potentially different semantics. +//! No C++ `ffi.StructuralVisitor` is constructed and no C++ default-visit +//! function is called. A non-container type with a foreign `__s_visit__` hook +//! is rejected instead of silently substituting reflection with potentially +//! different semantics; visit such a type's children explicitly from a +//! [`StructuralVisitor`], or skip the value in a walk. use std::ops::ControlFlow; use std::os::raw::c_void; use std::ptr::NonNull; -use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; use crate::any::{Any, AnyView}; use crate::error::{Error, Result, RUNTIME_ERROR, TYPE_ERROR}; @@ -98,15 +112,6 @@ impl IntoVisitResult for Result { } } -/// Whether a callback runs before or after a value's children. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Phase { - /// Before the value's children. - Enter, - /// After the value's children. - Exit, -} - /// Callback order for [`structural_walk`]. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum WalkOrder { @@ -146,11 +151,25 @@ const _: () = { ); }; -/// Result of a completed Rust walk. +/// Interrupt state of a traversal, mirroring C++ `ffi.VisitInterrupt`. /// -/// `Continue(())` means the whole graph was visited. `Break(payload)` means a -/// handler interrupted it; a payload-less interrupt carries `ffi::None`. -pub type VisitOutcome = ControlFlow; +/// Entry points and visitor-layer calls return +/// `Result>`: `Ok(None)` means the (sub)graph was +/// visited completely, `Ok(Some(..))` means a handler halted the traversal +/// with this interrupt, and `Err` means it failed. +pub struct VisitInterrupt { + /// Payload returned with the interrupt, or FFI `None` for no payload. + pub value: Any, +} + +impl VisitInterrupt { + /// Interrupt carrying an FFI-compatible payload. + pub fn with>(payload: T) -> Self { + Self { + value: payload.into(), + } + } +} /// Fallible result returned by generated typed dispatch. #[doc(hidden)] @@ -223,116 +242,255 @@ impl From for NativeHalt { type NativeResult = std::result::Result<(), NativeHalt>; -/// Typed dispatch implemented by the visitor object itself. +/// Typed dispatch implemented by a walk-layer observer. /// /// [`crate::dispatch`] tests the implementation's `visit_*` methods in source /// order. Borrowed node arguments use refcount-free subtype checks, owned /// FFI-compatible arguments use exact value casts, and `&VisitValue` is a /// catch-all. `None` asks the Rust walker to continue normally. /// -/// Every visitor carries a definition-region mirror field, named through -/// `#[dispatch(visit, def_region = )]`. The walker writes the field -/// before each dispatched handler and restores it afterwards, so during a -/// handler [`VisitDispatch::def_region_kind`] always reports the state at -/// the current value — including after nested [`VisitDispatch::subvisit`] -/// calls. Treat the field as read-only; the walker's own propagation never -/// reads it, so overwriting it only misleads your own reads and the -/// kind-inheriting `subvisit` forms. +/// This is the observer layer, mirroring C++ `StructuralWalk` callbacks: the +/// walker owns recursion, and a handler steers it only through the returned +/// [`WalkResult`]. A traversal that must visit children itself — selected +/// children, custom orders, explicit definition-region overrides — belongs in +/// a [`StructuralVisitor`] instead. +/// +/// The definition-region state active at the dispatched value arrives as the +/// `def_region_kind` argument. A `#[dispatch(visit)]` handler opts into it by +/// declaring a trailing `DefRegionKind` parameter — the analog of a C++ +/// `StructuralWalk` callback accepting `(value, def_region_kind)` instead of +/// `(value)`. pub trait VisitDispatch: Sized { - fn dispatch_visit(&mut self, value: &VisitValue) -> Option; + fn dispatch_visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option; +} - /// Return the definition-region state at the value being dispatched. - /// - /// Meaningful only while a handler is running; outside a walk the mirror - /// holds its initial or last-restored value. - fn def_region_kind(&self) -> DefRegionKind; +impl VisitDispatch for &mut V { + #[inline] + fn dispatch_visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option { + (**self).dispatch_visit(value, def_region_kind) + } +} - /// Mirror storage refreshed by the walker around each dispatch. +/// Conversion into the walker argument of [`structural_walk`]. +/// +/// The `Marker` parameter lets one entry point accept several handler +/// shapes without overlapping implementations — the Rust equivalent of the +/// C++ `StructuralWalk` callback overload set: +/// +/// * `&mut V` where `V: VisitDispatch` — a stateful typed visitor +/// (`#[dispatch(visit)]` or hand-written). +/// * `FnMut(&VisitValue) -> impl IntoVisitResult` — a bare observer closure, +/// the analog of a C++ `(value)` callback. +/// * `FnMut(&VisitValue, DefRegionKind) -> impl IntoVisitResult` — the +/// analog of a C++ `(value, def_region_kind)` callback. +/// +/// Closure arguments usually need explicit type annotations +/// (`|value: &VisitValue| ...`) for the marker to be inferred. +pub trait IntoWalker { #[doc(hidden)] - fn def_region_slot(&mut self) -> &mut DefRegionKind; + type Walker: NativeVisit; + #[doc(hidden)] + fn into_walker(self, order: WalkOrder) -> Self::Walker; +} - /// Visit `child` immediately, inheriting the current definition-region - /// state. - /// - /// This is how a pre-order handler takes over a value's children: visit - /// each selected child, then return [`WalkResult::Skip`]. Use - /// [`VisitDispatch::subvisit_with_def_region`] to override the state for - /// exactly this subtree. The nested traversal always dispatches - /// pre-order. - /// - /// `Err` carries a nested handler failure and should be propagated with - /// `?`. `Ok(ControlFlow::Break(payload))` reports a nested interrupt; - /// return [`WalkResult::InterruptWith`] with the payload to keep - /// halting. Dropping the result silently swallows both. - fn subvisit(&mut self, child: &T) -> Result - where - for<'x> AnyView<'x>: From<&'x T>, - { - let def_region_kind = self.def_region_kind(); - self.subvisit_with_def_region(child, def_region_kind) +#[doc(hidden)] +pub enum ByDispatch {} + +impl<'a, V: VisitDispatch> IntoWalker for &'a mut V { + type Walker = DispatchVisitor<&'a mut V>; + fn into_walker(self, order: WalkOrder) -> Self::Walker { + DispatchVisitor { + visitor: self, + order, + } } +} - /// Visit `child` immediately under an explicitly selected - /// definition-region state. - /// - /// The override is scoped to this recursive call; see - /// [`VisitDispatch::subvisit`] for the result contract. - fn subvisit_with_def_region( +/// Runs a catch-all closure at the phase selected by `order` — the closure +/// analog of `DispatchVisitor`, without the `Option` +/// no-handler-matched layer a dispatch chain needs. (Routing closures +/// through `DispatchVisitor` instead measures ~10-20% slower on the bare +/// closure walk: the wrapped-and-unwrapped `Option>` does not +/// fold away.) +#[doc(hidden)] +pub struct ClosureWalker { + callback: F, + order: WalkOrder, +} + +impl NativeVisit for ClosureWalker +where + F: FnMut(&VisitValue) -> O, + O: IntoVisitResult, +{ + fn enter(&mut self, value: &VisitValue, _def_region_kind: DefRegionKind) -> Result { + match self.order { + WalkOrder::PreOrder => (self.callback)(value).into_visit_result(), + WalkOrder::PostOrder => Ok(WalkResult::Advance), + } + } + + fn exit(&mut self, value: &VisitValue, _def_region_kind: DefRegionKind) -> Result { + match self.order { + WalkOrder::PreOrder => Ok(WalkResult::Advance), + WalkOrder::PostOrder => (self.callback)(value).into_visit_result(), + } + } +} + +#[doc(hidden)] +pub enum ByValueClosure {} + +impl IntoWalker for F +where + F: FnMut(&VisitValue) -> O, + O: IntoVisitResult, +{ + type Walker = ClosureWalker; + fn into_walker(self, order: WalkOrder) -> Self::Walker { + ClosureWalker { + callback: self, + order, + } + } +} + +/// `ClosureWalker` variant whose callback also receives the definition-region +/// state. +#[doc(hidden)] +pub struct ClosureKindWalker { + callback: F, + order: WalkOrder, +} + +impl NativeVisit for ClosureKindWalker +where + F: FnMut(&VisitValue, DefRegionKind) -> O, + O: IntoVisitResult, +{ + fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { + match self.order { + WalkOrder::PreOrder => (self.callback)(value, def_region_kind).into_visit_result(), + WalkOrder::PostOrder => Ok(WalkResult::Advance), + } + } + + fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { + match self.order { + WalkOrder::PreOrder => Ok(WalkResult::Advance), + WalkOrder::PostOrder => (self.callback)(value, def_region_kind).into_visit_result(), + } + } +} + +#[doc(hidden)] +pub enum ByValueKindClosure {} + +impl IntoWalker for F +where + F: FnMut(&VisitValue, DefRegionKind) -> O, + O: IntoVisitResult, +{ + type Walker = ClosureKindWalker; + fn into_walker(self, order: WalkOrder) -> Self::Walker { + ClosureKindWalker { + callback: self, + order, + } + } +} + +/// A visitor that drives recursion itself, mirroring C++ +/// `StructuralVisitorObj`. +/// +/// [`structural_visit`] calls [`StructuralVisitor::visit`] for the root; +/// after that the visitor is in control, exactly like a C++ visitor whose +/// vtable `visit` runs per value. A `visit` implementation descends only +/// where it chooses: +/// +/// * [`StructuralVisitor::default_visit_children`] delegates the default +/// child recursion — the analog of C++ +/// `StructuralVisitorObj::DefaultVisitExpected`. +/// * [`StructuralVisitor::visit_child`] visits one selected child — the +/// analog of C++ `visitor->Visit(child)`, with the explicit +/// `def_region_kind` argument playing the role of `WithDefRegionKind`. +/// +/// Returning without descending skips the value's children. There is no +/// [`WalkResult`] at this layer: control flow is what the implementation +/// visits, and `Ok(Some(interrupt))` halts the traversal — the analog of +/// returning a C++ `VisitInterrupt`. Nested `visit_child` and +/// `default_visit_children` calls report a nested interrupt through their +/// return value; propagate it (and errors, via `?`) upward instead of +/// dropping the result. +/// +/// The definition-region state is threaded explicitly, exactly like walk +/// handlers that declare the trailing argument: `visit` receives the state +/// active at the value and forwards it — or an override — when descending. +/// Reflected-field annotations override the forwarded state automatically +/// inside `default_visit_children`. +pub trait StructuralVisitor: Sized { + /// Visit one value under the definition-region state active at it. + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result>; + + /// Visit `child` now under `def_region_kind`, dispatching back into + /// [`StructuralVisitor::visit`]. An FFI `None` child is skipped without + /// a callback, matching the walk layer. + #[inline] + fn visit_child( &mut self, child: &T, def_region_kind: DefRegionKind, - ) -> Result + ) -> Result> where for<'x> AnyView<'x>: From<&'x T>, { - let walker = NativeWalker::new(); - let mut dispatch = DispatchVisitor { - visitor: self, - order: WalkOrder::PreOrder, - }; - finish(walker.visit_raw(raw_of(AnyView::from(child)), &mut dispatch, def_region_kind)) + let raw = raw_of(AnyView::from(child)); + if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 { + return Ok(None); + } + self.visit(&VisitValue::from_raw(raw), def_region_kind) } - /// Visit `value`'s children — not `value` itself — with the walker's - /// default rules, inheriting the current definition-region state. + /// Visit `value`'s children — not `value` itself — with the default + /// rules, dispatching each child back into [`StructuralVisitor::visit`]. /// /// Children are container contents for `Array`/`List`/`Map`/`Dict` and - /// reflected structural fields otherwise. This is the Rust analog of C++ - /// `StructuralVisitorObj::DefaultVisitExpected`: a handler may run its - /// enter logic, delegate the default child recursion explicitly, run its - /// exit logic with the same locals in scope, and return - /// [`WalkResult::Skip`]. Unlike [`VisitDispatch::subvisit`] it needs no - /// knowledge of the value's concrete type, so it also works from a - /// `&VisitValue` catch-all handler. - /// - /// The result contract matches [`VisitDispatch::subvisit`]. - fn subvisit_children(&mut self, value: &VisitValue) -> Result { - let def_region_kind = self.def_region_kind(); - self.subvisit_children_with_def_region(value, def_region_kind) - } - - /// Visit `value`'s children with the walker's default rules under an - /// explicitly selected definition-region state. - /// - /// See [`VisitDispatch::subvisit_children`]. - fn subvisit_children_with_def_region( + /// reflected structural fields otherwise. Field annotations override + /// `def_region_kind` for that field's recursive visit exactly like the + /// walk layer. + #[inline] + fn default_visit_children( &mut self, value: &VisitValue, def_region_kind: DefRegionKind, - ) -> Result { - let walker = NativeWalker::new(); - let mut dispatch = DispatchVisitor { - visitor: self, - order: WalkOrder::PreOrder, - }; - let result = walker - .visit_children_raw(value.0, &mut dispatch, def_region_kind) - .map_err(|halt| NativeWalker::with_value_context(halt, value.0)); + ) -> Result> { + let result = visit_children_raw( + value.0, + &mut UserChildren { visitor: self }, + def_region_kind, + ) + .map_err(|halt| with_value_context(halt, value.0)); finish(result) } } -trait NativeVisit { +/// Internal per-value protocol driven by the recursion engine. Public only +/// as the bound of [`IntoWalker::Walker`]; not meant to be implemented +/// outside this crate. +#[doc(hidden)] +pub trait NativeVisit { fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result; fn exit(&mut self, _value: &VisitValue, _def_region_kind: DefRegionKind) -> Result { @@ -340,36 +498,22 @@ trait NativeVisit { } } -struct DispatchVisitor<'a, V> { - visitor: &'a mut V, +/// Owns its walker so a closure's state stays inline and a `&mut` visitor +/// keeps a single level of indirection. Public only as an +/// [`IntoWalker::Walker`] projection. +#[doc(hidden)] +pub struct DispatchVisitor { + visitor: V, order: WalkOrder, } -impl DispatchVisitor<'_, V> { - /// Run one typed dispatch with the mirror scoped to `def_region_kind`. - /// - /// The save/restore pair keeps the mirror equal to the dispatched value's - /// state for the whole handler call, and transparently rewinds it after - /// nested `subvisit` recursion. - fn dispatch_scoped( - &mut self, - value: &VisitValue, - def_region_kind: DefRegionKind, - ) -> Result { - let saved = std::mem::replace(self.visitor.def_region_slot(), def_region_kind); - let result = self - .visitor - .dispatch_visit(value) - .unwrap_or(Ok(WalkResult::Advance)); - *self.visitor.def_region_slot() = saved; - result - } -} - -impl NativeVisit for DispatchVisitor<'_, V> { +impl NativeVisit for DispatchVisitor { fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { match self.order { - WalkOrder::PreOrder => self.dispatch_scoped(value, def_region_kind), + WalkOrder::PreOrder => self + .visitor + .dispatch_visit(value, def_region_kind) + .unwrap_or(Ok(WalkResult::Advance)), WalkOrder::PostOrder => Ok(WalkResult::Advance), } } @@ -377,472 +521,454 @@ impl NativeVisit for DispatchVisitor<'_, V> { fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { match self.order { WalkOrder::PreOrder => Ok(WalkResult::Advance), - WalkOrder::PostOrder => self.dispatch_scoped(value, def_region_kind), + WalkOrder::PostOrder => self + .visitor + .dispatch_visit(value, def_region_kind) + .unwrap_or(Ok(WalkResult::Advance)), } } } -struct CallbackVisitor(F); +/// Per-child action invoked by the shared child-iteration engine. +/// +/// The engine owns *finding* the children (container contents, reflected +/// fields) and computing each child's definition-region state; this trait +/// decides what happens at a child. The walk layer recurses +/// ([`WalkChildren`]); the visitor layer hands the child straight to user +/// code ([`UserChildren`]). +trait ChildVisit { + fn visit_child(&mut self, child: TVMFFIAny, def_region_kind: DefRegionKind) -> NativeResult; +} -impl NativeVisit for CallbackVisitor -where - F: FnMut(&VisitValue, Phase, DefRegionKind) -> O, - O: IntoVisitResult, -{ - fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { - (self.0)(value, Phase::Enter, def_region_kind).into_visit_result() - } +/// Walk-layer recursion: every child re-enters [`visit_raw`]. +struct WalkChildren<'a, V> { + visitor: &'a mut V, +} - fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { - (self.0)(value, Phase::Exit, def_region_kind).into_visit_result() +impl ChildVisit for WalkChildren<'_, V> { + fn visit_child(&mut self, child: TVMFFIAny, def_region_kind: DefRegionKind) -> NativeResult { + visit_raw(child, self.visitor, def_region_kind) } } -/// Stateless Rust recursion engine. -struct NativeWalker { - structural_visit: Option, +/// Visitor-layer dispatch: every child goes back into the user-driven +/// [`StructuralVisitor::visit`], which controls further descent itself. +struct UserChildren<'a, V> { + visitor: &'a mut V, } -impl NativeWalker { - fn new() -> Self { - Self { - structural_visit: type_attr_column(STRUCTURAL_VISIT_ATTR), - } - } - - fn visit_raw( - &self, - value: TVMFFIAny, - visitor: &mut V, - def_region_kind: DefRegionKind, - ) -> NativeResult { - if value.type_index == TVMFFITypeIndex::kTVMFFINone as i32 { +impl ChildVisit for UserChildren<'_, V> { + #[inline] + fn visit_child(&mut self, child: TVMFFIAny, def_region_kind: DefRegionKind) -> NativeResult { + if child.type_index == TVMFFITypeIndex::kTVMFFINone as i32 { return Ok(()); } - - let visit_value = VisitValue::from_raw(value); - let enter = match visitor.enter(&visit_value, def_region_kind) { - Ok(flow) => flow, - Err(error) => return Err(Self::with_value_context(error.into(), value)), - }; - match enter { - WalkResult::Advance => {} - WalkResult::Skip => return Ok(()), - WalkResult::Interrupt => return Err(NativeHalt::Interrupt(Any::new())), - WalkResult::InterruptWith(payload) => return Err(NativeHalt::Interrupt(payload)), + match self + .visitor + .visit(&VisitValue::from_raw(child), def_region_kind) + { + Ok(None) => Ok(()), + Ok(Some(interrupt)) => Err(NativeHalt::Interrupt(interrupt.value)), + Err(error) => Err(NativeHalt::Error(error)), } + } +} - if let Err(halt) = self.visit_children_raw(value, visitor, def_region_kind) { - return Err(Self::with_value_context(halt, value)); - } +/// Recurse into `value` on behalf of `visitor`: fire its enter hook, walk the +/// children, fire its exit hook. The engine below is stateless — these are +/// free functions, with the only shared piece (the `__s_visit__` attribute +/// column) cached process-wide. +fn visit_raw( + value: TVMFFIAny, + visitor: &mut V, + def_region_kind: DefRegionKind, +) -> NativeResult { + if value.type_index == TVMFFITypeIndex::kTVMFFINone as i32 { + return Ok(()); + } - let exit = match visitor.exit(&visit_value, def_region_kind) { - Ok(flow) => flow, - Err(error) => return Err(Self::with_value_context(error.into(), value)), - }; - match exit { - WalkResult::Interrupt => Err(NativeHalt::Interrupt(Any::new())), - WalkResult::InterruptWith(payload) => Err(NativeHalt::Interrupt(payload)), - WalkResult::Advance | WalkResult::Skip => Ok(()), - } + let visit_value = VisitValue::from_raw(value); + let enter = match visitor.enter(&visit_value, def_region_kind) { + Ok(flow) => flow, + Err(error) => return Err(with_value_context(error.into(), value)), + }; + match enter { + WalkResult::Advance => {} + WalkResult::Skip => return Ok(()), + WalkResult::Interrupt => return Err(NativeHalt::Interrupt(Any::new())), + WalkResult::InterruptWith(payload) => return Err(NativeHalt::Interrupt(payload)), } - #[inline] - fn visit_children_raw( - &self, - value: TVMFFIAny, - visitor: &mut V, - def_region_kind: DefRegionKind, - ) -> NativeResult { - match value.type_index { - x if x == TVMFFITypeIndex::kTVMFFIArray as i32 - || x == TVMFFITypeIndex::kTVMFFIList as i32 => - { - return self.visit_sequence(value, visitor, def_region_kind); - } - x if x == TVMFFITypeIndex::kTVMFFIMap as i32 - || x == TVMFFITypeIndex::kTVMFFIDict as i32 => - { - // Fast path: read the MapBaseObj storage layout directly, like - // the SeqPrefix path for arrays — zero FFI calls per entry. - // Dict entries are snapshotted first to keep the re-entrant - // mutation guard. If the one-time layout validation fails - // (e.g. an ABI-debug build), fall back to the packed-functor - // iteration protocol. - if map_layout_usable(value) { - let snapshot = x == TVMFFITypeIndex::kTVMFFIDict as i32; - return self.visit_map_layout(value, visitor, def_region_kind, snapshot); - } - return self.visit_map(value, visitor, def_region_kind); - } - _ => {} - } + let children = &mut WalkChildren { + visitor: &mut *visitor, + }; + if let Err(halt) = visit_children_raw(value, children, def_region_kind) { + return Err(with_value_context(halt, value)); + } - self.reject_foreign_structural_visit(value.type_index)?; - if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { - Ok(()) - } else { - self.visit_reflected_fields(value, visitor, def_region_kind) - } + let exit = match visitor.exit(&visit_value, def_region_kind) { + Ok(flow) => flow, + Err(error) => return Err(with_value_context(error.into(), value)), + }; + match exit { + WalkResult::Interrupt => Err(NativeHalt::Interrupt(Any::new())), + WalkResult::InterruptWith(payload) => Err(NativeHalt::Interrupt(payload)), + WalkResult::Advance | WalkResult::Skip => Ok(()), } +} - #[inline(never)] - fn visit_sequence( - &self, - value: TVMFFIAny, - visitor: &mut V, - def_region_kind: DefRegionKind, - ) -> NativeResult { - let seq = unsafe { &*(value.data_union.v_obj as *const SeqPrefix) }; - if seq.size < 0 { - return Err(runtime_error("native visitor: sequence reports a negative size").into()); - } - if seq.data.is_null() && seq.size != 0 { - return Err(runtime_error( - "native visitor: non-empty sequence has a null data pointer", - ) - .into()); - } - let size = usize::try_from(seq.size) - .map_err(|_| runtime_error("native visitor: sequence size does not fit usize"))?; - if size == 0 { - return Ok(()); +#[inline] +fn visit_children_raw( + value: TVMFFIAny, + visitor: &mut C, + def_region_kind: DefRegionKind, +) -> NativeResult { + match value.type_index { + x if x == TVMFFITypeIndex::kTVMFFIArray as i32 + || x == TVMFFITypeIndex::kTVMFFIList as i32 => + { + return visit_sequence(value, visitor, def_region_kind); } - - if value.type_index == TVMFFITypeIndex::kTVMFFIList as i32 { - // List storage may be invalidated by a re-entrant callback. Own a - // snapshot before running the first callback. - let children: Vec = { - let cells = unsafe { std::slice::from_raw_parts(seq.data, size) }; - cells - .iter() - .map(|cell| Any::from(unsafe { view_of(cell) })) - .collect() - }; - for (index, mut child) in children.into_iter().enumerate() { - let raw = raw_of_owned(&mut child); - self.visit_raw(raw, visitor, def_region_kind) - .map_err(|halt| { - with_error_context(halt, &format!("sequence item [{index}]")) - })?; + x if x == TVMFFITypeIndex::kTVMFFIMap as i32 + || x == TVMFFITypeIndex::kTVMFFIDict as i32 => + { + // Fast path: read the MapBaseObj storage layout directly, like + // the SeqPrefix path for arrays — zero FFI calls per entry. + // Dict entries are snapshotted first to keep the re-entrant + // mutation guard. If the one-time layout validation fails + // (e.g. an ABI-debug build), fall back to the packed-functor + // iteration protocol. + if map_layout_usable(value) { + let snapshot = x == TVMFFITypeIndex::kTVMFFIDict as i32; + return visit_map_layout(value, visitor, def_region_kind, snapshot); } - return Ok(()); + return visit_map(value, visitor, def_region_kind); } + _ => {} + } - // Array is immutable, so its element cells remain stable throughout - // recursive callbacks and need no refcounted snapshot. - let cells = unsafe { std::slice::from_raw_parts(seq.data, size) }; - for (index, child) in cells.iter().enumerate() { - self.visit_raw(*child, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("sequence item [{index}]")))?; - } + reject_foreign_structural_visit(value.type_index)?; + if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { Ok(()) + } else { + visit_reflected_fields(value, visitor, def_region_kind) + } +} + +#[inline(never)] +fn visit_sequence( + value: TVMFFIAny, + visitor: &mut C, + def_region_kind: DefRegionKind, +) -> NativeResult { + let seq = unsafe { &*(value.data_union.v_obj as *const SeqPrefix) }; + if seq.size < 0 { + return Err(runtime_error("native visitor: sequence reports a negative size").into()); + } + if seq.data.is_null() && seq.size != 0 { + return Err( + runtime_error("native visitor: non-empty sequence has a null data pointer").into(), + ); + } + let size = usize::try_from(seq.size) + .map_err(|_| runtime_error("native visitor: sequence size does not fit usize"))?; + if size == 0 { + return Ok(()); } - /// Walk map/dict entries by reading the `MapBaseObj` storage directly — - /// the map analog of the `SeqPrefix` array fast path. `snapshot` first - /// takes owned copies of all entries (Dict re-entrant mutation guard). - #[inline(never)] - fn visit_map_layout( - &self, - value: TVMFFIAny, - visitor: &mut V, - def_region_kind: DefRegionKind, - snapshot: bool, - ) -> NativeResult { - let map = unsafe { &*(value.data_union.v_obj as *const MapPrefix) }; - let size = map.size as usize; - if size == 0 { - return Ok(()); - } - let mut cursor = unsafe { MapCursor::new(map) }; - - if snapshot { - let mut entries: Vec<(Any, Any)> = Vec::with_capacity(size); - for _ in 0..size { - let Some((key, val)) = (unsafe { cursor.next() }) else { - return Err(runtime_error("native visitor: map iteration ended early").into()); - }; - entries.push(( - Any::from(unsafe { view_of(&key) }), - Any::from(unsafe { view_of(&val) }), - )); - } - for (index, (mut key, mut val)) in entries.into_iter().enumerate() { - let key_raw = raw_of_owned(&mut key); - self.visit_raw(key_raw, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("dict key [{index}]")))?; - let val_raw = raw_of_owned(&mut val); - self.visit_raw(val_raw, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("dict value [{index}]")))?; - } - return Ok(()); + if value.type_index == TVMFFITypeIndex::kTVMFFIList as i32 { + // List storage may be invalidated by a re-entrant callback. Own a + // snapshot before running the first callback. + let children: Vec = { + let cells = unsafe { std::slice::from_raw_parts(seq.data, size) }; + cells + .iter() + .map(|cell| Any::from(unsafe { view_of(cell) })) + .collect() + }; + for (index, child) in children.into_iter().enumerate() { + let raw = raw_of_owned(&child); + visitor + .visit_child(raw, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("sequence item [{index}]")))?; } + return Ok(()); + } + + // Array is immutable, so its element cells remain stable throughout + // recursive callbacks and need no refcounted snapshot. + let cells = unsafe { std::slice::from_raw_parts(seq.data, size) }; + for (index, child) in cells.iter().enumerate() { + visitor + .visit_child(*child, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("sequence item [{index}]")))?; + } + Ok(()) +} + +/// Walk map/dict entries by reading the `MapBaseObj` storage directly — +/// the map analog of the `SeqPrefix` array fast path. `snapshot` first +/// takes owned copies of all entries (Dict re-entrant mutation guard). +#[inline(never)] +fn visit_map_layout( + value: TVMFFIAny, + visitor: &mut C, + def_region_kind: DefRegionKind, + snapshot: bool, +) -> NativeResult { + let map = unsafe { &*(value.data_union.v_obj as *const MapPrefix) }; + let size = map.size as usize; + if size == 0 { + return Ok(()); + } + let mut cursor = unsafe { MapCursor::new(map) }; - // Immutable map: entry cells stay stable throughout recursive - // callbacks, so visit them in place. The `size` bound also guards the - // dense iteration list against corruption-induced cycles. - for index in 0..size { + if snapshot { + let mut entries: Vec<(Any, Any)> = Vec::with_capacity(size); + for _ in 0..size { let Some((key, val)) = (unsafe { cursor.next() }) else { return Err(runtime_error("native visitor: map iteration ended early").into()); }; - self.visit_raw(key, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("map key [{index}]")))?; - self.visit_raw(val, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("map value [{index}]")))?; + entries.push(( + Any::from(unsafe { view_of(&key) }), + Any::from(unsafe { view_of(&val) }), + )); } - Ok(()) + for (index, (key, val)) in entries.into_iter().enumerate() { + visitor + .visit_child(raw_of_owned(&key), def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("dict key [{index}]")))?; + visitor + .visit_child(raw_of_owned(&val), def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("dict value [{index}]")))?; + } + return Ok(()); } - fn visit_map( - &self, - value: TVMFFIAny, - visitor: &mut V, - def_region_kind: DefRegionKind, - ) -> NativeResult { - // Map storage is private C++. The Rust binding itself uses these public - // iterator functors; using them here does not invoke structural - // visiting or transfer traversal control out of Rust. - let is_dict = value.type_index == TVMFFITypeIndex::kTVMFFIDict as i32; - let (size_name, iter_name) = if is_dict { - ("ffi.DictSize", "ffi.DictForwardIterFunctor") - } else { - ("ffi.MapSize", "ffi.MapForwardIterFunctor") + // Immutable map: entry cells stay stable throughout recursive + // callbacks, so visit them in place. The `size` bound also guards the + // dense iteration list against corruption-induced cycles. + for index in 0..size { + let Some((key, val)) = (unsafe { cursor.next() }) else { + return Err(runtime_error("native visitor: map iteration ended early").into()); }; - let size = Function::get_global(size_name)? - .call_packed(&[unsafe { view_of(&value) }]) - .and_then(i64::try_from)?; - if size < 0 { - return Err(runtime_error("native visitor: map reports a negative size").into()); - } - let size = usize::try_from(size) - .map_err(|_| runtime_error("native visitor: map size does not fit usize"))?; - if size == 0 { - return Ok(()); - } + visitor + .visit_child(key, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("map key [{index}]")))?; + visitor + .visit_child(val, def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("map value [{index}]")))?; + } + Ok(()) +} - let iter_any = - Function::get_global(iter_name)?.call_packed(&[unsafe { view_of(&value) }])?; - let iter = Function::try_from(iter_any)?; - - if is_dict { - // Dict mutation invalidates its iterator, so snapshot all entries - // before dispatching to user code. - let mut entries = Vec::with_capacity(size); - for index in 0..size { - let key = iter.call_packed(&[AnyView::from(&0i64)])?; - let map_value = iter.call_packed(&[AnyView::from(&1i64)])?; - entries.push((key, map_value)); - if index + 1 != size { - iter.call_packed(&[AnyView::from(&2i64)])?; - } - } +/// Cold fallback used when the mirrored layout fails validation (e.g. an +/// ABI-debug build): iterate through the public packed functors. Map storage +/// is private C++; the Rust binding itself uses these iterator functors, so +/// no structural visiting or traversal control leaves Rust. Entries are +/// snapshotted before user callbacks run — required for Dict, whose mutation +/// invalidates the iterator, and harmless for immutable Map on this +/// non-performance path. +fn visit_map( + value: TVMFFIAny, + visitor: &mut C, + def_region_kind: DefRegionKind, +) -> NativeResult { + let is_dict = value.type_index == TVMFFITypeIndex::kTVMFFIDict as i32; + let (size_name, iter_name, kind) = if is_dict { + ("ffi.DictSize", "ffi.DictForwardIterFunctor", "dict") + } else { + ("ffi.MapSize", "ffi.MapForwardIterFunctor", "map") + }; + let size = Function::get_global(size_name)? + .call_packed(&[unsafe { view_of(&value) }]) + .and_then(i64::try_from)?; + if size < 0 { + return Err(runtime_error("native visitor: map reports a negative size").into()); + } + let size = usize::try_from(size) + .map_err(|_| runtime_error("native visitor: map size does not fit usize"))?; + if size == 0 { + return Ok(()); + } - for (index, (mut key, mut map_value)) in entries.into_iter().enumerate() { - let key_raw = raw_of_owned(&mut key); - self.visit_raw(key_raw, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("dict key [{index}]")))?; - let value_raw = raw_of_owned(&mut map_value); - self.visit_raw(value_raw, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("dict value [{index}]")))?; - } - return Ok(()); - } + let iter_any = Function::get_global(iter_name)?.call_packed(&[unsafe { view_of(&value) }])?; + let iter = Function::try_from(iter_any)?; - // Map is immutable. Retain only the current owned key/value pair. - for index in 0..size { - let mut key = iter.call_packed(&[AnyView::from(&0i64)])?; - let mut map_value = iter.call_packed(&[AnyView::from(&1i64)])?; - let key_raw = raw_of_owned(&mut key); - self.visit_raw(key_raw, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("map key [{index}]")))?; - let value_raw = raw_of_owned(&mut map_value); - self.visit_raw(value_raw, visitor, def_region_kind) - .map_err(|halt| with_error_context(halt, &format!("map value [{index}]")))?; - if index + 1 != size { - iter.call_packed(&[AnyView::from(&2i64)])?; - } + let mut entries = Vec::with_capacity(size); + for index in 0..size { + let key = iter.call_packed(&[AnyView::from(&0i64)])?; + let map_value = iter.call_packed(&[AnyView::from(&1i64)])?; + entries.push((key, map_value)); + if index + 1 != size { + iter.call_packed(&[AnyView::from(&2i64)])?; } - Ok(()) } - #[inline] - fn visit_reflected_fields( - &self, - value: TVMFFIAny, - visitor: &mut V, - def_region_kind: DefRegionKind, - ) -> NativeResult { - let type_info = unsafe { TVMFFIGetTypeInfo(value.type_index) }; - if type_info.is_null() { - return Err(runtime_error(&format!( - "native visitor: unregistered type index {}", - value.type_index - )) - .into()); - } - let seq_hash_kind = unsafe { - let metadata = (*type_info).metadata; - if metadata.is_null() { - TVMFFISEqHashKind::kTVMFFISEqHashKindUnsupported as i32 - } else { - (*metadata).structural_eq_hash_kind - } - }; - let def_region_kind = free_var_child_region(def_region_kind, seq_hash_kind); - let object = unsafe { value.data_union.v_obj } as *mut u8; - let halted = unsafe { - for_each_field(value.type_index, |field| { - match self.visit_reflected_field(object, field, visitor, def_region_kind) { - Ok(()) => ControlFlow::Continue(()), - Err(halt) => ControlFlow::Break(halt), - } - }) - }; - halted.map_or(Ok(()), Err) + for (index, (key, map_value)) in entries.into_iter().enumerate() { + visitor + .visit_child(raw_of_owned(&key), def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("{kind} key [{index}]")))?; + visitor + .visit_child(raw_of_owned(&map_value), def_region_kind) + .map_err(|halt| with_error_context(halt, &format!("{kind} value [{index}]")))?; } + Ok(()) +} - unsafe fn visit_reflected_field( - &self, - object: *mut u8, - field: &TVMFFIFieldInfo, - visitor: &mut V, - inherited_region: DefRegionKind, - ) -> NativeResult { - if field.flags & FLAG_SEQ_HASH_IGNORE != 0 { - return Ok(()); - } - - let Some(getter) = field.getter else { - return Err(NativeHalt::Error(runtime_error(&format!( - "native visitor: reflected field `{}` has no getter", - field.name.as_str() - )))); - }; - let address = object.offset(field.offset as isize) as *mut c_void; - let mut child_raw = TVMFFIAny::new(); - if getter(address, &mut child_raw) != 0 { - return Err(with_error_context( - NativeHalt::Error(Error::from_raised()), - &format!("field `{}`", field.name.as_str()), - )); +#[inline] +fn visit_reflected_fields( + value: TVMFFIAny, + visitor: &mut C, + def_region_kind: DefRegionKind, +) -> NativeResult { + let type_info = unsafe { TVMFFIGetTypeInfo(value.type_index) }; + if type_info.is_null() { + return Err(runtime_error(&format!( + "native visitor: unregistered type index {}", + value.type_index + )) + .into()); + } + let seq_hash_kind = unsafe { + let metadata = (*type_info).metadata; + if metadata.is_null() { + TVMFFISEqHashKind::kTVMFFISEqHashKindUnsupported as i32 + } else { + (*metadata).structural_eq_hash_kind } + }; + let def_region_kind = free_var_child_region(def_region_kind, seq_hash_kind); + let object = unsafe { value.data_union.v_obj } as *mut u8; + let halted = unsafe { + for_each_field(value.type_index, |field| { + match visit_reflected_field(object, field, visitor, def_region_kind) { + Ok(()) => ControlFlow::Continue(()), + Err(halt) => ControlFlow::Break(halt), + } + }) + }; + halted.map_or(Ok(()), Err) +} - // A reflection getter returns an owned Any. Keep it alive while the - // recursive walk borrows its raw cell. - let mut child = Any::from_raw_ffi_any(child_raw); - let borrowed = raw_of_owned(&mut child); - let child_region = field_def_region(field, inherited_region); - self.visit_raw(borrowed, visitor, child_region) - .map_err(|halt| with_error_context(halt, &format!("field `{}`", field.name.as_str()))) +unsafe fn visit_reflected_field( + object: *mut u8, + field: &TVMFFIFieldInfo, + visitor: &mut C, + inherited_region: DefRegionKind, +) -> NativeResult { + if field.flags & FLAG_SEQ_HASH_IGNORE != 0 { + return Ok(()); } - fn reject_foreign_structural_visit(&self, type_index: i32) -> Result<()> { - let Some(attr) = self - .structural_visit - .and_then(|column| column.get(type_index)) - else { - return Ok(()); - }; - match attr.type_index { - x if x == TVMFFITypeIndex::kTVMFFINone as i32 => Ok(()), - x if x == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 - || x == TVMFFITypeIndex::kTVMFFIFunction as i32 => - { - let value_type = if type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { - format!("type index {type_index}") - } else { - format!("type `{}`", type_key_of(type_index)) - }; - Err(runtime_error(&format!( - "native visitor: {value_type} registers foreign `{STRUCTURAL_VISIT_ATTR}`; \ - use a matching pre-order Rust handler, visit its children through \ - `VisitDispatch::subvisit`, and return `WalkResult::Skip`" - ))) - } - _ => Err(Error::new( - TYPE_ERROR, - &format!( - "{STRUCTURAL_VISIT_ATTR} must be an opaque function pointer or ffi.Function" - ), - "", - )), - } + let Some(getter) = field.getter else { + return Err(NativeHalt::Error(runtime_error(&format!( + "native visitor: reflected field `{}` has no getter", + field.name.as_str() + )))); + }; + let address = object.offset(field.offset as isize) as *mut c_void; + let mut child_raw = TVMFFIAny::new(); + if getter(address, &mut child_raw) != 0 { + return Err(with_error_context( + NativeHalt::Error(Error::from_raised()), + &format!("field `{}`", field.name.as_str()), + )); } - fn with_value_context(halt: NativeHalt, value: TVMFFIAny) -> NativeHalt { - if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { - halt - } else { - with_error_context(halt, &format!("object `{}`", type_key_of(value.type_index))) + // A reflection getter returns an owned Any. Keep it alive while the + // recursive walk borrows its raw cell. + let child = Any::from_raw_ffi_any(child_raw); + let borrowed = raw_of_owned(&child); + let child_region = field_def_region(field, inherited_region); + visitor + .visit_child(borrowed, child_region) + .map_err(|halt| with_error_context(halt, &format!("field `{}`", field.name.as_str()))) +} + +#[inline] +fn reject_foreign_structural_visit(type_index: i32) -> Result<()> { + let Some(attr) = structural_visit_column().and_then(|column| column.get(type_index)) else { + return Ok(()); + }; + match attr.type_index { + x if x == TVMFFITypeIndex::kTVMFFINone as i32 => Ok(()), + x if x == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 + || x == TVMFFITypeIndex::kTVMFFIFunction as i32 => + { + let value_type = if type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + format!("type index {type_index}") + } else { + format!("type `{}`", type_key_of(type_index)) + }; + Err(runtime_error(&format!( + "native visitor: {value_type} registers foreign `{STRUCTURAL_VISIT_ATTR}`; \ + visit its children explicitly from a `StructuralVisitor` \ + (`structural_visit`), or skip it with a pre-order `WalkResult::Skip` \ + handler" + ))) } + _ => Err(Error::new( + TYPE_ERROR, + &format!("{STRUCTURAL_VISIT_ATTR} must be an opaque function pointer or ffi.Function"), + "", + )), } } -/// Visit `root` in pre-order with typed handlers stored in `visitor`. -pub fn structural_visit(root: &R, visitor: &mut V) -> Result -where - V: VisitDispatch, - for<'x> AnyView<'x>: From<&'x R>, -{ - structural_walk(root, visitor, WalkOrder::PreOrder) +fn with_value_context(halt: NativeHalt, value: TVMFFIAny) -> NativeHalt { + if value.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + halt + } else { + with_error_context(halt, &format!("object `{}`", type_key_of(value.type_index))) + } } -/// Walk `root` with typed handlers and state stored in `walker`. +/// Visit `root` with a user-driven [`StructuralVisitor`]. /// -/// `walker` may use [`crate::dispatch`] exactly like a visitor. Each matching -/// handler runs once, before or after the value's children according to -/// `order`. -pub fn structural_walk(root: &R, walker: &mut W, order: WalkOrder) -> Result +/// The visitor's [`StructuralVisitor::visit`] runs for the root under +/// [`DefRegionKind::None`] and controls all further recursion itself. This is +/// the Rust analog of constructing a C++ `StructuralVisitorObj` and calling +/// `visitor->Visit(root)`. An FFI `None` root completes immediately. +pub fn structural_visit(root: &R, visitor: &mut V) -> Result> where - W: VisitDispatch, + V: StructuralVisitor, for<'x> AnyView<'x>: From<&'x R>, { - let native_walker = NativeWalker::new(); - let mut dispatch = DispatchVisitor { - visitor: walker, - order, - }; - finish(native_walker.visit_raw( - raw_of(AnyView::from(root)), - &mut dispatch, - DefRegionKind::None, - )) + visitor.visit_child(root, DefRegionKind::None) } -/// Native pre/post walk used by analyses that need to observe every raw value. -pub fn walk(root: &R, mut callback: F) -> Result -where - for<'x> AnyView<'x>: From<&'x R>, - F: FnMut(&VisitValue, Phase) -> O, - O: IntoVisitResult, -{ - walk_with_context(root, move |value, phase, _def_region_kind| { - callback(value, phase) - }) -} - -/// Native pre/post walk whose callback also receives definition-region state. -pub fn walk_with_context(root: &R, callback: F) -> Result +/// Walk `root` with an observer, the Rust analog of C++ +/// `StructuralWalk(root, callbacks...)`. +/// +/// `walker` is anything implementing [`IntoWalker`]: a `&mut` reference to a +/// stateful [`VisitDispatch`] visitor (`#[dispatch(visit)]`), or a bare +/// closure taking `&VisitValue` with an optional trailing [`DefRegionKind`] +/// — the C++ callback overloads. The walker owns recursion: the handler runs +/// once per value, before or after the value's children according to +/// `order`, and steers traversal through the returned [`WalkResult`]. +pub fn structural_walk( + root: &R, + walker: H, + order: WalkOrder, +) -> Result> where + H: IntoWalker, for<'x> AnyView<'x>: From<&'x R>, - F: FnMut(&VisitValue, Phase, DefRegionKind) -> O, - O: IntoVisitResult, { - let walker = NativeWalker::new(); - let mut callback = CallbackVisitor(callback); - finish(walker.visit_raw( + let mut dispatch = walker.into_walker(order); + finish(visit_raw( raw_of(AnyView::from(root)), - &mut callback, + &mut dispatch, DefRegionKind::None, )) } -fn finish(result: NativeResult) -> Result { +fn finish(result: NativeResult) -> Result> { match result { - Ok(()) => Ok(ControlFlow::Continue(())), + Ok(()) => Ok(None), Err(NativeHalt::Error(error)) => Err(error), - Err(NativeHalt::Interrupt(payload)) => Ok(ControlFlow::Break(payload)), + Err(NativeHalt::Interrupt(payload)) => Ok(Some(VisitInterrupt { value: payload })), } } @@ -976,7 +1102,9 @@ impl MapCursor { return None; } let block = data.add((*index / MAP_BLOCK_CAP) as usize * MAP_BLOCK_SIZE); - let item = block.add(MAP_BLOCK_CAP as usize + (*index % MAP_BLOCK_CAP) as usize * MAP_ITEM_SIZE); + let item = block.add( + MAP_BLOCK_CAP as usize + (*index % MAP_BLOCK_CAP) as usize * MAP_ITEM_SIZE, + ); let key = *(item as *const TVMFFIAny); let val = *(item.add(16) as *const TVMFFIAny); *index = *(item.add(MAP_ITEM_NEXT_OFFSET) as *const u64); @@ -1010,7 +1138,11 @@ fn map_layout_usable(value: TVMFFIAny) -> bool { fn validate_map_layout(value: TVMFFIAny) -> bool { let expected = (|| -> Result { let is_dict = value.type_index == TVMFFITypeIndex::kTVMFFIDict as i32; - let name = if is_dict { "ffi.DictSize" } else { "ffi.MapSize" }; + let name = if is_dict { + "ffi.DictSize" + } else { + "ffi.MapSize" + }; Function::get_global(name)? .call_packed(&[unsafe { view_of(&value) }]) .and_then(i64::try_from) @@ -1060,6 +1192,25 @@ fn type_attr_column(attr_name: &str) -> Option { } } +/// Cached `__s_visit__` column pointer (0 = not seen yet). A registry column +/// is stable once created — C++ `DefaultVisitExpected` caches the same +/// pointer in a function-local static — while an absent column is re-queried +/// because a later attr registration may create it. The cache keeps the +/// per-value foreign-hook check free of FFI lookups. +static STRUCTURAL_VISIT_COLUMN: AtomicUsize = AtomicUsize::new(0); + +#[inline] +fn structural_visit_column() -> Option { + let cached = STRUCTURAL_VISIT_COLUMN.load(Ordering::Relaxed); + if cached != 0 { + let pointer = cached as *mut TVMFFITypeAttrColumn; + return Some(TypeAttrColumn(unsafe { NonNull::new_unchecked(pointer) })); + } + let column = type_attr_column(STRUCTURAL_VISIT_ATTR)?; + STRUCTURAL_VISIT_COLUMN.store(column.0.as_ptr() as usize, Ordering::Relaxed); + Some(column) +} + fn type_key_of(type_index: i32) -> String { unsafe { let info = TVMFFIGetTypeInfo(type_index); @@ -1145,7 +1296,7 @@ fn raw_of(view: AnyView<'_>) -> TVMFFIAny { } #[inline] -fn raw_of_owned(any: &mut Any) -> TVMFFIAny { +fn raw_of_owned(any: &Any) -> TVMFFIAny { *any.as_raw_ffi_any() } @@ -1174,15 +1325,13 @@ mod tests { #[derive(Default)] struct TypedRegionProbe { - def_region: DefRegionKind, seen: Vec, } - #[crate::dispatch(visit, def_region = def_region)] + #[crate::dispatch(visit)] impl TypedRegionProbe { - fn visit_integer(&mut self, _value: i64) -> WalkResult { - let kind = self.def_region_kind(); - self.seen.push(kind); + fn visit_integer(&mut self, _value: i64, def_region_kind: DefRegionKind) -> WalkResult { + self.seen.push(def_region_kind); WalkResult::Advance } } @@ -1196,21 +1345,18 @@ mod tests { #[test] fn def_region_is_inherited_through_containers() { let root = Array::new(vec![1i64, 2]); - let walker = NativeWalker::new(); let mut probe = RegionProbe(Vec::new()); - assert!(walker - .visit_raw( - raw_of(AnyView::from(&root)), - &mut probe, - DefRegionKind::Recursive, - ) - .is_ok()); + assert!(visit_raw( + raw_of(AnyView::from(&root)), + &mut probe, + DefRegionKind::Recursive, + ) + .is_ok()); assert_eq!(probe.0, vec![DefRegionKind::Recursive; 3]); } #[test] - fn reflected_field_def_region_reaches_typed_handler_and_restores() { - let walker = NativeWalker::new(); + fn reflected_field_def_region_reaches_typed_handler() { let mut probe = TypedRegionProbe::default(); let mut dispatch = DispatchVisitor { visitor: &mut probe, @@ -1222,6 +1368,9 @@ mod tests { field.getter = Some(clone_any_field); let object = (&mut value as *mut Any).cast::(); + let mut children = WalkChildren { + visitor: &mut dispatch, + }; for flags in [ FLAG_SEQ_HASH_DEF_RECURSIVE, 0, @@ -1231,7 +1380,7 @@ mod tests { ] { field.flags = flags; assert!(unsafe { - walker.visit_reflected_field(object, &field, &mut dispatch, DefRegionKind::None) + visit_reflected_field(object, &field, &mut children, DefRegionKind::None) } .is_ok()); } diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index 987bb9480..70eff8df7 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -47,8 +47,8 @@ pub use crate::error::{ }; pub use crate::extra::module::Module; pub use crate::extra::structural_visit::{ - structural_visit, structural_walk, walk, walk_with_context, DefRegionKind, Phase, - VisitDispatch, VisitOutcome, VisitValue, WalkOrder, WalkResult, + structural_visit, structural_walk, DefRegionKind, IntoWalker, StructuralVisitor, VisitDispatch, + VisitInterrupt, VisitValue, WalkOrder, WalkResult, }; pub use crate::function::Function; pub use crate::object::ObjectRefCast; diff --git a/rust/tvm-ffi/tests/test_dispatch.rs b/rust/tvm-ffi/tests/test_dispatch.rs index b02da3650..49ff99e80 100644 --- a/rust/tvm-ffi/tests/test_dispatch.rs +++ b/rust/tvm-ffi/tests/test_dispatch.rs @@ -21,16 +21,15 @@ //! path emitted by the dispatch macro outside `tvm_ffi` itself. use tvm_ffi::{ - dispatch, structural_visit, Array, DefRegionKind, Object, VisitDispatch, WalkResult, + dispatch, structural_walk, Array, DefRegionKind, Object, VisitDispatch, WalkOrder, WalkResult, }; #[derive(Default)] struct ExternalCounter { - def_region: DefRegionKind, objects: usize, } -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] impl ExternalCounter { #[cfg(any(unix, windows))] #[cfg_attr(all(), inline)] @@ -45,11 +44,9 @@ fn assert_visit_dispatch() {} const _: fn() = assert_visit_dispatch::; #[derive(Default)] -struct CfgAttrCounter { - def_region: DefRegionKind, -} +struct CfgAttrCounter {} -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] impl CfgAttrCounter { #[cfg(any())] fn visit_disabled_catch_all(&mut self, _value: &tvm_ffi::VisitValue) -> WalkResult { @@ -71,7 +68,7 @@ const _: fn() = assert_visit_dispatch::; struct DisabledCounter; const _: usize = std::mem::size_of::(); -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] #[cfg(any())] impl DisabledCounter { fn visit_object(&mut self, _value: &Object) -> WalkResult { @@ -82,7 +79,7 @@ impl DisabledCounter { struct CfgAttrDisabledCounter; const _: usize = std::mem::size_of::(); -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] #[cfg_attr(all(), cfg(any()))] impl CfgAttrDisabledCounter { fn visit_object(&mut self, _value: &Object) -> WalkResult { @@ -90,10 +87,42 @@ impl CfgAttrDisabledCounter { } } +#[derive(Default)] +struct MixedArityCounter { + kinds: Vec, + objects: usize, +} + +#[dispatch(visit)] +impl MixedArityCounter { + fn visit_int(&mut self, _value: i64, kind: DefRegionKind) -> WalkResult { + self.kinds.push(kind); + WalkResult::Advance + } + + fn visit_object(&mut self, _value: &Object) -> WalkResult { + self.objects += 1; + WalkResult::Advance + } +} + +#[test] +fn handlers_may_mix_def_region_arity() { + let root = Array::new(vec![1i64, 2]); + let mut visitor = MixedArityCounter::default(); + assert!(structural_walk(&root, &mut visitor, WalkOrder::PreOrder) + .unwrap() + .is_none()); + assert_eq!(visitor.objects, 1); + assert_eq!(visitor.kinds, vec![DefRegionKind::None; 2]); +} + #[test] fn generated_dispatch_uses_public_downstream_paths() { let root = Array::new(vec![1i64, 2]); let mut visitor = ExternalCounter::default(); - assert!(structural_visit(&root, &mut visitor).unwrap().is_continue()); + assert!(structural_walk(&root, &mut visitor, WalkOrder::PreOrder) + .unwrap() + .is_none()); assert_eq!(visitor.objects, 1); } diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index 16ae9c930..d13ffa547 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -17,12 +17,10 @@ * under the License. */ -use std::ops::ControlFlow; - use tvm_ffi::tvm_ffi_sys::{TVMFFIByteArray, TVMFFITypeIndex, TVMFFITypeRegisterAttr}; use tvm_ffi::{ - dispatch, structural_visit, structural_walk, walk, walk_with_context, Any, AnyView, Array, - DefRegionKind, Error, Function, Map, Phase, Result, Shape, String as FfiString, VisitDispatch, + dispatch, structural_visit, structural_walk, Any, AnyView, Array, DefRegionKind, Error, + Function, Map, Result, Shape, String as FfiString, StructuralVisitor, VisitInterrupt, VisitValue, WalkOrder, WalkResult, RUNTIME_ERROR, }; @@ -34,14 +32,18 @@ fn runtime_error(message: &str) -> Error { fn plain_walk_uses_native_sequence_fallback() { let root = Array::new(vec![1i64, 2, 3]); let mut integers = 0; - assert!(walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::().is_some() { - integers += 1; - } - WalkResult::Advance - }) + assert!(structural_walk( + &root, + |value: &VisitValue| { + if value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap() - .is_continue()); + .is_none()); assert_eq!(integers, 3); } @@ -51,29 +53,52 @@ fn plain_walk_uses_native_map_fallback() { .into_iter() .collect(); let mut integers = 0; - assert!(walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::().is_some() { - integers += 1; - } - WalkResult::Advance - }) + assert!(structural_walk( + &root, + |value: &VisitValue| { + if value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap() - .is_continue()); + .is_none()); assert_eq!(integers, 2); } #[derive(Default)] -struct SkipForeignShape { - def_region: DefRegionKind, -} +struct SkipForeignShape {} -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] impl SkipForeignShape { fn visit_shape(&mut self, _shape: Shape) -> WalkResult { WalkResult::Skip } } +/// Visitor-layer handling of the foreign type: `visit` enumerates the +/// children itself (none, for a shape) instead of the default recursion. +#[derive(Default)] +struct ForeignShapeVisitor { + shapes: usize, +} + +impl StructuralVisitor for ForeignShapeVisitor { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + if value.cast::().is_some() { + self.shapes += 1; + return Ok(None); + } + self.default_visit_children(value, def_region_kind) + } +} + #[test] fn foreign_structural_visit_requires_explicit_rust_override() { let hook = Function::get_global("ffi.ArraySize").unwrap(); @@ -91,16 +116,28 @@ fn foreign_structural_visit_requires_explicit_rust_override() { ); let root = Shape::from([2i64, 3]); - let error = match walk(&root, |_value, _phase| WalkResult::Advance) { + let error = match structural_walk( + &root, + |_value: &VisitValue| WalkResult::Advance, + WalkOrder::PreOrder, + ) { Err(error) => error, Ok(_) => panic!("foreign structural visit unexpectedly used reflection"), }; assert!(error.message().contains("registers foreign `__s_visit__`")); - assert!(error.message().contains("return `WalkResult::Skip`")); + assert!(error.message().contains("StructuralVisitor")); - assert!(structural_visit(&root, &mut SkipForeignShape::default()) - .unwrap() - .is_continue()); + // Walk layer: a pre-order handler skips the foreign type. + assert!( + structural_walk(&root, &mut SkipForeignShape::default(), WalkOrder::PreOrder) + .unwrap() + .is_none() + ); + + // Visitor layer: take over the type's children explicitly instead. + let mut takeover = ForeignShapeVisitor::default(); + assert!(structural_visit(&root, &mut takeover).unwrap().is_none()); + assert_eq!(takeover.shapes, 1); } #[test] @@ -114,8 +151,9 @@ fn mutable_list_is_snapshotted_before_callbacks() { let mut appended = false; let mut integers = Vec::new(); - assert!(walk(&root, |value, phase| { - if phase == Phase::Enter { + assert!(structural_walk( + &root, + |value: &VisitValue| { if let Some(integer) = value.cast::() { integers.push(integer); if !appended { @@ -125,11 +163,12 @@ fn mutable_list_is_snapshotted_before_callbacks() { appended = true; } } - } - WalkResult::Advance - }) + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap() - .is_continue()); + .is_none()); assert_eq!(integers, vec![1, 2]); let size = Function::get_global("ffi.ListSize") @@ -156,8 +195,9 @@ fn mutable_dict_is_snapshotted_before_callbacks() { let mut inserted = false; let mut integers = Vec::new(); - assert!(walk(&root, |value, phase| { - if phase == Phase::Enter { + assert!(structural_walk( + &root, + |value: &VisitValue| { if let Some(integer) = value.cast::() { integers.push(integer); if !inserted { @@ -171,11 +211,12 @@ fn mutable_dict_is_snapshotted_before_callbacks() { inserted = true; } } - } - WalkResult::Advance - }) + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap() - .is_continue()); + .is_none()); integers.sort_unstable(); assert_eq!(integers, vec![1, 2]); @@ -195,18 +236,20 @@ fn dense_map_layout_is_traversed_completely() { .collect(); let mut sum = 0; let mut strings = 0; - assert!(walk(&root, |value, phase| { - if phase == Phase::Enter { + assert!(structural_walk( + &root, + |value: &VisitValue| { if let Some(integer) = value.cast::() { sum += integer; } else if value.cast::().is_some() { strings += 1; } - } - WalkResult::Advance - }) + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap() - .is_continue()); + .is_none()); assert_eq!(sum, (0..9).sum::()); assert_eq!(strings, 9); } @@ -216,29 +259,37 @@ fn interrupt_payload_crosses_map_traversal() { let root: Map = [(FfiString::from("a"), 1i64), (FfiString::from("b"), 2i64)] .into_iter() .collect(); - let outcome = walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::().is_some() { - return WalkResult::interrupt_with(99i64); - } - WalkResult::Advance - }) + let outcome = structural_walk( + &root, + |value: &VisitValue| { + if value.cast::().is_some() { + return WalkResult::interrupt_with(99i64); + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap(); - let ControlFlow::Break(payload) = outcome else { + let Some(interrupt) = outcome else { panic!("map walk unexpectedly completed"); }; - assert_eq!(i64::try_from(payload).unwrap(), 99); + assert_eq!(i64::try_from(interrupt.value).unwrap(), 99); } #[test] fn handler_error_crosses_map_traversal() { let root: Map = [(FfiString::from("a"), 1i64)].into_iter().collect(); - let error = match walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::().is_some() { - Err(runtime_error("map handler failed")) - } else { - Ok(WalkResult::Advance) - } - }) { + let error = match structural_walk( + &root, + |value: &VisitValue| -> Result { + if value.cast::().is_some() { + Err(runtime_error("map handler failed")) + } else { + Ok(WalkResult::Advance) + } + }, + WalkOrder::PreOrder, + ) { Err(error) => error, Ok(_) => panic!("map handler unexpectedly succeeded"), }; @@ -250,85 +301,58 @@ fn handler_error_crosses_map_traversal() { fn interrupt_stops_without_running_remaining_callbacks() { let root = Array::new(vec![1i64, 2, 3]); let mut integers = 0; - let outcome = walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::().is_some() { - integers += 1; - return WalkResult::Interrupt; - } - WalkResult::Advance - }) + let outcome = structural_walk( + &root, + |value: &VisitValue| { + if value.cast::().is_some() { + integers += 1; + return WalkResult::Interrupt; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap(); - assert!(outcome.is_break()); + assert!(outcome.is_some()); assert_eq!(integers, 1); } +/// Visitor-layer traversal that overrides the def-region for one child and +/// inherits it for the next, mirroring a C++ visitor using +/// `WithDefRegionKind`. #[derive(Default)] -struct ManualRegionProbe { - def_region: DefRegionKind, +struct ManualRegionVisitor { seen: Vec, } -#[dispatch(visit, def_region = def_region)] -impl ManualRegionProbe { - fn visit_array(&mut self, array: Array) -> Result { - let before = self.def_region_kind(); - let overridden = array.get(0).unwrap(); - if let ControlFlow::Break(payload) = - self.subvisit_with_def_region(&overridden, DefRegionKind::NonRecursive)? - { - return Ok(WalkResult::InterruptWith(payload)); +impl StructuralVisitor for ManualRegionVisitor { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + if let Some(array) = value.cast::>() { + // Override the state for exactly this child's subtree... + let overridden = array.get(0).unwrap(); + if let Some(interrupt) = self.visit_child(&overridden, DefRegionKind::NonRecursive)? { + return Ok(Some(interrupt)); + } + // ...and forward the received state to inherit it. + let inherited = array.get(1).unwrap(); + return self.visit_child(&inherited, def_region_kind); } - // The walker rewinds the mirror after nested dispatches. - assert_eq!(self.def_region_kind(), before); - let inherited = array.get(1).unwrap(); - if let ControlFlow::Break(payload) = self.subvisit(&inherited)? { - return Ok(WalkResult::InterruptWith(payload)); + if value.cast::().is_some() { + self.seen.push(def_region_kind); } - Ok(WalkResult::Skip) - } - - fn visit_integer(&mut self, _value: i64) -> WalkResult { - let kind = self.def_region_kind(); - self.seen.push(kind); - WalkResult::Advance + Ok(None) } } -#[derive(Default)] -struct MirrorCorruptionProbe { - def_region: DefRegionKind, - seen: Vec, -} - -#[dispatch(visit, def_region = def_region)] -impl MirrorCorruptionProbe { - fn visit_array(&mut self, _array: Array) -> WalkResult { - // Overwriting the mirror must not leak into the walker's own - // by-value propagation. - self.def_region = DefRegionKind::Recursive; - WalkResult::Advance - } - - fn visit_integer(&mut self, _value: i64) -> WalkResult { - let kind = self.def_region_kind(); - self.seen.push(kind); - WalkResult::Advance - } -} - -#[test] -fn mirror_corruption_does_not_affect_walker_propagation() { - let root = Array::new(vec![7i64]); - let mut probe = MirrorCorruptionProbe::default(); - assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); - assert_eq!(probe.seen, vec![DefRegionKind::None]); -} - #[test] fn manual_child_visit_can_override_def_region() { let root = Array::new(vec![7i64, 8]); - let mut probe = ManualRegionProbe::default(); - assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + let mut probe = ManualRegionVisitor::default(); + assert!(structural_visit(&root, &mut probe).unwrap().is_none()); assert_eq!( probe.seen, vec![DefRegionKind::NonRecursive, DefRegionKind::None] @@ -337,13 +361,12 @@ fn manual_child_visit_can_override_def_region() { #[derive(Default)] struct GenericDispatchProbe { - def_region: DefRegionKind, integers: Vec, objects: usize, catch_all: usize, } -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] impl GenericDispatchProbe { fn visit_integer(&mut self, value: i64) -> WalkResult { self.integers.push(value); @@ -365,43 +388,52 @@ impl GenericDispatchProbe { fn generated_dispatch_supports_pod_and_ordered_catch_all() { let root = Array::new(vec![1i64, 2]); let mut probe = GenericDispatchProbe::default(); - assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + assert!(structural_walk(&root, &mut probe, WalkOrder::PreOrder) + .unwrap() + .is_none()); assert_eq!(probe.integers, vec![1, 2]); assert_eq!(probe.objects, 1); let floats = Array::new(vec![1.0f64, 2.0]); - assert!(structural_visit(&floats, &mut probe).unwrap().is_continue()); + assert!(structural_walk(&floats, &mut probe, WalkOrder::PreOrder) + .unwrap() + .is_none()); assert_eq!(probe.objects, 2); assert_eq!(probe.catch_all, 2); } +/// Visitor-layer enter/exit straddling: run enter logic, delegate the +/// default child recursion, then run exit logic with the same locals in +/// scope — the C++ `DefaultVisitExpected` pattern. #[derive(Default)] -struct StraddleProbe { - def_region: DefRegionKind, +struct StraddleVisitor { events: Vec, } -#[dispatch(visit, def_region = def_region)] -impl StraddleProbe { - fn visit_any(&mut self, value: &VisitValue) -> Result { +impl StructuralVisitor for StraddleVisitor { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { let label = match value.cast::() { Some(integer) => format!("int:{integer}"), None => "node".to_string(), }; self.events.push(format!("enter:{label}")); - if let ControlFlow::Break(payload) = self.subvisit_children(value)? { - return Ok(WalkResult::InterruptWith(payload)); + if let Some(interrupt) = self.default_visit_children(value, def_region_kind)? { + return Ok(Some(interrupt)); } self.events.push(format!("exit:{label}")); - Ok(WalkResult::Skip) + Ok(None) } } #[test] -fn catch_all_handler_can_straddle_default_children() { +fn visitor_can_straddle_default_children() { let root = Array::new(vec![1i64, 2]); - let mut probe = StraddleProbe::default(); - assert!(structural_visit(&root, &mut probe).unwrap().is_continue()); + let mut probe = StraddleVisitor::default(); + assert!(structural_visit(&root, &mut probe).unwrap().is_none()); assert_eq!( probe.events, vec![ @@ -417,11 +449,10 @@ fn catch_all_handler_can_straddle_default_children() { #[derive(Default)] struct OrderProbe { - def_region: DefRegionKind, events: Vec, } -#[dispatch(visit, def_region = def_region)] +#[dispatch(visit)] impl OrderProbe { fn visit_array(&mut self, _array: Array) -> WalkResult { self.events.push("array".to_string()); @@ -440,36 +471,44 @@ fn stateful_structural_walk_supports_post_order() { let mut probe = OrderProbe::default(); assert!(structural_walk(&root, &mut probe, WalkOrder::PostOrder) .unwrap() - .is_continue()); + .is_none()); assert_eq!(probe.events, vec!["int:1", "int:2", "array"]); } #[test] fn interrupt_payload_is_returned_to_the_caller() { let root = Array::new(vec![1i64, 2]); - let outcome = walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::() == Some(1) { - return WalkResult::interrupt_with(42i64); - } - WalkResult::Advance - }) + let outcome = structural_walk( + &root, + |value: &VisitValue| { + if value.cast::() == Some(1) { + return WalkResult::interrupt_with(42i64); + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) .unwrap(); - let ControlFlow::Break(payload) = outcome else { + let Some(interrupt) = outcome else { panic!("walk unexpectedly completed"); }; - assert_eq!(i64::try_from(payload).unwrap(), 42); + assert_eq!(i64::try_from(interrupt.value).unwrap(), 42); } #[test] fn handler_errors_include_native_visit_path() { let root = Array::new(vec![1i64]); - let error = match walk(&root, |value, phase| { - if phase == Phase::Enter && value.cast::().is_some() { - Err(runtime_error("handler failed")) - } else { - Ok(WalkResult::Advance) - } - }) { + let error = match structural_walk( + &root, + |value: &VisitValue| -> Result { + if value.cast::().is_some() { + Err(runtime_error("handler failed")) + } else { + Ok(WalkResult::Advance) + } + }, + WalkOrder::PreOrder, + ) { Err(error) => error, Ok(_) => panic!("handler unexpectedly succeeded"), }; @@ -479,16 +518,165 @@ fn handler_errors_include_native_visit_path() { } #[test] -fn raw_walk_context_receives_def_region() { +fn visitor_errors_include_native_visit_path() { + struct FailingVisitor; + + impl StructuralVisitor for FailingVisitor { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + if value.cast::().is_some() { + return Err(runtime_error("visitor failed")); + } + self.default_visit_children(value, def_region_kind) + } + } + let root = Array::new(vec![1i64]); - let mut regions = Vec::new(); - assert!(walk_with_context(&root, |_value, phase, region| { - if phase == Phase::Enter { - regions.push(region); + let error = match structural_visit(&root, &mut FailingVisitor) { + Err(error) => error, + Ok(_) => panic!("visitor unexpectedly succeeded"), + }; + assert_eq!(error.message(), "visitor failed"); + assert!(error.backtrace().contains("sequence item [0]")); + assert!(error.backtrace().contains("object `ffi.Array`")); +} + +#[test] +fn visitor_interrupt_propagates_through_default_children() { + struct InterruptingVisitor; + + impl StructuralVisitor for InterruptingVisitor { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + if value.cast::() == Some(2) { + return Ok(Some(VisitInterrupt::with(7i64))); + } + self.default_visit_children(value, def_region_kind) } - WalkResult::Advance - }) + } + + let root = Array::new(vec![1i64, 2, 3]); + let outcome = structural_visit(&root, &mut InterruptingVisitor).unwrap(); + let Some(interrupt) = outcome else { + panic!("visitor traversal unexpectedly completed"); + }; + assert_eq!(i64::try_from(interrupt.value).unwrap(), 7); +} + +#[test] +fn closure_walk_observes_values() { + // C++: StructuralWalk(root, [&](AnyView value) { ... }) + let root = Array::new(vec![1i64, 2, 3]); + let mut integers = 0; + assert!(structural_walk( + &root, + |value: &VisitValue| { + if value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(integers, 3); +} + +#[test] +fn closure_walk_receives_def_region_kind() { + // C++: StructuralWalk(root, + // [&](const TVarObj* var, TVMFFIDefRegionKind kind) { ... }) + let root = Array::new(vec![1i64, 2]); + let mut kinds = Vec::new(); + assert!(structural_walk( + &root, + |value: &VisitValue, kind: DefRegionKind| { + if value.cast::().is_some() { + kinds.push(kind); + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(kinds, vec![DefRegionKind::None; 2]); +} + +#[test] +fn closure_walk_interrupts_and_propagates_errors() { + let root = Array::new(vec![1i64, 2, 3]); + let outcome = structural_walk( + &root, + |value: &VisitValue| -> Result { + if value.cast::() == Some(2) { + return Ok(WalkResult::interrupt_with(2i64)); + } + Ok(WalkResult::Advance) + }, + WalkOrder::PreOrder, + ) + .unwrap(); + let Some(interrupt) = outcome else { + panic!("closure walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(interrupt.value).unwrap(), 2); + + let error = match structural_walk( + &root, + |value: &VisitValue| -> Result { + if value.cast::().is_some() { + Err(runtime_error("closure failed")) + } else { + Ok(WalkResult::Advance) + } + }, + WalkOrder::PreOrder, + ) { + Err(error) => error, + Ok(_) => panic!("closure walk unexpectedly succeeded"), + }; + assert_eq!(error.message(), "closure failed"); + assert!(error.backtrace().contains("object `ffi.Array`")); +} + +#[test] +fn closure_walk_supports_post_order_and_skip() { + let root = Array::new(vec![1i64, 2]); + let mut order_probe = Vec::new(); + assert!(structural_walk( + &root, + |value: &VisitValue| { + order_probe.push(value.cast::()); + WalkResult::Advance + }, + WalkOrder::PostOrder, + ) + .unwrap() + .is_none()); + assert_eq!(order_probe, vec![Some(1), Some(2), None]); + + let mut visited = 0; + assert!(structural_walk( + &root, + |value: &VisitValue| { + visited += 1; + if value.cast::().is_none() { + WalkResult::Skip + } else { + WalkResult::Advance + } + }, + WalkOrder::PreOrder, + ) .unwrap() - .is_continue()); - assert_eq!(regions, vec![DefRegionKind::None; 2]); + .is_none()); + assert_eq!(visited, 1); } diff --git a/rust/tvm-ffi/tests/test_visitor_alignment.rs b/rust/tvm-ffi/tests/test_visitor_alignment.rs index 62c74d54b..2d964bdff 100644 --- a/rust/tvm-ffi/tests/test_visitor_alignment.rs +++ b/rust/tvm-ffi/tests/test_visitor_alignment.rs @@ -20,23 +20,20 @@ //! Rust mirror of the C++ visitor example. //! //! `RecordingVisitor` lines up member-for-member with the C++ -//! `TestVisitorObj` (tests/cpp/extra/test_structural_visit.cc), and its -//! `visit_array` handler plays the role of the C++ `TFuncObj::StructuralVisit` -//! hook (tests/cpp/testing_object.h): the first element is visited as a -//! recursive definition region, the rest inherit the surrounding state. - -use std::ops::ControlFlow; +//! `TestVisitorObj` (tests/cpp/extra/test_structural_visit.cc): its `visit` +//! plays the role of the C++ vtable `VisitImpl`, and the array arm plays the +//! role of the C++ `TFuncObj::StructuralVisit` hook +//! (tests/cpp/testing_object.h): the first element is visited as a recursive +//! definition region, the rest inherit the surrounding state. use tvm_ffi::{ - dispatch, structural_visit, Array, DefRegionKind, Result, String as FfiString, VisitDispatch, - VisitValue, WalkResult, + structural_visit, Array, DefRegionKind, Result, String as FfiString, StructuralVisitor, + VisitInterrupt, VisitValue, }; /// C++: class TestVisitorObj : public StructuralVisitorObj #[derive(Default)] struct RecordingVisitor { - /// C++: `def_region_mode_` (base-class member maintained by the walker). - def_region: DefRegionKind, /// C++: `std::vector visited;` visited: Vec, /// C++: `std::vector modes;` @@ -45,56 +42,50 @@ struct RecordingVisitor { interrupt_on: Option, } -#[dispatch(visit, def_region = def_region)] -impl RecordingVisitor { - /// C++ analog: `TFuncObj::StructuralVisit` — visit "params" (element 0) - /// under a recursive definition region, then the "body" (element 1) - /// under the inherited state, and skip the default recursion. - fn visit_array(&mut self, array: Array) -> Result { - self.visited.push("array".to_string()); - self.modes.push(self.def_region_kind()); - - // C++: visitor->WithDefRegionKind(kTVMFFIDefRegionKindRecursive, - // [&] { return visitor->VisitExpected(self->params); }) - let params = array.get(0).unwrap(); - if let ControlFlow::Break(payload) = - self.subvisit_with_def_region(¶ms, DefRegionKind::Recursive)? - { - return Ok(WalkResult::InterruptWith(payload)); - } - - // C++: visitor->VisitExpected(self->body) (inherits the current state) - let body = array.get(1).unwrap(); - if let ControlFlow::Break(payload) = self.subvisit(&body)? { - return Ok(WalkResult::InterruptWith(payload)); - } - Ok(WalkResult::Skip) - } - +impl StructuralVisitor for RecordingVisitor { /// C++ analog: `TestVisitorObj::VisitImpl` — record every value together /// with the active def-region state, optionally interrupt with a payload, - /// otherwise delegate the default recursion. - fn visit_any(&mut self, value: &VisitValue) -> Result { - let label = match value.cast::() { + /// otherwise delegate recursion explicitly. The array arm mirrors + /// `TFuncObj::StructuralVisit`. + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + let integer = value.cast::(); + let label = match integer { Some(integer) => integer.to_string(), + None if value.cast::>().is_some() => "array".to_string(), None => "obj".to_string(), }; // C++: visited.push_back(value_ref); // modes.push_back(def_region_mode_); self.visited.push(label); - self.modes.push(self.def_region_kind()); + self.modes.push(def_region_kind); // C++: if (value_ref.same_as(interrupt_on)) // return VisitInterrupt(String("stop")); - if self.interrupt_on.is_some() && value.cast::() == self.interrupt_on { - return Ok(WalkResult::interrupt_with(FfiString::from("stop"))); + if self.interrupt_on.is_some() && integer == self.interrupt_on { + return Ok(Some(VisitInterrupt::with(FfiString::from("stop")))); } - // C++: return DefaultVisitExpected(value); - if let ControlFlow::Break(payload) = self.subvisit_children(value)? { - return Ok(WalkResult::InterruptWith(payload)); + // C++ analog: `TFuncObj::StructuralVisit` — visit "params" + // (element 0) under a recursive definition region, then the "body" + // (element 1) under the inherited state. + if let Some(array) = value.cast::>() { + // C++: visitor->WithDefRegionKind(kTVMFFIDefRegionKindRecursive, + // [&] { return visitor->VisitExpected(self->params); }) + let params = array.get(0).unwrap(); + if let Some(interrupt) = self.visit_child(¶ms, DefRegionKind::Recursive)? { + return Ok(Some(interrupt)); + } + // C++: visitor->VisitExpected(self->body) (inherits the state) + let body = array.get(1).unwrap(); + return self.visit_child(&body, def_region_kind); } - Ok(WalkResult::Skip) + + // C++: return DefaultVisitExpected(value); + self.default_visit_children(value, def_region_kind) } } @@ -107,7 +98,7 @@ fn records_values_and_def_region_modes() { let outcome = structural_visit(&root, &mut visitor).unwrap(); - assert!(outcome.is_continue()); + assert!(outcome.is_none()); assert_eq!(visitor.visited, vec!["array", "10", "20"]); assert_eq!( visitor.modes, @@ -131,9 +122,12 @@ fn stops_on_interrupt_with_payload() { let outcome = structural_visit(&root, &mut visitor).unwrap(); - let ControlFlow::Break(payload) = outcome else { + let Some(interrupt) = outcome else { panic!("traversal unexpectedly completed"); }; - assert_eq!(FfiString::try_from(payload).unwrap().as_str(), "stop"); + assert_eq!( + FfiString::try_from(interrupt.value).unwrap().as_str(), + "stop" + ); assert_eq!(visitor.visited, vec!["array", "10", "20"]); } From 94b9971cc6d72470dde52abe4c840ffe128112c7 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 12:53:16 -0400 Subject: [PATCH 11/20] update readme. Signed-off-by: yuchuan --- rust/README.md | 91 ++++++++++++++------------------------------------ 1 file changed, 25 insertions(+), 66 deletions(-) diff --git a/rust/README.md b/rust/README.md index 226e082ae..47fbf2773 100644 --- a/rust/README.md +++ b/rust/README.md @@ -31,22 +31,29 @@ efficiency while maintaining interoperability. ## Structural Visitors and Walkers The `tvm-ffi` crate provides native Rust structural traversal over FFI values, -built-in containers, and reflected object fields, split into the same two -layers as the C++ API: +built-in containers, and reflected object fields, mirroring the two layers of +the C++ API: - **Walk layer (observer)** — `structural_walk`, the analog of C++ - `StructuralWalk`: the walker owns recursion; handlers observe each value - and steer traversal by returning a `WalkResult`. + `StructuralWalk`: the walker owns recursion; handlers observe each value in + pre- or post-order (`WalkOrder`) and steer traversal through the returned + `WalkResult` (`Advance`, `Skip`, `Interrupt`/`InterruptWith`). - **Visitor layer (user-driven)** — `structural_visit` with a - `StructuralVisitor`, the analog of a hand-written C++ - `StructuralVisitorObj`: your `visit` method runs for the root and then - controls all recursion itself. + `StructuralVisitor`, the analog of a C++ `StructuralVisitorObj`: your + `visit` method runs for the root and controls all recursion itself. + +Every traversal returns `Result>`: `Ok(None)` means +the whole graph was visited, `Ok(Some(interrupt))` carries an interrupting +handler's payload. Handlers may also return `Result` to propagate +errors with `?`. ### Observer walks -`#[dispatch(visit)]` turns the `visit_*` methods in an inherent -implementation into a typed, stateful observer. Each value is automatically -dispatched to the handler matching its runtime type: +`#[dispatch(visit)]` turns the `visit_*` methods of an inherent impl into a +typed, stateful observer. Each value dispatches to the first handler matching +its runtime type (a `&VisitValue` handler acts as the catch-all); a handler +that needs the definition-region state declares a trailing `DefRegionKind` +argument: ```rust use tvm_ffi::{dispatch, structural_walk, Array, WalkOrder, WalkResult}; @@ -77,56 +84,16 @@ assert!(structural_walk(&values, &mut calculator, WalkOrder::PreOrder) assert_eq!(calculator.value, 12.0); ``` -Typed handlers are tested in source order: borrowed `ObjectCore` node types use -runtime subtype checks, owned arguments use `AnyCompatible` casts, and a final -`&VisitValue` handler acts as a catch-all. A handler that needs the -definition-region state declares a trailing `DefRegionKind` argument and the -generated dispatch forwards it by arity — the analog of a C++ `StructuralWalk` -callback accepting `(value, def_region_kind)` instead of `(value)`. -`WalkOrder` selects pre-order or post-order dispatch. - -`structural_walk` also accepts a bare closure — a catch-all observer taking -`&VisitValue` with an optional trailing `DefRegionKind`, mirroring the C++ -callback overloads (annotate the closure arguments so the handler shape can -be inferred): - -```rust -use tvm_ffi::{structural_walk, Array, VisitValue, WalkOrder, WalkResult}; - -let values = Array::new(vec![1_i64, 2, 3]); -let mut integers = 0; -assert!(structural_walk( - &values, - |value: &VisitValue| { - if value.cast::().is_some() { - integers += 1; - } - WalkResult::Advance - }, - WalkOrder::PreOrder, -) -.unwrap() -.is_none()); -assert_eq!(integers, 3); -``` - -`WalkResult::Advance` visits container or reflected children, `Skip` suppresses -the current value's default recursion, and `Interrupt`/`InterruptWith` halt the -walk. Every traversal returns `Result>`, matching the -C++ `Expected>`: `Ok(None)` means the whole graph was -visited, `Ok(Some(interrupt))` carries the interrupting handler's payload. -Handlers and callbacks may also return `Result` to propagate -errors with `?`. +`structural_walk` also accepts a bare closure taking `&VisitValue` (with an +optional trailing `DefRegionKind`) as a catch-all observer. ### User-driven visitors When traversal itself is part of the analysis — visiting selected children, custom orders, definition-region overrides — implement `StructuralVisitor`. -`visit` receives each value with its definition-region state and descends only -where it chooses: `default_visit_children` delegates the default child -recursion (the analog of C++ `DefaultVisitExpected`), and `visit_child` visits -one child under an explicit state (the analog of `Visit` under -`WithDefRegionKind`): +`visit` receives each value with its definition-region state and descends +only where it chooses: `default_visit_children` delegates the default child +recursion, and `visit_child` visits one child under an explicit state: ```rust,ignore use tvm_ffi::{DefRegionKind, Result, StructuralVisitor, VisitInterrupt, VisitValue}; @@ -151,18 +118,10 @@ impl StructuralVisitor for FuncVisitor { } ``` -Returning without descending skips a value's children; -`Ok(Some(VisitInterrupt))` halts the traversal. Nested -`visit_child`/`default_visit_children` calls report a nested interrupt through -their return value — propagate it (and errors, via `?`) instead of dropping -the result. - Recursion runs natively in Rust; no C++ visitor is constructed. Mutable -`List`/`Dict` contents are snapshotted before callbacks run, so re-entrant -mutation cannot invalidate a traversal. A non-container type that registers a -foreign `__s_visit__` hook is rejected rather than silently walked through -reflection; visit its children explicitly from a `StructuralVisitor`, or skip -it in a walk with a pre-order `WalkResult::Skip` handler. +`List`/`Dict` contents are snapshotted before callbacks run, and a +non-container type with a foreign `__s_visit__` hook is rejected rather than +silently walked through reflection. ## Installation From 81727c89458049c78b75b9196640cdc07d8ff9a6 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 15:33:38 -0400 Subject: [PATCH 12/20] update the doc. Signed-off-by: yuchuan --- docs/guides/rust_lang_guide.md | 106 +++++++++++++++++++++++++++++++++ rust/README.md | 95 ----------------------------- 2 files changed, 106 insertions(+), 95 deletions(-) diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md index f4c19e786..8686d6040 100644 --- a/docs/guides/rust_lang_guide.md +++ b/docs/guides/rust_lang_guide.md @@ -180,6 +180,112 @@ fn may_fail(value: i32) -> Result<()> { } ``` +### Structural Walk and Visit + +Rust provides equivalents of the C++ `StructuralWalk`/`StructuralVisitor` +APIs. Put `#[dispatch(visit)]` on an impl to turn its `visit_*` methods into +typed handlers, then pass it to `structural_walk`; each handler returns a +`WalkResult` (`Advance`, `Skip`, or `Interrupt`) to steer the traversal. +Handlers dispatch on their argument type and may take an optional trailing +`DefRegionKind` argument: + +```rust +use tvm_ffi::{dispatch, structural_walk, Array, DefRegionKind, WalkOrder, WalkResult}; + +#[derive(Default)] +struct Probe { + total: i64, + floats: usize, +} + +#[dispatch(visit)] +impl Probe { + fn visit_integer(&mut self, value: i64) -> WalkResult { + self.total += value; + WalkResult::Advance + } + + fn visit_float(&mut self, _value: f64, _kind: DefRegionKind) -> WalkResult { + self.floats += 1; + WalkResult::Advance + } +} + +let values = Array::new(vec![1_i64, 2, 3]); +let mut probe = Probe::default(); +structural_walk(&values, &mut probe, WalkOrder::PreOrder)?; +assert_eq!(probe.total, 6); +``` + +A closure works as a catch-all walker, over `&VisitValue` alone or with the +definition-region state as a second argument: + +```rust +use tvm_ffi::{structural_walk, Array, DefRegionKind, VisitValue, WalkOrder, WalkResult}; + +let values = Array::new(vec![1_i64, 2, 3]); +let mut integers = 0; +structural_walk( + &values, + |value: &VisitValue| { + if value.cast::().is_some() { + integers += 1; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, +)?; +assert_eq!(integers, 3); + +let mut uses = 0; +structural_walk( + &values, + |value: &VisitValue, kind: DefRegionKind| { + if value.cast::().is_some() && kind == DefRegionKind::None { + uses += 1; + } + WalkResult::Advance + }, + WalkOrder::PreOrder, +)?; +assert_eq!(uses, 3); +``` + +To drive recursion yourself, implement `StructuralVisitor` and call +`structural_visit`; `visit` runs for each value and descends through +`default_visit_children` (or `visit_child` for selected children): + +```rust +use tvm_ffi::{ + structural_visit, Array, DefRegionKind, Result, StructuralVisitor, VisitInterrupt, VisitValue, +}; + +#[derive(Default)] +struct Depth { + max: usize, + current: usize, +} + +impl StructuralVisitor for Depth { + fn visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result> { + self.current += 1; + self.max = self.max.max(self.current); + let interrupt = self.default_visit_children(value, def_region_kind)?; + self.current -= 1; + Ok(interrupt) + } +} + +let values = Array::new(vec![1_i64, 2]); +let mut depth = Depth::default(); +structural_visit(&values, &mut depth)?; +assert_eq!(depth.max, 2); +``` + ## Examples The repository includes a complete example in `rust/tvm-ffi/examples/load_library.rs`. diff --git a/rust/README.md b/rust/README.md index 47fbf2773..31f3fb01b 100644 --- a/rust/README.md +++ b/rust/README.md @@ -28,101 +28,6 @@ This workspace contains three crates: The overall project focuses on low-level, direct access to the ABI when possible for maximum efficiency while maintaining interoperability. -## Structural Visitors and Walkers - -The `tvm-ffi` crate provides native Rust structural traversal over FFI values, -built-in containers, and reflected object fields, mirroring the two layers of -the C++ API: - -- **Walk layer (observer)** — `structural_walk`, the analog of C++ - `StructuralWalk`: the walker owns recursion; handlers observe each value in - pre- or post-order (`WalkOrder`) and steer traversal through the returned - `WalkResult` (`Advance`, `Skip`, `Interrupt`/`InterruptWith`). -- **Visitor layer (user-driven)** — `structural_visit` with a - `StructuralVisitor`, the analog of a C++ `StructuralVisitorObj`: your - `visit` method runs for the root and controls all recursion itself. - -Every traversal returns `Result>`: `Ok(None)` means -the whole graph was visited, `Ok(Some(interrupt))` carries an interrupting -handler's payload. Handlers may also return `Result` to propagate -errors with `?`. - -### Observer walks - -`#[dispatch(visit)]` turns the `visit_*` methods of an inherent impl into a -typed, stateful observer. Each value dispatches to the first handler matching -its runtime type (a `&VisitValue` handler acts as the catch-all); a handler -that needs the definition-region state declares a trailing `DefRegionKind` -argument: - -```rust -use tvm_ffi::{dispatch, structural_walk, Array, WalkOrder, WalkResult}; - -#[derive(Default)] -struct Calculator { - value: f64, -} - -#[dispatch(visit)] -impl Calculator { - fn visit_integer(&mut self, value: i64) -> WalkResult { - self.value += value as f64; - WalkResult::Advance - } - - fn visit_float(&mut self, value: f64) -> WalkResult { - self.value -= value; - WalkResult::Advance - } -} - -let values = Array::new(vec![10_i64, 2]); -let mut calculator = Calculator::default(); -assert!(structural_walk(&values, &mut calculator, WalkOrder::PreOrder) - .unwrap() - .is_none()); -assert_eq!(calculator.value, 12.0); -``` - -`structural_walk` also accepts a bare closure taking `&VisitValue` (with an -optional trailing `DefRegionKind`) as a catch-all observer. - -### User-driven visitors - -When traversal itself is part of the analysis — visiting selected children, -custom orders, definition-region overrides — implement `StructuralVisitor`. -`visit` receives each value with its definition-region state and descends -only where it chooses: `default_visit_children` delegates the default child -recursion, and `visit_child` visits one child under an explicit state: - -```rust,ignore -use tvm_ffi::{DefRegionKind, Result, StructuralVisitor, VisitInterrupt, VisitValue}; - -struct FuncVisitor; - -impl StructuralVisitor for FuncVisitor { - fn visit( - &mut self, - value: &VisitValue, - def_region_kind: DefRegionKind, - ) -> Result> { - if let Some(func) = value.as_node::() { - // Parameters bind recursively; the body inherits the state. - if let Some(interrupt) = self.visit_child(&func.params, DefRegionKind::Recursive)? { - return Ok(Some(interrupt)); - } - return self.visit_child(&func.body, def_region_kind); - } - self.default_visit_children(value, def_region_kind) - } -} -``` - -Recursion runs natively in Rust; no C++ visitor is constructed. Mutable -`List`/`Dict` contents are snapshotted before callbacks run, and a -non-container type with a foreign `__s_visit__` hook is rejected rather than -silently walked through reflection. - ## Installation The Rust support depends on `libtvm_ffi`. From a8f7b59ddeaac35be468acc94e64846db522ef92 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 15:58:14 -0400 Subject: [PATCH 13/20] support lambda tuple. Signed-off-by: yuchuan --- docs/guides/rust_lang_guide.md | 43 ++- rust/tvm-ffi/src/extra/structural_visit.rs | 342 ++++++++++++++++++++- rust/tvm-ffi/src/lib.rs | 2 +- 3 files changed, 360 insertions(+), 27 deletions(-) diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md index 8686d6040..d529cd800 100644 --- a/docs/guides/rust_lang_guide.md +++ b/docs/guides/rust_lang_guide.md @@ -217,38 +217,47 @@ structural_walk(&values, &mut probe, WalkOrder::PreOrder)?; assert_eq!(probe.total, 6); ``` -A closure works as a catch-all walker, over `&VisitValue` alone or with the -definition-region state as a second argument: +Lambdas also work — pass a single typed lambda, or a tuple of them tried in +order with the first matching argument type winning, like the variadic C++ +`StructuralWalk(root, callbacks...)` chain. Unmatched values simply advance +(a `&VisitValue` lambda acts as a catch-all), and each lambda may take a +trailing `DefRegionKind` argument: ```rust -use tvm_ffi::{structural_walk, Array, DefRegionKind, VisitValue, WalkOrder, WalkResult}; +use tvm_ffi::{structural_walk, Array, DefRegionKind, Object, WalkOrder, WalkResult}; let values = Array::new(vec![1_i64, 2, 3]); -let mut integers = 0; + +let mut total = 0; structural_walk( &values, - |value: &VisitValue| { - if value.cast::().is_some() { - integers += 1; - } + |value: i64| { + total += value; WalkResult::Advance }, WalkOrder::PreOrder, )?; -assert_eq!(integers, 3); +assert_eq!(total, 6); -let mut uses = 0; +let mut evens = 0; +let mut objects = 0; structural_walk( &values, - |value: &VisitValue, kind: DefRegionKind| { - if value.cast::().is_some() && kind == DefRegionKind::None { - uses += 1; - } - WalkResult::Advance - }, + ( + |value: i64| { + if value % 2 == 0 { + evens += 1; + } + WalkResult::Advance + }, + |_object: &Object, _kind: DefRegionKind| { + objects += 1; + WalkResult::Advance + }, + ), WalkOrder::PreOrder, )?; -assert_eq!(uses, 3); +assert_eq!((evens, objects), (1, 1)); ``` To drive recursion yourself, implement `StructuralVisitor` and call diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index 6fccdd4ed..cf269c7df 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -50,6 +50,7 @@ //! different semantics; visit such a type's children explicitly from a //! [`StructuralVisitor`], or skip the value in a walk. +use std::marker::PhantomData; use std::ops::ControlFlow; use std::os::raw::c_void; use std::ptr::NonNull; @@ -287,10 +288,13 @@ impl VisitDispatch for &mut V { /// /// * `&mut V` where `V: VisitDispatch` — a stateful typed visitor /// (`#[dispatch(visit)]` or hand-written). -/// * `FnMut(&VisitValue) -> impl IntoVisitResult` — a bare observer closure, -/// the analog of a C++ `(value)` callback. -/// * `FnMut(&VisitValue, DefRegionKind) -> impl IntoVisitResult` — the -/// analog of a C++ `(value, def_region_kind)` callback. +/// * A bare closure in any [`WalkChainLink`] shape — catch-all +/// `FnMut(&VisitValue)`, typed `FnMut(T)`, node `FnMut(&N)`, each with an +/// optional trailing [`DefRegionKind`] argument — the analog of a single +/// C++ callback. Values a typed closure does not match advance normally. +/// * A tuple of typed links `(link1, link2, ...)` — the analog of the C++ +/// variadic callback chain; see [`WalkChainLink`] for the accepted link +/// shapes. /// /// Closure arguments usually need explicit type annotations /// (`|value: &VisitValue| ...`) for the marker to be inferred. @@ -408,6 +412,324 @@ where } } +/// One typed link of a tuple walker — a single callback of the C++ variadic +/// `StructuralWalk(root, callbacks...)` chain. +/// +/// A tuple of links passed to [`structural_walk`] is tried in order and the +/// first link whose argument type matches the value runs, exactly like the +/// C++ callback chain and the Python `(type, callback)` entries. Accepted +/// link shapes mirror `#[dispatch(visit)]` handlers: +/// +/// * `FnMut(T) -> impl IntoVisitResult` for an FFI-convertible `T` — exact +/// value cast via [`VisitValue::cast`]. +/// * `FnMut(&N) -> impl IntoVisitResult` for an object node `N` — +/// refcount-free subtype check via [`VisitValue::as_node`]. +/// * `FnMut(&VisitValue) -> impl IntoVisitResult` — catch-all; place it +/// last, links after it never run. +/// * `&mut V` where `V: VisitDispatch` — splice a typed visitor into the +/// chain. +/// +/// Every closure shape may declare a trailing [`DefRegionKind`] argument, +/// and a single typed closure may also be passed to [`structural_walk`] +/// bare, without the tuple. Closure arguments need explicit type +/// annotations for the marker to be inferred. Borrow rules apply per link, +/// so state shared across links goes through a `Cell`/`RefCell` — or in a +/// single `#[dispatch(visit)]` visitor, which shares `&mut self` between +/// its handlers. +pub trait WalkChainLink { + /// Run this link if `value` matches its argument type; `None` hands the + /// value to the next link. + #[doc(hidden)] + fn try_call( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option; +} + +#[doc(hidden)] +pub struct ByOwnedLink(PhantomData); + +impl WalkChainLink> for F +where + F: FnMut(T) -> O, + T: crate::type_traits::AnyCompatible, + O: IntoVisitResult, +{ + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + _def_region_kind: DefRegionKind, + ) -> Option { + value + .cast::() + .map(|typed| self(typed).into_visit_result()) + } +} + +#[doc(hidden)] +pub struct ByOwnedKindLink(PhantomData); + +impl WalkChainLink> for F +where + F: FnMut(T, DefRegionKind) -> O, + T: crate::type_traits::AnyCompatible, + O: IntoVisitResult, +{ + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option { + value + .cast::() + .map(|typed| self(typed, def_region_kind).into_visit_result()) + } +} + +#[doc(hidden)] +pub struct ByNodeLink(PhantomData); + +impl WalkChainLink> for F +where + F: for<'a> FnMut(&'a N) -> O, + N: ObjectCore, + O: IntoVisitResult, +{ + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + _def_region_kind: DefRegionKind, + ) -> Option { + value + .as_node::() + .map(|node| self(node).into_visit_result()) + } +} + +#[doc(hidden)] +pub struct ByNodeKindLink(PhantomData); + +impl WalkChainLink> for F +where + F: for<'a> FnMut(&'a N, DefRegionKind) -> O, + N: ObjectCore, + O: IntoVisitResult, +{ + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option { + value + .as_node::() + .map(|node| self(node, def_region_kind).into_visit_result()) + } +} + +#[doc(hidden)] +pub enum ByCatchAllLink {} + +impl WalkChainLink for F +where + F: for<'a> FnMut(&'a VisitValue) -> O, + O: IntoVisitResult, +{ + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + _def_region_kind: DefRegionKind, + ) -> Option { + Some(self(value).into_visit_result()) + } +} + +#[doc(hidden)] +pub enum ByCatchAllKindLink {} + +impl WalkChainLink for F +where + F: for<'a> FnMut(&'a VisitValue, DefRegionKind) -> O, + O: IntoVisitResult, +{ + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option { + Some(self(value, def_region_kind).into_visit_result()) + } +} + +#[doc(hidden)] +pub enum ByDispatchLink {} + +impl WalkChainLink for &mut V { + #[inline] + fn try_call( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option { + self.dispatch_visit(value, def_region_kind) + } +} + +/// Runs a tuple of [`WalkChainLink`]s at the phase selected by `order`, +/// trying links in order and short-circuiting on the first whose type +/// matches — the Rust analog of C++ `StructuralWalkCallbackChain`. Static +/// dispatch throughout: each link's type test inlines to the same code the +/// `#[dispatch(visit)]` macro generates for a `visit_*` chain. +#[doc(hidden)] +pub struct ChainWalker { + links: Links, + order: WalkOrder, + markers: PhantomData, +} + +macro_rules! impl_chain_walker { + ($(($F:ident, $M:ident, $idx:tt)),+) => { + impl<$($F, $M,)+> ChainWalker<($($F,)+), ($($M,)+)> + where + $($F: WalkChainLink<$M>,)+ + { + #[inline] + fn dispatch( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result { + $( + if let Some(result) = self.links.$idx.try_call(value, def_region_kind) { + return result; + } + )+ + Ok(WalkResult::Advance) + } + } + + impl<$($F, $M,)+> NativeVisit for ChainWalker<($($F,)+), ($($M,)+)> + where + $($F: WalkChainLink<$M>,)+ + { + fn enter( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result { + match self.order { + WalkOrder::PreOrder => self.dispatch(value, def_region_kind), + WalkOrder::PostOrder => Ok(WalkResult::Advance), + } + } + + fn exit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Result { + match self.order { + WalkOrder::PreOrder => Ok(WalkResult::Advance), + WalkOrder::PostOrder => self.dispatch(value, def_region_kind), + } + } + } + + impl<$($F, $M,)+> IntoWalker<($($M,)+)> for ($($F,)+) + where + $($F: WalkChainLink<$M>,)+ + { + type Walker = ChainWalker<($($F,)+), ($($M,)+)>; + fn into_walker(self, order: WalkOrder) -> Self::Walker { + ChainWalker { + links: self, + order, + markers: PhantomData, + } + } + } + }; +} + +impl_chain_walker!((F0, M0, 0)); +impl_chain_walker!((F0, M0, 0), (F1, M1, 1)); +impl_chain_walker!((F0, M0, 0), (F1, M1, 1), (F2, M2, 2)); +impl_chain_walker!((F0, M0, 0), (F1, M1, 1), (F2, M2, 2), (F3, M3, 3)); +impl_chain_walker!( + (F0, M0, 0), + (F1, M1, 1), + (F2, M2, 2), + (F3, M3, 3), + (F4, M4, 4) +); +impl_chain_walker!( + (F0, M0, 0), + (F1, M1, 1), + (F2, M2, 2), + (F3, M3, 3), + (F4, M4, 4), + (F5, M5, 5) +); +impl_chain_walker!( + (F0, M0, 0), + (F1, M1, 1), + (F2, M2, 2), + (F3, M3, 3), + (F4, M4, 4), + (F5, M5, 5), + (F6, M6, 6) +); +impl_chain_walker!( + (F0, M0, 0), + (F1, M1, 1), + (F2, M2, 2), + (F3, M3, 3), + (F4, M4, 4), + (F5, M5, 5), + (F6, M6, 6), + (F7, M7, 7) +); + +// A bare typed closure — `FnMut(T)` or `FnMut(&N)`, optionally with a +// trailing `DefRegionKind` — walks as a single-link chain, so a lone typed +// handler needs no tuple wrapping; values that do not match its argument +// type advance normally. `&VisitValue` catch-all closures keep their +// dedicated `ClosureWalker`/`ClosureKindWalker` path above. +macro_rules! impl_bare_link_walker { + ($(($marker:ident, $($fn_args:ty),+)),+ $(,)?) => { + $( + impl IntoWalker<$marker> for F + where + F: FnMut($($fn_args),+) -> O, + Self: WalkChainLink<$marker>, + O: IntoVisitResult, + { + type Walker = ChainWalker<(F,), ($marker,)>; + fn into_walker(self, order: WalkOrder) -> Self::Walker { + ChainWalker { + links: (self,), + order, + markers: PhantomData, + } + } + } + )+ + }; +} + +impl_bare_link_walker!( + (ByOwnedLink, T), + (ByOwnedKindLink, T, DefRegionKind), + (ByNodeLink, &T), + (ByNodeKindLink, &T, DefRegionKind), +); + /// A visitor that drives recursion itself, mirroring C++ /// `StructuralVisitorObj`. /// @@ -942,11 +1264,13 @@ where /// `StructuralWalk(root, callbacks...)`. /// /// `walker` is anything implementing [`IntoWalker`]: a `&mut` reference to a -/// stateful [`VisitDispatch`] visitor (`#[dispatch(visit)]`), or a bare -/// closure taking `&VisitValue` with an optional trailing [`DefRegionKind`] -/// — the C++ callback overloads. The walker owns recursion: the handler runs -/// once per value, before or after the value's children according to -/// `order`, and steers traversal through the returned [`WalkResult`]. +/// stateful [`VisitDispatch`] visitor (`#[dispatch(visit)]`), a bare closure +/// in any [`WalkChainLink`] shape (catch-all `&VisitValue`, typed, or node, +/// with an optional trailing [`DefRegionKind`]), or a tuple of such +/// callbacks tried in order — the C++ callback overloads and variadic +/// chain. The walker owns recursion: the handler runs once per value, +/// before or after the value's children according to `order`, and steers +/// traversal through the returned [`WalkResult`]. pub fn structural_walk( root: &R, walker: H, diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index 70eff8df7..f05eb0608 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -48,7 +48,7 @@ pub use crate::error::{ pub use crate::extra::module::Module; pub use crate::extra::structural_visit::{ structural_visit, structural_walk, DefRegionKind, IntoWalker, StructuralVisitor, VisitDispatch, - VisitInterrupt, VisitValue, WalkOrder, WalkResult, + VisitInterrupt, VisitValue, WalkChainLink, WalkOrder, WalkResult, }; pub use crate::function::Function; pub use crate::object::ObjectRefCast; From c8a5abc35f40a8f4c136a37547acc0a174101baa Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 16:46:11 -0400 Subject: [PATCH 14/20] doc and tests. Signed-off-by: yuchuan --- docs/guides/rust_lang_guide.md | 45 ++- rust/tvm-ffi/src/extra/structural_visit.rs | 102 +++++- rust/tvm-ffi/src/lib.rs | 5 +- rust/tvm-ffi/tests/test_walk_chain.rs | 401 +++++++++++++++++++++ 4 files changed, 532 insertions(+), 21 deletions(-) create mode 100644 rust/tvm-ffi/tests/test_walk_chain.rs diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md index d529cd800..f2ae5eff5 100644 --- a/docs/guides/rust_lang_guide.md +++ b/docs/guides/rust_lang_guide.md @@ -217,11 +217,12 @@ structural_walk(&values, &mut probe, WalkOrder::PreOrder)?; assert_eq!(probe.total, 6); ``` -Lambdas also work — pass a single typed lambda, or a tuple of them tried in -order with the first matching argument type winning, like the variadic C++ -`StructuralWalk(root, callbacks...)` chain. Unmatched values simply advance -(a `&VisitValue` lambda acts as a catch-all), and each lambda may take a -trailing `DefRegionKind` argument: +Lambdas also work — pass a single typed lambda, or a tuple of them (up to 8) +tried in order with the first matching argument type winning, like the +variadic C++ `StructuralWalk(root, callbacks...)` chain. Unmatched values +simply advance; a `&VisitValue` lambda acts as a catch-all and must come +last, since links after an always-matching one never run. Each lambda may +take a trailing `DefRegionKind` argument: ```rust use tvm_ffi::{structural_walk, Array, DefRegionKind, Object, WalkOrder, WalkResult}; @@ -260,9 +261,34 @@ structural_walk( assert_eq!((evens, objects), (1, 1)); ``` +Both entry points return `Result>`: `Ok(None)` means +the whole graph was visited, and a handler stops the walk early by returning +`WalkResult::interrupt_with(payload)`, which comes back to the caller as +`Ok(Some(interrupt))`. Handlers may also return `Result` and +propagate errors with `?`: + +```rust +use tvm_ffi::{structural_walk, Array, WalkOrder, WalkResult}; + +let values = Array::new(vec![1_i64, 2, 3]); +let found = structural_walk( + &values, + |value: i64| { + if value == 2 { + return WalkResult::interrupt_with(value); + } + WalkResult::Advance + }, + WalkOrder::PreOrder, +)?; +assert_eq!(found.map(|i| i64::try_from(i.value).unwrap()), Some(2)); +``` + To drive recursion yourself, implement `StructuralVisitor` and call `structural_visit`; `visit` runs for each value and descends through -`default_visit_children` (or `visit_child` for selected children): +`default_visit_children`, or through `visit_child`, which visits one +selected child and can override the def-region state for it (e.g. +`DefRegionKind::Recursive` when descending into a binder's parameters): ```rust use tvm_ffi::{ @@ -295,6 +321,13 @@ structural_visit(&values, &mut depth)?; assert_eq!(depth.max, 2); ``` +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 +than silently walked through reflection — visit such a type's children +explicitly from a `StructuralVisitor`, or skip it with a pre-order +`WalkResult::Skip`. + ## Examples The repository includes a complete example in `rust/tvm-ffi/examples/load_library.rs`. diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index cf269c7df..45d0bfac0 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -44,6 +44,10 @@ //! object-description boundary: traversal, control flow, typed dispatch, //! visitor state, and definition-region propagation remain in Rust. //! +//! Mutable `List`/`Dict` contents are snapshotted before callbacks run, so a +//! callback mutating the container it was reached through cannot invalidate +//! the traversal; the walk sees the pre-mutation contents. +//! //! No C++ `ffi.StructuralVisitor` is constructed and no C++ default-visit //! function is called. A non-container type with a foreign `__s_visit__` hook //! is rejected instead of silently substituting reflection with potentially @@ -248,7 +252,9 @@ type NativeResult = std::result::Result<(), NativeHalt>; /// [`crate::dispatch`] tests the implementation's `visit_*` methods in source /// order. Borrowed node arguments use refcount-free subtype checks, owned /// FFI-compatible arguments use exact value casts, and `&VisitValue` is a -/// catch-all. `None` asks the Rust walker to continue normally. +/// catch-all. `None` reports that no handler matched: a standalone walk then +/// advances normally, while a tuple chain hands the value to the next link — +/// so a spliced visitor that "handled" a value must not return `None`. /// /// This is the observer layer, mirroring C++ `StructuralWalk` callbacks: the /// walker owns recursion, and a handler steers it only through the returned @@ -292,12 +298,23 @@ impl VisitDispatch for &mut V { /// `FnMut(&VisitValue)`, typed `FnMut(T)`, node `FnMut(&N)`, each with an /// optional trailing [`DefRegionKind`] argument — the analog of a single /// C++ callback. Values a typed closure does not match advance normally. -/// * A tuple of typed links `(link1, link2, ...)` — the analog of the C++ -/// variadic callback chain; see [`WalkChainLink`] for the accepted link -/// shapes. +/// * A tuple of typed links `(link1, link2, ...)`, up to 8 — the analog of +/// the C++ variadic callback chain; see [`WalkChainLink`] for the +/// accepted link shapes. Larger handler sets belong in one +/// `#[dispatch(visit)]` visitor, which itself splices into a tuple as a +/// single link. /// /// Closure arguments usually need explicit type annotations /// (`|value: &VisitValue| ...`) for the marker to be inferred. +#[diagnostic::on_unimplemented( + message = "`{Self}` is not a supported `structural_walk` walker", + note = "accepted walkers: `&mut V` where `V: VisitDispatch`; a closure over `&VisitValue`, \ + an FFI value type `T`, or `&N` of an object node type (`N: ObjectCore`, e.g. \ + `&Object`), optionally with a trailing `DefRegionKind` argument; or a tuple of \ + up to 8 such links", + note = "closure arguments need explicit type annotations; ObjectRef wrappers like `String` \ + or `Array` are FFI value types — take them by value, not by reference" +)] pub trait IntoWalker { #[doc(hidden)] type Walker: NativeVisit; @@ -415,19 +432,29 @@ where /// One typed link of a tuple walker — a single callback of the C++ variadic /// `StructuralWalk(root, callbacks...)` chain. /// -/// A tuple of links passed to [`structural_walk`] is tried in order and the -/// first link whose argument type matches the value runs, exactly like the -/// C++ callback chain and the Python `(type, callback)` entries. Accepted -/// link shapes mirror `#[dispatch(visit)]` handlers: +/// A tuple of up to 8 links passed to [`structural_walk`] is tried in order +/// and the first link whose argument type matches the value runs, exactly +/// like the C++ callback chain. (Python's `structural_walk` differs on one +/// point: it keeps `callbacks` and `with_def_region_kind` as two separately +/// ordered groups, trying every plain entry before any kind-taking entry, +/// so a mixed Rust tuple's single interleaved order has no exact Python +/// equivalent.) Accepted link shapes mirror `#[dispatch(visit)]` handlers: /// /// * `FnMut(T) -> impl IntoVisitResult` for an FFI-convertible `T` — exact -/// value cast via [`VisitValue::cast`]. +/// value cast via [`VisitValue::cast`]. Exact includes the value itself: +/// a numeric link matches only losslessly representable values, so an +/// `Int(300)` falls through a `u8` link to the next one rather than +/// truncating. /// * `FnMut(&N) -> impl IntoVisitResult` for an object node `N` — /// refcount-free subtype check via [`VisitValue::as_node`]. -/// * `FnMut(&VisitValue) -> impl IntoVisitResult` — catch-all; place it -/// last, links after it never run. +/// * `FnMut(&VisitValue) -> impl IntoVisitResult` — catch-all. /// * `&mut V` where `V: VisitDispatch` — splice a typed visitor into the -/// chain. +/// chain; it claims every value one of its handlers matches. +/// +/// Links after one that matches every value never run: place a catch-all +/// closure — or a spliced visitor whose own chain ends in a `&VisitValue` +/// handler — last. Unlike the in-visitor ordering check, misordering a +/// tuple is not a compile error. /// /// Every closure shape may declare a trailing [`DefRegionKind`] argument, /// and a single typed closure may also be passed to [`structural_walk`] @@ -436,7 +463,10 @@ where /// so state shared across links goes through a `Cell`/`RefCell` — or in a /// single `#[dispatch(visit)]` visitor, which shares `&mut self` between /// its handlers. -pub trait WalkChainLink { +/// +/// This trait is sealed: the link shapes above are the complete set, and +/// the dispatch method is an internal detail. +pub trait WalkChainLink: sealed::SealedLink { /// Run this link if `value` matches its argument type; `None` hands the /// value to the next link. #[doc(hidden)] @@ -447,6 +477,52 @@ pub trait WalkChainLink { ) -> Option; } +mod sealed { + use super::{DefRegionKind, IntoVisitResult, ObjectCore, VisitDispatch, VisitValue}; + + /// Seal for [`super::WalkChainLink`]: one impl per accepted link shape, + /// mirroring the `WalkChainLink` impl set exactly. + pub trait SealedLink {} + + impl SealedLink> for F + where + F: FnMut(T) -> O, + O: IntoVisitResult, + { + } + impl SealedLink> for F + where + F: FnMut(T, DefRegionKind) -> O, + O: IntoVisitResult, + { + } + impl SealedLink> for F + where + F: for<'a> FnMut(&'a N) -> O, + O: IntoVisitResult, + { + } + impl SealedLink> for F + where + F: for<'a> FnMut(&'a N, DefRegionKind) -> O, + O: IntoVisitResult, + { + } + impl SealedLink for F + where + F: for<'a> FnMut(&'a VisitValue) -> O, + O: IntoVisitResult, + { + } + impl SealedLink for F + where + F: for<'a> FnMut(&'a VisitValue, DefRegionKind) -> O, + O: IntoVisitResult, + { + } + impl SealedLink for &mut V {} +} + #[doc(hidden)] pub struct ByOwnedLink(PhantomData); diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs index f05eb0608..865798576 100644 --- a/rust/tvm-ffi/src/lib.rs +++ b/rust/tvm-ffi/src/lib.rs @@ -47,8 +47,9 @@ pub use crate::error::{ }; pub use crate::extra::module::Module; pub use crate::extra::structural_visit::{ - structural_visit, structural_walk, DefRegionKind, IntoWalker, StructuralVisitor, VisitDispatch, - VisitInterrupt, VisitValue, WalkChainLink, WalkOrder, WalkResult, + structural_visit, structural_walk, DefRegionKind, IntoVisitResult, IntoWalker, + StructuralVisitor, VisitDispatch, VisitInterrupt, VisitValue, WalkChainLink, WalkOrder, + WalkResult, }; pub use crate::function::Function; pub use crate::object::ObjectRefCast; diff --git a/rust/tvm-ffi/tests/test_walk_chain.rs b/rust/tvm-ffi/tests/test_walk_chain.rs new file mode 100644 index 000000000..16ba349d4 --- /dev/null +++ b/rust/tvm-ffi/tests/test_walk_chain.rs @@ -0,0 +1,401 @@ +/* + * 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. + */ + +//! Tuple walkers: `structural_walk(root, (link1, link2, ...), order)`, the +//! Rust analog of the C++ variadic `StructuralWalk(root, callbacks...)` +//! chain. Links are tried in order and the first whose argument type matches +//! the value runs. + +use tvm_ffi::{ + dispatch, structural_walk, Array, DefRegionKind, Error, Object, Result, VisitValue, WalkOrder, + WalkResult, RUNTIME_ERROR, +}; + +fn runtime_error(message: &str) -> Error { + Error::new(RUNTIME_ERROR, message, "") +} + +#[test] +fn chain_dispatches_first_matching_link() { + // C++: StructuralWalk(root, [&](int64_t v) {...}, + // [&](const ObjectRef& o) {...}, [&](AnyView v) {...}) + let root = Array::new(vec![1i64, 2, 3]); + let mut integers = Vec::new(); + let mut objects = 0; + let mut others = 0; + assert!(structural_walk( + &root, + ( + |value: i64| { + integers.push(value); + WalkResult::Advance + }, + |_object: &Object| { + objects += 1; + WalkResult::Advance + }, + |_value: &VisitValue| { + others += 1; + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(integers, vec![1, 2, 3]); + assert_eq!(objects, 1); // the array itself; integers matched earlier + assert_eq!(others, 0); // every value matched an earlier link +} + +#[test] +fn chain_accepts_owned_object_ref_links() { + let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64, 3])]); + let mut lengths = Vec::new(); + assert!(structural_walk( + &root, + ( + |array: Array| { + lengths.push(array.len()); + WalkResult::Advance + }, + |_value: i64| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + // The outer Array> fails the strict element check and falls + // through the chain; only the inner arrays match the typed link. + assert_eq!(lengths, vec![1, 2]); +} + +#[test] +fn chain_links_may_mix_def_region_arity() { + // Like #[dispatch(visit)] handlers, each link independently opts into + // the trailing DefRegionKind argument. + let root = Array::new(vec![1i64, 2]); + let mut kinds = Vec::new(); + let mut objects = 0; + assert!(structural_walk( + &root, + ( + |_value: i64, kind: DefRegionKind| { + kinds.push(kind); + WalkResult::Advance + }, + |_value: &VisitValue, kind: DefRegionKind| { + assert_eq!(kind, DefRegionKind::None); + objects += 1; + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(kinds, vec![DefRegionKind::None; 2]); + assert_eq!(objects, 1); +} + +#[test] +fn chain_links_can_skip_children() { + let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); + let mut arrays = 0; + let mut integers = 0; + assert!(structural_walk( + &root, + ( + |_array: Array| { + arrays += 1; + WalkResult::Skip + }, + |_value: i64| { + integers += 1; + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(arrays, 2); + assert_eq!(integers, 0); // both inner arrays were skipped +} + +#[test] +fn chain_links_can_interrupt_with_payload() { + let root = Array::new(vec![1i64, 2, 3]); + let mut seen = 0; + let outcome = structural_walk( + &root, + ( + |value: i64| { + seen += 1; + if value == 2 { + return WalkResult::interrupt_with(value * 10); + } + WalkResult::Advance + }, + |_value: &VisitValue| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) + .unwrap(); + let Some(interrupt) = outcome else { + panic!("walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(interrupt.value).unwrap(), 20); + assert_eq!(seen, 2); +} + +#[test] +fn chain_link_errors_include_native_visit_path() { + let root = Array::new(vec![1i64]); + let error = match structural_walk( + &root, + ( + |_value: i64| -> Result { Err(runtime_error("link failed")) }, + |_value: &VisitValue| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) { + Err(error) => error, + Ok(_) => panic!("link unexpectedly succeeded"), + }; + assert_eq!(error.message(), "link failed"); + assert!(error.backtrace().contains("sequence item [0]")); + assert!(error.backtrace().contains("object `ffi.Array`")); +} + +#[test] +fn chain_supports_post_order() { + // Rust borrow rules apply per link: state shared across links goes + // through a RefCell (or a single #[dispatch(visit)] visitor). + let root = Array::new(vec![1i64, 2]); + let events = std::cell::RefCell::new(Vec::new()); + assert!(structural_walk( + &root, + ( + |value: i64| { + events.borrow_mut().push(format!("int:{value}")); + WalkResult::Advance + }, + |_object: &Object| { + events.borrow_mut().push("array".to_string()); + WalkResult::Advance + }, + ), + WalkOrder::PostOrder, + ) + .unwrap() + .is_none()); + assert_eq!(events.into_inner(), vec!["int:1", "int:2", "array"]); +} + +#[test] +fn single_link_tuple_walks() { + let root = Array::new(vec![1i64, 2, 3]); + let mut total = 0; + assert!(structural_walk( + &root, + (|value: i64| { + total += value; + WalkResult::Advance + },), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(total, 6); +} + +#[derive(Default)] +struct ObjectCounter { + objects: usize, +} + +#[dispatch(visit)] +impl ObjectCounter { + fn visit_object(&mut self, _value: &Object) -> WalkResult { + self.objects += 1; + WalkResult::Advance + } +} + +#[test] +fn chain_splices_dispatch_visitors_between_closures() { + // A `&mut` typed visitor participates in the chain like any other link, + // keeping its own no-match fall-through semantics. + let root = Array::new(vec![1i64, 2]); + let mut counter = ObjectCounter::default(); + let mut integers = 0; + assert!(structural_walk( + &root, + (&mut counter, |_value: i64| { + integers += 1; + WalkResult::Advance + },), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(counter.objects, 1); + assert_eq!(integers, 2); +} + +#[test] +fn chain_supports_full_arity() { + let root = Array::new(vec![1i64]); + let mut int_hits = 0; + let mut object_hits = 0; + assert!(structural_walk( + &root, + ( + |_value: f64| WalkResult::Advance, + |_value: bool| WalkResult::Advance, + |_value: tvm_ffi::String| WalkResult::Advance, + |_value: Array| WalkResult::Advance, + |_value: i64| { + int_hits += 1; + WalkResult::Advance + }, + |_value: &Object, _kind: DefRegionKind| { + object_hits += 1; + WalkResult::Advance + }, + |_value: &VisitValue, _kind: DefRegionKind| WalkResult::Advance, + |_value: &VisitValue| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(int_hits, 1); + assert_eq!(object_hits, 1); +} + +#[test] +fn narrow_numeric_links_match_only_exact_values() { + // An out-of-range Int falls through a narrow numeric link to the next + // one instead of silently truncating into it. + let root = Array::new(vec![200i64, 300, -1]); + let mut narrow = Vec::new(); + let mut wide = Vec::new(); + assert!(structural_walk( + &root, + ( + |value: u8| { + narrow.push(value); + WalkResult::Advance + }, + |value: i64| { + wide.push(value); + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(narrow, vec![200u8]); + assert_eq!(wide, vec![300, -1]); +} + +#[test] +fn f32_links_match_only_lossless_values() { + let root = Array::new(vec![1.5f64, 1e300]); + let mut narrow = Vec::new(); + let mut wide = Vec::new(); + assert!(structural_walk( + &root, + ( + |value: f32| { + narrow.push(value); + WalkResult::Advance + }, + |value: f64| { + wide.push(value); + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(narrow, vec![1.5f32]); + assert_eq!(wide, vec![1e300]); +} + +#[test] +fn bare_typed_lambda_walks_without_tuple() { + // A lone typed handler needs no tuple: unmatched values (the array + // itself) advance normally. + let root = Array::new(vec![1i64, 2, 3]); + let mut total = 0; + assert!(structural_walk( + &root, + |value: i64| { + total += value; + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(total, 6); +} + +#[test] +fn bare_node_lambda_takes_def_region_kind() { + let root = Array::new(vec![1i64, 2]); + let mut objects = 0; + assert!(structural_walk( + &root, + |_object: &Object, kind: DefRegionKind| { + assert_eq!(kind, DefRegionKind::None); + objects += 1; + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(objects, 1); +} + +#[test] +fn bare_owned_object_lambda_interrupts() { + let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); + let outcome = structural_walk( + &root, + |array: Array| { + if array.len() == 1 { + return WalkResult::interrupt_with(array.len() as i64); + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap(); + let Some(interrupt) = outcome else { + panic!("walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(interrupt.value).unwrap(), 1); +} From 98bf890cb9616b5289c76cbf29a5791849a15790 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 16:58:28 -0400 Subject: [PATCH 15/20] [FIX][RUST] Drop lossless numeric link tests, align rustdoc (#693) The walk-chain tests asserted lossless numeric matching, but VisitValue::cast matches on the FFI type tag and converts with `as` semantics. Remove the two tests and fix the WalkChainLink rustdoc that claimed exact-value matching, recommending i64/f64 links unless a deliberate narrowing is wanted. Co-Authored-By: Claude Fable 5 --- rust/tvm-ffi/src/extra/structural_visit.rs | 10 ++--- rust/tvm-ffi/tests/test_walk_chain.rs | 52 ---------------------- 2 files changed, 5 insertions(+), 57 deletions(-) diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index 45d0bfac0..b74e4e442 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -440,11 +440,11 @@ where /// so a mixed Rust tuple's single interleaved order has no exact Python /// equivalent.) Accepted link shapes mirror `#[dispatch(visit)]` handlers: /// -/// * `FnMut(T) -> impl IntoVisitResult` for an FFI-convertible `T` — exact -/// value cast via [`VisitValue::cast`]. Exact includes the value itself: -/// a numeric link matches only losslessly representable values, so an -/// `Int(300)` falls through a `u8` link to the next one rather than -/// truncating. +/// * `FnMut(T) -> impl IntoVisitResult` for an FFI-convertible `T` — value +/// cast via [`VisitValue::cast`], which matches on the FFI type tag: a +/// numeric link claims every `Int` (or `Float`) regardless of width and +/// converts with `as` semantics, so prefer `i64`/`f64` links unless a +/// deliberate narrowing is wanted. /// * `FnMut(&N) -> impl IntoVisitResult` for an object node `N` — /// refcount-free subtype check via [`VisitValue::as_node`]. /// * `FnMut(&VisitValue) -> impl IntoVisitResult` — catch-all. diff --git a/rust/tvm-ffi/tests/test_walk_chain.rs b/rust/tvm-ffi/tests/test_walk_chain.rs index 16ba349d4..4a5dab9f6 100644 --- a/rust/tvm-ffi/tests/test_walk_chain.rs +++ b/rust/tvm-ffi/tests/test_walk_chain.rs @@ -291,58 +291,6 @@ fn chain_supports_full_arity() { assert_eq!(object_hits, 1); } -#[test] -fn narrow_numeric_links_match_only_exact_values() { - // An out-of-range Int falls through a narrow numeric link to the next - // one instead of silently truncating into it. - let root = Array::new(vec![200i64, 300, -1]); - let mut narrow = Vec::new(); - let mut wide = Vec::new(); - assert!(structural_walk( - &root, - ( - |value: u8| { - narrow.push(value); - WalkResult::Advance - }, - |value: i64| { - wide.push(value); - WalkResult::Advance - }, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(narrow, vec![200u8]); - assert_eq!(wide, vec![300, -1]); -} - -#[test] -fn f32_links_match_only_lossless_values() { - let root = Array::new(vec![1.5f64, 1e300]); - let mut narrow = Vec::new(); - let mut wide = Vec::new(); - assert!(structural_walk( - &root, - ( - |value: f32| { - narrow.push(value); - WalkResult::Advance - }, - |value: f64| { - wide.push(value); - WalkResult::Advance - }, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(narrow, vec![1.5f32]); - assert_eq!(wide, vec![1e300]); -} - #[test] fn bare_typed_lambda_walks_without_tuple() { // A lone typed handler needs no tuple: unmatched values (the array From 5d922f614ab33315191e10a87ad3711b8bac18b5 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 17:31:16 -0400 Subject: [PATCH 16/20] [PERF][RUST] Inline the foreign-hook check, fold WalkResult drop glue (#693) Callgrind on the nested-array shape showed two per-visit costs C++ does not pay: reject_foreign_structural_visit stayed a real call (~20% of the container fast path) because its cold error-formatting body defeated the inline hint, and the split Result/WalkResult matches in visit_raw left a partially-moved temporary whose drop glue could not fold away. Split the cold body behind #[cold] #[inline(never)] and match the handler results by value in one step. Nested-array walks drop from ~1.26-1.30x of C++ to ~1.09-1.17x; reflected-object and map shapes stay at or below parity. Co-Authored-By: Claude Fable 5 --- rust/tvm-ffi/src/extra/structural_visit.rs | 79 ++++++++++++---------- 1 file changed, 44 insertions(+), 35 deletions(-) diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index b74e4e442..e17944351 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -986,15 +986,16 @@ fn visit_raw( } let visit_value = VisitValue::from_raw(value); - let enter = match visitor.enter(&visit_value, def_region_kind) { - Ok(flow) => flow, + // Single by-value matches: splitting the Result match from the + // WalkResult match leaves a partially-moved temporary whose drop glue + // the compiler cannot fold away (measurably so on the container fast + // path). + match visitor.enter(&visit_value, def_region_kind) { + Ok(WalkResult::Advance) => {} + Ok(WalkResult::Skip) => return Ok(()), + Ok(WalkResult::Interrupt) => return Err(NativeHalt::Interrupt(Any::new())), + Ok(WalkResult::InterruptWith(payload)) => return Err(NativeHalt::Interrupt(payload)), Err(error) => return Err(with_value_context(error.into(), value)), - }; - match enter { - WalkResult::Advance => {} - WalkResult::Skip => return Ok(()), - WalkResult::Interrupt => return Err(NativeHalt::Interrupt(Any::new())), - WalkResult::InterruptWith(payload) => return Err(NativeHalt::Interrupt(payload)), } let children = &mut WalkChildren { @@ -1004,14 +1005,11 @@ fn visit_raw( return Err(with_value_context(halt, value)); } - let exit = match visitor.exit(&visit_value, def_region_kind) { - Ok(flow) => flow, - Err(error) => return Err(with_value_context(error.into(), value)), - }; - match exit { - WalkResult::Interrupt => Err(NativeHalt::Interrupt(Any::new())), - WalkResult::InterruptWith(payload) => Err(NativeHalt::Interrupt(payload)), - WalkResult::Advance | WalkResult::Skip => Ok(()), + match visitor.exit(&visit_value, def_region_kind) { + Ok(WalkResult::Interrupt) => Err(NativeHalt::Interrupt(Any::new())), + Ok(WalkResult::InterruptWith(payload)) => Err(NativeHalt::Interrupt(payload)), + Ok(WalkResult::Advance | WalkResult::Skip) => Ok(()), + Err(error) => Err(with_value_context(error.into(), value)), } } @@ -1284,33 +1282,44 @@ unsafe fn visit_reflected_field( .map_err(|halt| with_error_context(halt, &format!("field `{}`", field.name.as_str()))) } +// Runs once per visited value: keep the no-hook fast path small enough to +// actually inline (one cached-column load and a tag compare) and the error +// formatting out of line — with the cold body inside, the `#[inline]` hint +// was declined and the call cost ~20% of the container fast path. #[inline] fn reject_foreign_structural_visit(type_index: i32) -> Result<()> { let Some(attr) = structural_visit_column().and_then(|column| column.get(type_index)) else { return Ok(()); }; - match attr.type_index { - x if x == TVMFFITypeIndex::kTVMFFINone as i32 => Ok(()), - x if x == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 - || x == TVMFFITypeIndex::kTVMFFIFunction as i32 => - { - let value_type = if type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { - format!("type index {type_index}") - } else { - format!("type `{}`", type_key_of(type_index)) - }; - Err(runtime_error(&format!( - "native visitor: {value_type} registers foreign `{STRUCTURAL_VISIT_ATTR}`; \ - visit its children explicitly from a `StructuralVisitor` \ - (`structural_visit`), or skip it with a pre-order `WalkResult::Skip` \ - handler" - ))) - } - _ => Err(Error::new( + if attr.type_index == TVMFFITypeIndex::kTVMFFINone as i32 { + return Ok(()); + } + reject_foreign_structural_visit_cold(type_index, attr.type_index) +} + +#[cold] +#[inline(never)] +fn reject_foreign_structural_visit_cold(type_index: i32, attr_type_index: i32) -> Result<()> { + if attr_type_index == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 + || attr_type_index == TVMFFITypeIndex::kTVMFFIFunction as i32 + { + let value_type = if type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 { + format!("type index {type_index}") + } else { + format!("type `{}`", type_key_of(type_index)) + }; + Err(runtime_error(&format!( + "native visitor: {value_type} registers foreign `{STRUCTURAL_VISIT_ATTR}`; \ + visit its children explicitly from a `StructuralVisitor` \ + (`structural_visit`), or skip it with a pre-order `WalkResult::Skip` \ + handler" + ))) + } else { + Err(Error::new( TYPE_ERROR, &format!("{STRUCTURAL_VISIT_ATTR} must be an opaque function pointer or ffi.Function"), "", - )), + )) } } From 2bbd9f7e0afad50faf8c3b4e86fb04f233d83954 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 18:36:47 -0400 Subject: [PATCH 17/20] [TEST][RUST] Prune redundant structural visit tests (#693) Remove test_dispatch.rs (its downstream-path and cfg checks duplicate coverage in test_structural_visit.rs; the mixed-arity handler case moves into GenericDispatchProbe as a trailing DefRegionKind argument) and drop two closure-walk tests whose behavior is already asserted by the sequence fallback, interrupt-payload, and error-context tests. Co-Authored-By: Claude Fable 5 --- rust/tvm-ffi/tests/test_dispatch.rs | 128 -------------------- rust/tvm-ffi/tests/test_structural_visit.rs | 61 +--------- 2 files changed, 3 insertions(+), 186 deletions(-) delete mode 100644 rust/tvm-ffi/tests/test_dispatch.rs diff --git a/rust/tvm-ffi/tests/test_dispatch.rs b/rust/tvm-ffi/tests/test_dispatch.rs deleted file mode 100644 index 49ff99e80..000000000 --- a/rust/tvm-ffi/tests/test_dispatch.rs +++ /dev/null @@ -1,128 +0,0 @@ -/* - * 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. - */ - -//! This integration test compiles as a downstream crate, checking every public -//! path emitted by the dispatch macro outside `tvm_ffi` itself. - -use tvm_ffi::{ - dispatch, structural_walk, Array, DefRegionKind, Object, VisitDispatch, WalkOrder, WalkResult, -}; - -#[derive(Default)] -struct ExternalCounter { - objects: usize, -} - -#[dispatch(visit)] -impl ExternalCounter { - #[cfg(any(unix, windows))] - #[cfg_attr(all(), inline)] - fn visit_object(&mut self, _value: &Object) -> WalkResult { - self.objects += 1; - WalkResult::Advance - } -} - -fn assert_visit_dispatch() {} - -const _: fn() = assert_visit_dispatch::; - -#[derive(Default)] -struct CfgAttrCounter {} - -#[dispatch(visit)] -impl CfgAttrCounter { - #[cfg(any())] - fn visit_disabled_catch_all(&mut self, _value: &tvm_ffi::VisitValue) -> WalkResult { - WalkResult::Advance - } - - #[cfg_attr(all(), cfg(any()))] - fn visit_disabled(&mut self, _value: &Object) -> WalkResult { - WalkResult::Advance - } - - fn visit_object(&mut self, _value: &Object) -> WalkResult { - WalkResult::Advance - } -} - -const _: fn() = assert_visit_dispatch::; - -struct DisabledCounter; -const _: usize = std::mem::size_of::(); - -#[dispatch(visit)] -#[cfg(any())] -impl DisabledCounter { - fn visit_object(&mut self, _value: &Object) -> WalkResult { - WalkResult::Advance - } -} - -struct CfgAttrDisabledCounter; -const _: usize = std::mem::size_of::(); - -#[dispatch(visit)] -#[cfg_attr(all(), cfg(any()))] -impl CfgAttrDisabledCounter { - fn visit_object(&mut self, _value: &Object) -> WalkResult { - WalkResult::Advance - } -} - -#[derive(Default)] -struct MixedArityCounter { - kinds: Vec, - objects: usize, -} - -#[dispatch(visit)] -impl MixedArityCounter { - fn visit_int(&mut self, _value: i64, kind: DefRegionKind) -> WalkResult { - self.kinds.push(kind); - WalkResult::Advance - } - - fn visit_object(&mut self, _value: &Object) -> WalkResult { - self.objects += 1; - WalkResult::Advance - } -} - -#[test] -fn handlers_may_mix_def_region_arity() { - let root = Array::new(vec![1i64, 2]); - let mut visitor = MixedArityCounter::default(); - assert!(structural_walk(&root, &mut visitor, WalkOrder::PreOrder) - .unwrap() - .is_none()); - assert_eq!(visitor.objects, 1); - assert_eq!(visitor.kinds, vec![DefRegionKind::None; 2]); -} - -#[test] -fn generated_dispatch_uses_public_downstream_paths() { - let root = Array::new(vec![1i64, 2]); - let mut visitor = ExternalCounter::default(); - assert!(structural_walk(&root, &mut visitor, WalkOrder::PreOrder) - .unwrap() - .is_none()); - assert_eq!(visitor.objects, 1); -} diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index d13ffa547..933a8e3e5 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -373,7 +373,9 @@ impl GenericDispatchProbe { WalkResult::Advance } - fn visit_object(&mut self, _value: &tvm_ffi::Object) -> WalkResult { + // Trailing DefRegionKind: handlers may mix arities within one impl. + fn visit_object(&mut self, _value: &tvm_ffi::Object, kind: DefRegionKind) -> WalkResult { + assert_eq!(kind, DefRegionKind::None); self.objects += 1; WalkResult::Advance } @@ -569,26 +571,6 @@ fn visitor_interrupt_propagates_through_default_children() { assert_eq!(i64::try_from(interrupt.value).unwrap(), 7); } -#[test] -fn closure_walk_observes_values() { - // C++: StructuralWalk(root, [&](AnyView value) { ... }) - let root = Array::new(vec![1i64, 2, 3]); - let mut integers = 0; - assert!(structural_walk( - &root, - |value: &VisitValue| { - if value.cast::().is_some() { - integers += 1; - } - WalkResult::Advance - }, - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(integers, 3); -} - #[test] fn closure_walk_receives_def_region_kind() { // C++: StructuralWalk(root, @@ -610,43 +592,6 @@ fn closure_walk_receives_def_region_kind() { assert_eq!(kinds, vec![DefRegionKind::None; 2]); } -#[test] -fn closure_walk_interrupts_and_propagates_errors() { - let root = Array::new(vec![1i64, 2, 3]); - let outcome = structural_walk( - &root, - |value: &VisitValue| -> Result { - if value.cast::() == Some(2) { - return Ok(WalkResult::interrupt_with(2i64)); - } - Ok(WalkResult::Advance) - }, - WalkOrder::PreOrder, - ) - .unwrap(); - let Some(interrupt) = outcome else { - panic!("closure walk unexpectedly completed"); - }; - assert_eq!(i64::try_from(interrupt.value).unwrap(), 2); - - let error = match structural_walk( - &root, - |value: &VisitValue| -> Result { - if value.cast::().is_some() { - Err(runtime_error("closure failed")) - } else { - Ok(WalkResult::Advance) - } - }, - WalkOrder::PreOrder, - ) { - Err(error) => error, - Ok(_) => panic!("closure walk unexpectedly succeeded"), - }; - assert_eq!(error.message(), "closure failed"); - assert!(error.backtrace().contains("object `ffi.Array`")); -} - #[test] fn closure_walk_supports_post_order_and_skip() { let root = Array::new(vec![1i64, 2]); From 839a8e24f33c2bc1e59bc40f61c93a0605788a3d Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 18:44:44 -0400 Subject: [PATCH 18/20] format. Signed-off-by: yuchuan --- rust/tvm-ffi/tests/test_structural_visit.rs | 324 +++++++++++++++++- rust/tvm-ffi/tests/test_walk_chain.rs | 349 -------------------- 2 files changed, 323 insertions(+), 350 deletions(-) delete mode 100644 rust/tvm-ffi/tests/test_walk_chain.rs diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index 933a8e3e5..b59dade60 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -20,7 +20,7 @@ use tvm_ffi::tvm_ffi_sys::{TVMFFIByteArray, TVMFFITypeIndex, TVMFFITypeRegisterAttr}; use tvm_ffi::{ dispatch, structural_visit, structural_walk, Any, AnyView, Array, DefRegionKind, Error, - Function, Map, Result, Shape, String as FfiString, StructuralVisitor, VisitInterrupt, + Function, Map, Object, Result, Shape, String as FfiString, StructuralVisitor, VisitInterrupt, VisitValue, WalkOrder, WalkResult, RUNTIME_ERROR, }; @@ -625,3 +625,325 @@ fn closure_walk_supports_post_order_and_skip() { .is_none()); assert_eq!(visited, 1); } + +// --------------------------------------------------------------------------- +// Tuple walkers: structural_walk(root, (link1, link2, ...), order) — links +// are tried in order and the first whose argument type matches the value +// runs, the Rust analog of the variadic C++ StructuralWalk callback chain. +// --------------------------------------------------------------------------- +#[test] +fn chain_dispatches_first_matching_link() { + // C++: StructuralWalk(root, [&](int64_t v) {...}, + // [&](const ObjectRef& o) {...}, [&](AnyView v) {...}) + let root = Array::new(vec![1i64, 2, 3]); + let mut integers = Vec::new(); + let mut objects = 0; + let mut others = 0; + assert!(structural_walk( + &root, + ( + |value: i64| { + integers.push(value); + WalkResult::Advance + }, + |_object: &Object| { + objects += 1; + WalkResult::Advance + }, + |_value: &VisitValue| { + others += 1; + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(integers, vec![1, 2, 3]); + assert_eq!(objects, 1); // the array itself; integers matched earlier + assert_eq!(others, 0); // every value matched an earlier link +} + +#[test] +fn chain_accepts_owned_object_ref_links() { + let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64, 3])]); + let mut lengths = Vec::new(); + assert!(structural_walk( + &root, + ( + |array: Array| { + lengths.push(array.len()); + WalkResult::Advance + }, + |_value: i64| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + // The outer Array> fails the strict element check and falls + // through the chain; only the inner arrays match the typed link. + assert_eq!(lengths, vec![1, 2]); +} + +#[test] +fn chain_links_may_mix_def_region_arity() { + // Like #[dispatch(visit)] handlers, each link independently opts into + // the trailing DefRegionKind argument. + let root = Array::new(vec![1i64, 2]); + let mut kinds = Vec::new(); + let mut objects = 0; + assert!(structural_walk( + &root, + ( + |_value: i64, kind: DefRegionKind| { + kinds.push(kind); + WalkResult::Advance + }, + |_value: &VisitValue, kind: DefRegionKind| { + assert_eq!(kind, DefRegionKind::None); + objects += 1; + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(kinds, vec![DefRegionKind::None; 2]); + assert_eq!(objects, 1); +} + +#[test] +fn chain_links_can_skip_children() { + let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); + let mut arrays = 0; + let mut integers = 0; + assert!(structural_walk( + &root, + ( + |_array: Array| { + arrays += 1; + WalkResult::Skip + }, + |_value: i64| { + integers += 1; + WalkResult::Advance + }, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(arrays, 2); + assert_eq!(integers, 0); // both inner arrays were skipped +} + +#[test] +fn chain_links_can_interrupt_with_payload() { + let root = Array::new(vec![1i64, 2, 3]); + let mut seen = 0; + let outcome = structural_walk( + &root, + ( + |value: i64| { + seen += 1; + if value == 2 { + return WalkResult::interrupt_with(value * 10); + } + WalkResult::Advance + }, + |_value: &VisitValue| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) + .unwrap(); + let Some(interrupt) = outcome else { + panic!("walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(interrupt.value).unwrap(), 20); + assert_eq!(seen, 2); +} + +#[test] +fn chain_link_errors_include_native_visit_path() { + let root = Array::new(vec![1i64]); + let error = match structural_walk( + &root, + ( + |_value: i64| -> Result { Err(runtime_error("link failed")) }, + |_value: &VisitValue| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) { + Err(error) => error, + Ok(_) => panic!("link unexpectedly succeeded"), + }; + assert_eq!(error.message(), "link failed"); + assert!(error.backtrace().contains("sequence item [0]")); + assert!(error.backtrace().contains("object `ffi.Array`")); +} + +#[test] +fn chain_supports_post_order() { + // Rust borrow rules apply per link: state shared across links goes + // through a RefCell (or a single #[dispatch(visit)] visitor). + let root = Array::new(vec![1i64, 2]); + let events = std::cell::RefCell::new(Vec::new()); + assert!(structural_walk( + &root, + ( + |value: i64| { + events.borrow_mut().push(format!("int:{value}")); + WalkResult::Advance + }, + |_object: &Object| { + events.borrow_mut().push("array".to_string()); + WalkResult::Advance + }, + ), + WalkOrder::PostOrder, + ) + .unwrap() + .is_none()); + assert_eq!(events.into_inner(), vec!["int:1", "int:2", "array"]); +} + +#[test] +fn single_link_tuple_walks() { + let root = Array::new(vec![1i64, 2, 3]); + let mut total = 0; + assert!(structural_walk( + &root, + (|value: i64| { + total += value; + WalkResult::Advance + },), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(total, 6); +} + +#[derive(Default)] +struct ObjectCounter { + objects: usize, +} + +#[dispatch(visit)] +impl ObjectCounter { + fn visit_object(&mut self, _value: &Object) -> WalkResult { + self.objects += 1; + WalkResult::Advance + } +} + +#[test] +fn chain_splices_dispatch_visitors_between_closures() { + // A `&mut` typed visitor participates in the chain like any other link, + // keeping its own no-match fall-through semantics. + let root = Array::new(vec![1i64, 2]); + let mut counter = ObjectCounter::default(); + let mut integers = 0; + assert!(structural_walk( + &root, + (&mut counter, |_value: i64| { + integers += 1; + WalkResult::Advance + },), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(counter.objects, 1); + assert_eq!(integers, 2); +} + +#[test] +fn chain_supports_full_arity() { + let root = Array::new(vec![1i64]); + let mut int_hits = 0; + let mut object_hits = 0; + assert!(structural_walk( + &root, + ( + |_value: f64| WalkResult::Advance, + |_value: bool| WalkResult::Advance, + |_value: tvm_ffi::String| WalkResult::Advance, + |_value: Array| WalkResult::Advance, + |_value: i64| { + int_hits += 1; + WalkResult::Advance + }, + |_value: &Object, _kind: DefRegionKind| { + object_hits += 1; + WalkResult::Advance + }, + |_value: &VisitValue, _kind: DefRegionKind| WalkResult::Advance, + |_value: &VisitValue| WalkResult::Advance, + ), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(int_hits, 1); + assert_eq!(object_hits, 1); +} + +#[test] +fn bare_typed_lambda_walks_without_tuple() { + // A lone typed handler needs no tuple: unmatched values (the array + // itself) advance normally. + let root = Array::new(vec![1i64, 2, 3]); + let mut total = 0; + assert!(structural_walk( + &root, + |value: i64| { + total += value; + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(total, 6); +} + +#[test] +fn bare_node_lambda_takes_def_region_kind() { + let root = Array::new(vec![1i64, 2]); + let mut objects = 0; + assert!(structural_walk( + &root, + |_object: &Object, kind: DefRegionKind| { + assert_eq!(kind, DefRegionKind::None); + objects += 1; + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!(objects, 1); +} + +#[test] +fn bare_owned_object_lambda_interrupts() { + let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); + let outcome = structural_walk( + &root, + |array: Array| { + if array.len() == 1 { + return WalkResult::interrupt_with(array.len() as i64); + } + WalkResult::Advance + }, + WalkOrder::PreOrder, + ) + .unwrap(); + let Some(interrupt) = outcome else { + panic!("walk unexpectedly completed"); + }; + assert_eq!(i64::try_from(interrupt.value).unwrap(), 1); +} diff --git a/rust/tvm-ffi/tests/test_walk_chain.rs b/rust/tvm-ffi/tests/test_walk_chain.rs deleted file mode 100644 index 4a5dab9f6..000000000 --- a/rust/tvm-ffi/tests/test_walk_chain.rs +++ /dev/null @@ -1,349 +0,0 @@ -/* - * 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. - */ - -//! Tuple walkers: `structural_walk(root, (link1, link2, ...), order)`, the -//! Rust analog of the C++ variadic `StructuralWalk(root, callbacks...)` -//! chain. Links are tried in order and the first whose argument type matches -//! the value runs. - -use tvm_ffi::{ - dispatch, structural_walk, Array, DefRegionKind, Error, Object, Result, VisitValue, WalkOrder, - WalkResult, RUNTIME_ERROR, -}; - -fn runtime_error(message: &str) -> Error { - Error::new(RUNTIME_ERROR, message, "") -} - -#[test] -fn chain_dispatches_first_matching_link() { - // C++: StructuralWalk(root, [&](int64_t v) {...}, - // [&](const ObjectRef& o) {...}, [&](AnyView v) {...}) - let root = Array::new(vec![1i64, 2, 3]); - let mut integers = Vec::new(); - let mut objects = 0; - let mut others = 0; - assert!(structural_walk( - &root, - ( - |value: i64| { - integers.push(value); - WalkResult::Advance - }, - |_object: &Object| { - objects += 1; - WalkResult::Advance - }, - |_value: &VisitValue| { - others += 1; - WalkResult::Advance - }, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(integers, vec![1, 2, 3]); - assert_eq!(objects, 1); // the array itself; integers matched earlier - assert_eq!(others, 0); // every value matched an earlier link -} - -#[test] -fn chain_accepts_owned_object_ref_links() { - let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64, 3])]); - let mut lengths = Vec::new(); - assert!(structural_walk( - &root, - ( - |array: Array| { - lengths.push(array.len()); - WalkResult::Advance - }, - |_value: i64| WalkResult::Advance, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - // The outer Array> fails the strict element check and falls - // through the chain; only the inner arrays match the typed link. - assert_eq!(lengths, vec![1, 2]); -} - -#[test] -fn chain_links_may_mix_def_region_arity() { - // Like #[dispatch(visit)] handlers, each link independently opts into - // the trailing DefRegionKind argument. - let root = Array::new(vec![1i64, 2]); - let mut kinds = Vec::new(); - let mut objects = 0; - assert!(structural_walk( - &root, - ( - |_value: i64, kind: DefRegionKind| { - kinds.push(kind); - WalkResult::Advance - }, - |_value: &VisitValue, kind: DefRegionKind| { - assert_eq!(kind, DefRegionKind::None); - objects += 1; - WalkResult::Advance - }, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(kinds, vec![DefRegionKind::None; 2]); - assert_eq!(objects, 1); -} - -#[test] -fn chain_links_can_skip_children() { - let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); - let mut arrays = 0; - let mut integers = 0; - assert!(structural_walk( - &root, - ( - |_array: Array| { - arrays += 1; - WalkResult::Skip - }, - |_value: i64| { - integers += 1; - WalkResult::Advance - }, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(arrays, 2); - assert_eq!(integers, 0); // both inner arrays were skipped -} - -#[test] -fn chain_links_can_interrupt_with_payload() { - let root = Array::new(vec![1i64, 2, 3]); - let mut seen = 0; - let outcome = structural_walk( - &root, - ( - |value: i64| { - seen += 1; - if value == 2 { - return WalkResult::interrupt_with(value * 10); - } - WalkResult::Advance - }, - |_value: &VisitValue| WalkResult::Advance, - ), - WalkOrder::PreOrder, - ) - .unwrap(); - let Some(interrupt) = outcome else { - panic!("walk unexpectedly completed"); - }; - assert_eq!(i64::try_from(interrupt.value).unwrap(), 20); - assert_eq!(seen, 2); -} - -#[test] -fn chain_link_errors_include_native_visit_path() { - let root = Array::new(vec![1i64]); - let error = match structural_walk( - &root, - ( - |_value: i64| -> Result { Err(runtime_error("link failed")) }, - |_value: &VisitValue| WalkResult::Advance, - ), - WalkOrder::PreOrder, - ) { - Err(error) => error, - Ok(_) => panic!("link unexpectedly succeeded"), - }; - assert_eq!(error.message(), "link failed"); - assert!(error.backtrace().contains("sequence item [0]")); - assert!(error.backtrace().contains("object `ffi.Array`")); -} - -#[test] -fn chain_supports_post_order() { - // Rust borrow rules apply per link: state shared across links goes - // through a RefCell (or a single #[dispatch(visit)] visitor). - let root = Array::new(vec![1i64, 2]); - let events = std::cell::RefCell::new(Vec::new()); - assert!(structural_walk( - &root, - ( - |value: i64| { - events.borrow_mut().push(format!("int:{value}")); - WalkResult::Advance - }, - |_object: &Object| { - events.borrow_mut().push("array".to_string()); - WalkResult::Advance - }, - ), - WalkOrder::PostOrder, - ) - .unwrap() - .is_none()); - assert_eq!(events.into_inner(), vec!["int:1", "int:2", "array"]); -} - -#[test] -fn single_link_tuple_walks() { - let root = Array::new(vec![1i64, 2, 3]); - let mut total = 0; - assert!(structural_walk( - &root, - (|value: i64| { - total += value; - WalkResult::Advance - },), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(total, 6); -} - -#[derive(Default)] -struct ObjectCounter { - objects: usize, -} - -#[dispatch(visit)] -impl ObjectCounter { - fn visit_object(&mut self, _value: &Object) -> WalkResult { - self.objects += 1; - WalkResult::Advance - } -} - -#[test] -fn chain_splices_dispatch_visitors_between_closures() { - // A `&mut` typed visitor participates in the chain like any other link, - // keeping its own no-match fall-through semantics. - let root = Array::new(vec![1i64, 2]); - let mut counter = ObjectCounter::default(); - let mut integers = 0; - assert!(structural_walk( - &root, - (&mut counter, |_value: i64| { - integers += 1; - WalkResult::Advance - },), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(counter.objects, 1); - assert_eq!(integers, 2); -} - -#[test] -fn chain_supports_full_arity() { - let root = Array::new(vec![1i64]); - let mut int_hits = 0; - let mut object_hits = 0; - assert!(structural_walk( - &root, - ( - |_value: f64| WalkResult::Advance, - |_value: bool| WalkResult::Advance, - |_value: tvm_ffi::String| WalkResult::Advance, - |_value: Array| WalkResult::Advance, - |_value: i64| { - int_hits += 1; - WalkResult::Advance - }, - |_value: &Object, _kind: DefRegionKind| { - object_hits += 1; - WalkResult::Advance - }, - |_value: &VisitValue, _kind: DefRegionKind| WalkResult::Advance, - |_value: &VisitValue| WalkResult::Advance, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(int_hits, 1); - assert_eq!(object_hits, 1); -} - -#[test] -fn bare_typed_lambda_walks_without_tuple() { - // A lone typed handler needs no tuple: unmatched values (the array - // itself) advance normally. - let root = Array::new(vec![1i64, 2, 3]); - let mut total = 0; - assert!(structural_walk( - &root, - |value: i64| { - total += value; - WalkResult::Advance - }, - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(total, 6); -} - -#[test] -fn bare_node_lambda_takes_def_region_kind() { - let root = Array::new(vec![1i64, 2]); - let mut objects = 0; - assert!(structural_walk( - &root, - |_object: &Object, kind: DefRegionKind| { - assert_eq!(kind, DefRegionKind::None); - objects += 1; - WalkResult::Advance - }, - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(objects, 1); -} - -#[test] -fn bare_owned_object_lambda_interrupts() { - let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); - let outcome = structural_walk( - &root, - |array: Array| { - if array.len() == 1 { - return WalkResult::interrupt_with(array.len() as i64); - } - WalkResult::Advance - }, - WalkOrder::PreOrder, - ) - .unwrap(); - let Some(interrupt) = outcome else { - panic!("walk unexpectedly completed"); - }; - assert_eq!(i64::try_from(interrupt.value).unwrap(), 1); -} From 7766d427e394cec9ff37951bc15f01407dacdcc4 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 18:48:52 -0400 Subject: [PATCH 19/20] format. Signed-off-by: yuchuan --- rust/tvm-ffi/tests/test_structural_visit.rs | 146 +++++--------------- 1 file changed, 35 insertions(+), 111 deletions(-) diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs b/rust/tvm-ffi/tests/test_structural_visit.rs index b59dade60..723de8079 100644 --- a/rust/tvm-ffi/tests/test_structural_visit.rs +++ b/rust/tvm-ffi/tests/test_structural_visit.rs @@ -631,39 +631,6 @@ fn closure_walk_supports_post_order_and_skip() { // are tried in order and the first whose argument type matches the value // runs, the Rust analog of the variadic C++ StructuralWalk callback chain. // --------------------------------------------------------------------------- -#[test] -fn chain_dispatches_first_matching_link() { - // C++: StructuralWalk(root, [&](int64_t v) {...}, - // [&](const ObjectRef& o) {...}, [&](AnyView v) {...}) - let root = Array::new(vec![1i64, 2, 3]); - let mut integers = Vec::new(); - let mut objects = 0; - let mut others = 0; - assert!(structural_walk( - &root, - ( - |value: i64| { - integers.push(value); - WalkResult::Advance - }, - |_object: &Object| { - objects += 1; - WalkResult::Advance - }, - |_value: &VisitValue| { - others += 1; - WalkResult::Advance - }, - ), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(integers, vec![1, 2, 3]); - assert_eq!(objects, 1); // the array itself; integers matched earlier - assert_eq!(others, 0); // every value matched an earlier link -} - #[test] fn chain_accepts_owned_object_ref_links() { let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64, 3])]); @@ -739,32 +706,6 @@ fn chain_links_can_skip_children() { assert_eq!(integers, 0); // both inner arrays were skipped } -#[test] -fn chain_links_can_interrupt_with_payload() { - let root = Array::new(vec![1i64, 2, 3]); - let mut seen = 0; - let outcome = structural_walk( - &root, - ( - |value: i64| { - seen += 1; - if value == 2 { - return WalkResult::interrupt_with(value * 10); - } - WalkResult::Advance - }, - |_value: &VisitValue| WalkResult::Advance, - ), - WalkOrder::PreOrder, - ) - .unwrap(); - let Some(interrupt) = outcome else { - panic!("walk unexpectedly completed"); - }; - assert_eq!(i64::try_from(interrupt.value).unwrap(), 20); - assert_eq!(seen, 2); -} - #[test] fn chain_link_errors_include_native_visit_path() { let root = Array::new(vec![1i64]); @@ -809,23 +750,6 @@ fn chain_supports_post_order() { assert_eq!(events.into_inner(), vec!["int:1", "int:2", "array"]); } -#[test] -fn single_link_tuple_walks() { - let root = Array::new(vec![1i64, 2, 3]); - let mut total = 0; - assert!(structural_walk( - &root, - (|value: i64| { - total += value; - WalkResult::Advance - },), - WalkOrder::PreOrder, - ) - .unwrap() - .is_none()); - assert_eq!(total, 6); -} - #[derive(Default)] struct ObjectCounter { objects: usize, @@ -862,9 +786,13 @@ fn chain_splices_dispatch_visitors_between_closures() { #[test] fn chain_supports_full_arity() { - let root = Array::new(vec![1i64]); - let mut int_hits = 0; - let mut object_hits = 0; + // Doubles as the first-match ordering probe: earlier misses fall + // through, the first matching link claims the value, later links + // never run. + let root = Array::new(vec![1i64, 2, 3]); + let mut integers = Vec::new(); + let mut objects = 0; + let mut others = 0; assert!(structural_walk( &root, ( @@ -872,42 +800,58 @@ fn chain_supports_full_arity() { |_value: bool| WalkResult::Advance, |_value: tvm_ffi::String| WalkResult::Advance, |_value: Array| WalkResult::Advance, - |_value: i64| { - int_hits += 1; + |value: i64| { + integers.push(value); WalkResult::Advance }, - |_value: &Object, _kind: DefRegionKind| { - object_hits += 1; + |_object: &Object, _kind: DefRegionKind| { + objects += 1; + WalkResult::Advance + }, + |_value: &VisitValue, _kind: DefRegionKind| { + others += 1; WalkResult::Advance }, - |_value: &VisitValue, _kind: DefRegionKind| WalkResult::Advance, |_value: &VisitValue| WalkResult::Advance, ), WalkOrder::PreOrder, ) .unwrap() .is_none()); - assert_eq!(int_hits, 1); - assert_eq!(object_hits, 1); + assert_eq!(integers, vec![1, 2, 3]); + assert_eq!(objects, 1); // the array itself; integers matched earlier + assert_eq!(others, 0); // every value matched an earlier link } #[test] -fn bare_typed_lambda_walks_without_tuple() { +fn typed_lambda_walks_bare_and_as_single_link_tuple() { // A lone typed handler needs no tuple: unmatched values (the array - // itself) advance normally. + // itself) advance normally. The 1-tuple spelling routes through the + // chain impls instead and must agree. let root = Array::new(vec![1i64, 2, 3]); - let mut total = 0; + let mut bare = 0; assert!(structural_walk( &root, |value: i64| { - total += value; + bare += value; WalkResult::Advance }, WalkOrder::PreOrder, ) .unwrap() .is_none()); - assert_eq!(total, 6); + let mut tupled = 0; + assert!(structural_walk( + &root, + (|value: i64| { + tupled += value; + WalkResult::Advance + },), + WalkOrder::PreOrder, + ) + .unwrap() + .is_none()); + assert_eq!((bare, tupled), (6, 6)); } #[test] @@ -927,23 +871,3 @@ fn bare_node_lambda_takes_def_region_kind() { .is_none()); assert_eq!(objects, 1); } - -#[test] -fn bare_owned_object_lambda_interrupts() { - let root = Array::new(vec![Array::new(vec![1i64]), Array::new(vec![2i64])]); - let outcome = structural_walk( - &root, - |array: Array| { - if array.len() == 1 { - return WalkResult::interrupt_with(array.len() as i64); - } - WalkResult::Advance - }, - WalkOrder::PreOrder, - ) - .unwrap(); - let Some(interrupt) = outcome else { - panic!("walk unexpectedly completed"); - }; - assert_eq!(i64::try_from(interrupt.value).unwrap(), 1); -} From 72d9e4c96fb2da36d172900ff449030ba8afeff3 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 2 Aug 2026 19:19:44 -0400 Subject: [PATCH 20/20] refac. Signed-off-by: yuchuan --- rust/tvm-ffi/src/extra/dispatch.rs | 114 ++++++++++++++++++ rust/tvm-ffi/src/extra/mod.rs | 1 + rust/tvm-ffi/src/extra/structural_visit.rs | 92 +------------- ...s => test_structural_visitor_alignment.rs} | 0 4 files changed, 121 insertions(+), 86 deletions(-) create mode 100644 rust/tvm-ffi/src/extra/dispatch.rs rename rust/tvm-ffi/tests/{test_visitor_alignment.rs => test_structural_visitor_alignment.rs} (100%) diff --git a/rust/tvm-ffi/src/extra/dispatch.rs b/rust/tvm-ffi/src/extra/dispatch.rs new file mode 100644 index 000000000..2b3a49f0b --- /dev/null +++ b/rust/tvm-ffi/src/extra/dispatch.rs @@ -0,0 +1,114 @@ +/* + * 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. + */ + +//! Typed visitor dispatch for [`super::structural_visit::structural_walk`]: +//! the [`VisitDispatch`] trait targeted by `#[dispatch(visit)]`, and the +//! walker adapter that runs such a visitor at the phase selected by the walk +//! order. The traversal engine, closure walkers, and tuple chains live in +//! [`super::structural_visit`], which re-exports these items to keep its +//! public paths stable. + +use crate::error::Result; + +use super::structural_visit::{ + DefRegionKind, IntoWalker, NativeVisit, VisitResult, VisitValue, WalkOrder, WalkResult, +}; + +/// Typed dispatch implemented by a walk-layer observer. +/// +/// [`crate::dispatch`] tests the implementation's `visit_*` methods in source +/// order. Borrowed node arguments use refcount-free subtype checks, owned +/// FFI-compatible arguments use exact value casts, and `&VisitValue` is a +/// catch-all. `None` reports that no handler matched: a standalone walk then +/// advances normally, while a tuple chain hands the value to the next link — +/// so a spliced visitor that "handled" a value must not return `None`. +/// +/// This is the observer layer, mirroring C++ `StructuralWalk` callbacks: the +/// walker owns recursion, and a handler steers it only through the returned +/// [`WalkResult`]. A traversal that must visit children itself — selected +/// children, custom orders, explicit definition-region overrides — belongs in +/// a [`super::structural_visit::StructuralVisitor`] instead. +/// +/// The definition-region state active at the dispatched value arrives as the +/// `def_region_kind` argument. A `#[dispatch(visit)]` handler opts into it by +/// declaring a trailing `DefRegionKind` parameter — the analog of a C++ +/// `StructuralWalk` callback accepting `(value, def_region_kind)` instead of +/// `(value)`. +pub trait VisitDispatch: Sized { + fn dispatch_visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option; +} + +impl VisitDispatch for &mut V { + #[inline] + fn dispatch_visit( + &mut self, + value: &VisitValue, + def_region_kind: DefRegionKind, + ) -> Option { + (**self).dispatch_visit(value, def_region_kind) + } +} + +#[doc(hidden)] +pub enum ByDispatch {} + +impl<'a, V: VisitDispatch> IntoWalker for &'a mut V { + type Walker = DispatchVisitor<&'a mut V>; + fn into_walker(self, order: WalkOrder) -> Self::Walker { + DispatchVisitor { + visitor: self, + order, + } + } +} + +/// Owns its walker so a closure's state stays inline and a `&mut` visitor +/// keeps a single level of indirection. Public only as an +/// [`IntoWalker::Walker`] projection. +#[doc(hidden)] +pub struct DispatchVisitor { + visitor: V, + order: WalkOrder, +} + +impl NativeVisit for DispatchVisitor { + fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { + match self.order { + WalkOrder::PreOrder => self + .visitor + .dispatch_visit(value, def_region_kind) + .unwrap_or(Ok(WalkResult::Advance)), + WalkOrder::PostOrder => Ok(WalkResult::Advance), + } + } + + fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { + match self.order { + WalkOrder::PreOrder => Ok(WalkResult::Advance), + WalkOrder::PostOrder => self + .visitor + .dispatch_visit(value, def_region_kind) + .unwrap_or(Ok(WalkResult::Advance)), + } + } +} diff --git a/rust/tvm-ffi/src/extra/mod.rs b/rust/tvm-ffi/src/extra/mod.rs index 32fed9fa3..0489fe27f 100644 --- a/rust/tvm-ffi/src/extra/mod.rs +++ b/rust/tvm-ffi/src/extra/mod.rs @@ -16,5 +16,6 @@ * specific language governing permissions and limitations * under the License. */ +pub mod dispatch; pub mod module; pub mod structural_visit; diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs b/rust/tvm-ffi/src/extra/structural_visit.rs index e17944351..1d15efd89 100644 --- a/rust/tvm-ffi/src/extra/structural_visit.rs +++ b/rust/tvm-ffi/src/extra/structural_visit.rs @@ -247,44 +247,11 @@ impl From for NativeHalt { type NativeResult = std::result::Result<(), NativeHalt>; -/// Typed dispatch implemented by a walk-layer observer. -/// -/// [`crate::dispatch`] tests the implementation's `visit_*` methods in source -/// order. Borrowed node arguments use refcount-free subtype checks, owned -/// FFI-compatible arguments use exact value casts, and `&VisitValue` is a -/// catch-all. `None` reports that no handler matched: a standalone walk then -/// advances normally, while a tuple chain hands the value to the next link — -/// so a spliced visitor that "handled" a value must not return `None`. -/// -/// This is the observer layer, mirroring C++ `StructuralWalk` callbacks: the -/// walker owns recursion, and a handler steers it only through the returned -/// [`WalkResult`]. A traversal that must visit children itself — selected -/// children, custom orders, explicit definition-region overrides — belongs in -/// a [`StructuralVisitor`] instead. -/// -/// The definition-region state active at the dispatched value arrives as the -/// `def_region_kind` argument. A `#[dispatch(visit)]` handler opts into it by -/// declaring a trailing `DefRegionKind` parameter — the analog of a C++ -/// `StructuralWalk` callback accepting `(value, def_region_kind)` instead of -/// `(value)`. -pub trait VisitDispatch: Sized { - fn dispatch_visit( - &mut self, - value: &VisitValue, - def_region_kind: DefRegionKind, - ) -> Option; -} - -impl VisitDispatch for &mut V { - #[inline] - fn dispatch_visit( - &mut self, - value: &VisitValue, - def_region_kind: DefRegionKind, - ) -> Option { - (**self).dispatch_visit(value, def_region_kind) - } -} +// The typed-dispatch layer (`VisitDispatch`, its walker adapter, and the +// `&mut V` IntoWalker form) lives in `super::dispatch`; re-exported here so +// the module's public paths — which `#[dispatch(visit)]`-generated code +// names — stay stable. +pub use super::dispatch::{ByDispatch, DispatchVisitor, VisitDispatch}; /// Conversion into the walker argument of [`structural_walk`]. /// @@ -322,19 +289,6 @@ pub trait IntoWalker { fn into_walker(self, order: WalkOrder) -> Self::Walker; } -#[doc(hidden)] -pub enum ByDispatch {} - -impl<'a, V: VisitDispatch> IntoWalker for &'a mut V { - type Walker = DispatchVisitor<&'a mut V>; - fn into_walker(self, order: WalkOrder) -> Self::Walker { - DispatchVisitor { - visitor: self, - order, - } - } -} - /// Runs a catch-all closure at the phase selected by `order` — the closure /// analog of `DispatchVisitor`, without the `Option` /// no-handler-matched layer a dispatch chain needs. (Routing closures @@ -896,37 +850,6 @@ pub trait NativeVisit { } } -/// Owns its walker so a closure's state stays inline and a `&mut` visitor -/// keeps a single level of indirection. Public only as an -/// [`IntoWalker::Walker`] projection. -#[doc(hidden)] -pub struct DispatchVisitor { - visitor: V, - order: WalkOrder, -} - -impl NativeVisit for DispatchVisitor { - fn enter(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { - match self.order { - WalkOrder::PreOrder => self - .visitor - .dispatch_visit(value, def_region_kind) - .unwrap_or(Ok(WalkResult::Advance)), - WalkOrder::PostOrder => Ok(WalkResult::Advance), - } - } - - fn exit(&mut self, value: &VisitValue, def_region_kind: DefRegionKind) -> Result { - match self.order { - WalkOrder::PreOrder => Ok(WalkResult::Advance), - WalkOrder::PostOrder => self - .visitor - .dispatch_visit(value, def_region_kind) - .unwrap_or(Ok(WalkResult::Advance)), - } - } -} - /// Per-child action invoked by the shared child-iteration engine. /// /// The engine owns *finding* the children (container contents, reflected @@ -1767,10 +1690,7 @@ mod tests { #[test] fn reflected_field_def_region_reaches_typed_handler() { let mut probe = TypedRegionProbe::default(); - let mut dispatch = DispatchVisitor { - visitor: &mut probe, - order: WalkOrder::PreOrder, - }; + let mut dispatch = (&mut probe).into_walker(WalkOrder::PreOrder); let mut value = Any::from(7i64); let mut field: TVMFFIFieldInfo = unsafe { std::mem::zeroed() }; field.name = unsafe { TVMFFIByteArray::from_str("value") }; diff --git a/rust/tvm-ffi/tests/test_visitor_alignment.rs b/rust/tvm-ffi/tests/test_structural_visitor_alignment.rs similarity index 100% rename from rust/tvm-ffi/tests/test_visitor_alignment.rs rename to rust/tvm-ffi/tests/test_structural_visitor_alignment.rs