This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new b19448e2 [Feat][Rust] Add leaf lookup dispatch for match_any (#685)
b19448e2 is described below
commit b19448e2b9f5e4ec499cbfe21bc6cbefa1081b09
Author: Shushi Hong <[email protected]>
AuthorDate: Mon Jul 27 18:18:05 2026 -0400
[Feat][Rust] Add leaf lookup dispatch for match_any (#685)
This PR adds a leaf-only lookup fast path to Rust match_any!. For at
least two unguarded arms with simple bindings, when every pattern
corresponds to one stable runtime TypeIndex, it lazily maps each
TypeIndex to a source-order ArmId and uses a native Rust match to select
only the relevant conversion. The table stores no closures or arm
bodies, duplicate types preserve first-arm semantics, and misses use the
_ fallback. Single-arm, guarded, non-leaf/category, or content-dependent
matches retain the Phase 1 ordered path. The PR also adds opt-in
final-type metadata and focused integration tests. The two-arm threshold
is based on release benchmarks of standard ObjectRef conversions.
---
rust/tvm-ffi-macros/src/lib.rs | 2 +-
rust/tvm-ffi-macros/src/match_any.rs | 344 ++++++++++++++++++++++++++++---
rust/tvm-ffi-macros/src/object_macros.rs | 35 ++++
rust/tvm-ffi/src/extra/module.rs | 1 +
rust/tvm-ffi/src/lib.rs | 2 +
rust/tvm-ffi/src/match_any_internal.rs | 149 +++++++++++++
rust/tvm-ffi/src/object.rs | 5 +
rust/tvm-ffi/src/type_traits.rs | 17 ++
rust/tvm-ffi/tests/test_match_any.rs | 89 +++++++-
9 files changed, 617 insertions(+), 27 deletions(-)
diff --git a/rust/tvm-ffi-macros/src/lib.rs b/rust/tvm-ffi-macros/src/lib.rs
index aeffad1f..03ecb050 100644
--- a/rust/tvm-ffi-macros/src/lib.rs
+++ b/rust/tvm-ffi-macros/src/lib.rs
@@ -37,7 +37,7 @@ pub fn match_any(input: TokenStream) -> TokenStream {
}
#[proc_macro_error]
-#[proc_macro_derive(Object, attributes(type_key, type_index))]
+#[proc_macro_derive(Object, attributes(type_key, type_index, type_final))]
pub fn derive_object(input: TokenStream) -> TokenStream {
TokenStream::from(object_macros::derive_object(input))
}
diff --git a/rust/tvm-ffi-macros/src/match_any.rs
b/rust/tvm-ffi-macros/src/match_any.rs
index 14839210..51a479bc 100644
--- a/rust/tvm-ffi-macros/src/match_any.rs
+++ b/rust/tvm-ffi-macros/src/match_any.rs
@@ -24,6 +24,11 @@ use syn::{braced, parenthesized, Expr, Pat, Path, Result,
Token};
use crate::utils::get_tvm_ffi_crate;
+// Keep single-arm matches ordered. Release benchmarks of the standard
ObjectRef
+// conversion path show that lookup pays off from two arms when a later arm or
+// the fallback is reached.
+const MIN_LOOKUP_TABLE_ARMS: usize = 2;
+
struct MatchAnyInput {
scrutinee: Expr,
arms: Vec<TypedArm>,
@@ -105,41 +110,306 @@ pub fn expand(input: proc_macro::TokenStream) ->
proc_macro::TokenStream {
fn expand_match_any(input: MatchAnyInput) -> TokenStream {
let tvm_ffi = get_tvm_ffi_crate();
+ let scrutinee = input.scrutinee;
+ let fallback = input.fallback;
+ let arms = input.arms;
+ let can_attempt_leaf_lookup = arms.len() >= MIN_LOOKUP_TABLE_ARMS
+ && arms
+ .iter()
+ .all(|arm| arm.guard.is_none() && is_simple_binding(&arm.binding));
+
+ if can_attempt_leaf_lookup {
+ expand_leaf_lookup_match(&tvm_ffi, &scrutinee, &arms, &fallback)
+ } else {
+ expand_ordered_match(&tvm_ffi, &scrutinee, &arms, &fallback)
+ }
+}
+
+fn expand_ordered_match(
+ tvm_ffi: &TokenStream,
+ scrutinee: &Expr,
+ arms: &[TypedArm],
+ fallback: &Expr,
+) -> TokenStream {
let span = Span::mixed_site();
let source = Ident::new("__tvm_ffi_match_any_source", span);
let converted = Ident::new("__tvm_ffi_match_any_converted", span);
let view = Ident::new("__tvm_ffi_match_any_view", span);
let rejected = Ident::new("__tvm_ffi_match_any_rejected", span);
- let scrutinee = input.scrutinee;
- let fallback = input.fallback;
- let dispatch_fallback = fallback.clone();
- let arms = input.arms;
- let dispatch = arms
- .into_iter()
- .rev()
- .fold(quote!({ #dispatch_fallback }), |next, arm| {
- let matcher = arm.matcher;
- let binding = arm.binding;
- let body = arm.body;
- let matched = if let Some(guard) = arm.guard {
- quote!(::core::result::Result::Ok(#binding) if #guard)
- } else {
- quote!(::core::result::Result::Ok(#binding))
+ let dispatch = expand_ordered_dispatch(arms, fallback, &view, &rejected);
+
+ quote! {
+ {
+ let #source = &(#scrutinee);
+ let #converted: ::core::result::Result<
+ #tvm_ffi::AnyView<'_>,
+ ::core::convert::Infallible,
+ > =
::core::convert::TryInto::<#tvm_ffi::AnyView<'_>>::try_into(#source);
+ let #view = match #converted {
+ ::core::result::Result::Ok(view) => view,
+ ::core::result::Result::Err(error) => match error {},
};
+ if #view.type_index()
+ >= #tvm_ffi::TypeIndex::kTVMFFIStaticObjectBegin as i32
+ {
+ #dispatch
+ } else {
+ #fallback
+ }
+ }
+ }
+}
+
+fn expand_ordered_dispatch(
+ arms: &[TypedArm],
+ fallback: &Expr,
+ view: &Ident,
+ rejected: &Ident,
+) -> TokenStream {
+ expand_ordered_try_into_chain(arms, quote!({ #fallback }), view, rejected,
|_, arm| {
+ let binding = &arm.binding;
+ let body = &arm.body;
+ if let Some(guard) = &arm.guard {
+ quote!(::core::result::Result::Ok(#binding) if #guard => { #body })
+ } else {
+ quote!(::core::result::Result::Ok(#binding) => { #body })
+ }
+ })
+}
+
+fn expand_ordered_try_into_chain<F>(
+ arms: &[TypedArm],
+ fallback: TokenStream,
+ view: &Ident,
+ rejected: &Ident,
+ mut matched_arm: F,
+) -> TokenStream
+where
+ F: FnMut(usize, &TypedArm) -> TokenStream,
+{
+ arms.iter()
+ .enumerate()
+ .rev()
+ .fold(fallback, |next, (arm_id, arm)| {
+ let matcher = &arm.matcher;
+ let matched = matched_arm(arm_id, arm);
quote! {
match ::core::convert::TryInto::<#matcher>::try_into(#view) {
- #matched => { #body },
+ #matched,
#rejected => {
::core::mem::drop(#rejected);
#next
- },
+ }
+ }
+ }
+ })
+}
+
+fn expand_leaf_table_lookup(
+ tvm_ffi: &TokenStream,
+ arms: &[TypedArm],
+ arm_constants: &[Ident],
+ view: &Ident,
+) -> TokenStream {
+ let span = Span::mixed_site();
+ let probe = Ident::new("__tvm_ffi_match_any_probe", span);
+ let pattern_list_id =
Ident::new("__tvm_ffi_match_any_leaf_pattern_list_id", span);
+ let type_indices = Ident::new("__tvm_ffi_match_any_type_indices", span);
+ let static_table = Ident::new("__TVM_FFI_MATCH_ANY_LEAF_TABLE", span);
+ let table = Ident::new("__tvm_ffi_match_any_leaf_table", span);
+ let arm_count = arms.len();
+ let lookup_entries = arm_constants
+ .iter()
+ .enumerate()
+ .map(|(arm_id, arm_constant)| quote!((#type_indices[#arm_id],
#arm_constant)));
+ let pattern_list = arms
+ .iter()
+ .map(|arm| &arm.matcher)
+ .rev()
+ .fold(quote!(()), |tail, matcher| quote!((#matcher, #tail)));
+
+ quote! {
+ {
+ use #tvm_ffi::match_any_internal::LeafPatternMetadata as _;
+
+ let #probe =
+
#tvm_ffi::match_any_internal::LeafPatternProbe::<#pattern_list>::new();
+ match (&#probe).leaf_pattern_list_id() {
+ ::core::option::Option::Some(#pattern_list_id) => {
+ static #static_table: ::std::sync::OnceLock<
+ #tvm_ffi::match_any_internal::LeafLookupTable,
+ > = ::std::sync::OnceLock::new();
+ let #table = #static_table.get_or_init(|| {
+ let mut #type_indices = [0_i32; #arm_count];
+ (&#probe).fill_leaf_type_indices(&mut #type_indices);
+ #tvm_ffi::match_any_internal::LeafLookupTable::build(
+ #pattern_list_id,
+ &[#(#lookup_entries),*],
+ )
+ });
+ #table.lookup(#pattern_list_id, #view.type_index())
+ }
+ ::core::option::Option::None => {
+ ::core::result::Result::Err(())
+ }
+ }
+ }
+ }
+}
+
+fn expand_direct_leaf_selection(
+ arms: &[TypedArm],
+ arm_constants: &[Ident],
+ arm_variants: &[Ident],
+ arm_id: &Ident,
+ view: &Ident,
+ rejected: &Ident,
+ selected_enum: &Ident,
+ selected_value: &Ident,
+) -> TokenStream {
+ let selections = arms.iter().enumerate().map(|(arm_id, arm)| {
+ let matcher = &arm.matcher;
+ let variant = &arm_variants[arm_id];
+ let arm_constant = &arm_constants[arm_id];
+
+ quote! {
+ #arm_constant => {
+ match ::core::convert::TryInto::<#matcher>::try_into(#view) {
+ ::core::result::Result::Ok(#selected_value) => {
+ #selected_enum::#variant(#selected_value)
+ }
+ #rejected => {
+ ::core::mem::drop(#rejected);
+ ::core::panic!(
+ "match_any! leaf lookup selected an incompatible
arm"
+ )
+ }
}
}
- });
+ }
+ });
+
+ quote! {
+ match #arm_id {
+ #(#selections,)*
+ _ => ::core::unreachable!(),
+ }
+ }
+}
+
+fn expand_leaf_body_dispatch(
+ arms: &[TypedArm],
+ arm_variants: &[Ident],
+ selected_enum: &Ident,
+ selected: &Ident,
+ fallback_variant: &Ident,
+ fallback: &Expr,
+) -> TokenStream {
+ let body_arms = arms.iter().enumerate().map(|(arm_id, arm)| {
+ let binding = &arm.binding;
+ let body = &arm.body;
+ let variant = &arm_variants[arm_id];
+
+ quote! {
+ #selected_enum::#variant(#binding) => {
+ #body
+ }
+ }
+ });
+
+ quote! {
+ match #selected {
+ #(#body_arms,)*
+ #selected_enum::#fallback_variant => {
+ #fallback
+ }
+ }
+ }
+}
+
+fn expand_leaf_lookup_match(
+ tvm_ffi: &TokenStream,
+ scrutinee: &Expr,
+ arms: &[TypedArm],
+ fallback: &Expr,
+) -> TokenStream {
+ let span = Span::mixed_site();
+ let source = Ident::new("__tvm_ffi_match_any_source", span);
+ let converted = Ident::new("__tvm_ffi_match_any_converted", span);
+ let view = Ident::new("__tvm_ffi_match_any_view", span);
+ let rejected = Ident::new("__tvm_ffi_match_any_rejected", span);
+ let arm_id = Ident::new("__tvm_ffi_match_any_arm_id", span);
+ let selected = Ident::new("__tvm_ffi_match_any_selected", span);
+ let selected_value = Ident::new("__tvm_ffi_match_any_selected_value",
span);
+ let selected_enum = Ident::new("__TvmFfiMatchAnyArm", span);
+ let fallback_variant = Ident::new("Fallback", span);
+ let arm_count = arms.len();
+ let arm_types = (0..arm_count)
+ .map(|arm_id| Ident::new(&format!("__TvmFfiMatchAnyType{arm_id}"),
span))
+ .collect::<Vec<_>>();
+ let arm_variants = (0..arm_count)
+ .map(|arm_id| Ident::new(&format!("Arm{arm_id}"), span))
+ .collect::<Vec<_>>();
+ let arm_constants = (0..arm_count)
+ .map(|arm_id| Ident::new(&format!("__TVM_FFI_MATCH_ANY_ARM_{arm_id}"),
span))
+ .collect::<Vec<_>>();
+ let arm_constant_definitions =
+ arm_constants
+ .iter()
+ .enumerate()
+ .map(|(arm_id, arm_constant)| {
+ quote! {
+ const #arm_constant: #tvm_ffi::match_any_internal::ArmId =
+ #arm_id as #tvm_ffi::match_any_internal::ArmId;
+ }
+ });
+ let lookup_arm_id = expand_leaf_table_lookup(tvm_ffi, arms,
&arm_constants, &view);
+
+ let ordered_selection = expand_ordered_try_into_chain(
+ arms,
+ quote!(#selected_enum::#fallback_variant),
+ &view,
+ &rejected,
+ |arm_id, _| {
+ let variant = &arm_variants[arm_id];
+
+ quote!(
+ ::core::result::Result::Ok(#selected_value) => {
+ #selected_enum::#variant(#selected_value)
+ }
+ )
+ },
+ );
+
+ let direct_selection = expand_direct_leaf_selection(
+ arms,
+ &arm_constants,
+ &arm_variants,
+ &arm_id,
+ &view,
+ &rejected,
+ &selected_enum,
+ &selected_value,
+ );
+ let body_dispatch = expand_leaf_body_dispatch(
+ arms,
+ &arm_variants,
+ &selected_enum,
+ &selected,
+ &fallback_variant,
+ fallback,
+ );
quote! {
{
+ enum #selected_enum<#(#arm_types),*> {
+ #(#arm_variants(#arm_types),)*
+ #fallback_variant,
+ }
+
+ #(#arm_constant_definitions)*
+
let #source = &(#scrutinee);
let #converted: ::core::result::Result<
#tvm_ffi::AnyView<'_>,
@@ -149,13 +419,37 @@ fn expand_match_any(input: MatchAnyInput) -> TokenStream {
::core::result::Result::Ok(view) => view,
::core::result::Result::Err(error) => match error {},
};
- if #view.type_index()
- >= #tvm_ffi::TypeIndex::kTVMFFIStaticObjectBegin as i32
- {
- #dispatch
- } else {
- #fallback
- }
+ let #selected =
+ if #view.type_index()
+ >= #tvm_ffi::TypeIndex::kTVMFFIStaticObjectBegin as i32
+ {
+ match #lookup_arm_id {
+ ::core::result::Result::Ok(
+ ::core::option::Option::Some(#arm_id),
+ ) => {
+ #direct_selection
+ }
+ ::core::result::Result::Ok(
+ ::core::option::Option::None,
+ ) => {
+ #selected_enum::#fallback_variant
+ }
+ ::core::result::Result::Err(()) => {
+ #ordered_selection
+ }
+ }
+ } else {
+ #selected_enum::#fallback_variant
+ };
+ #body_dispatch
}
}
}
+
+fn is_simple_binding(binding: &Pat) -> bool {
+ match binding {
+ Pat::Ident(binding) => binding.subpat.is_none(),
+ Pat::Wild(_) => true,
+ _ => false,
+ }
+}
diff --git a/rust/tvm-ffi-macros/src/object_macros.rs
b/rust/tvm-ffi-macros/src/object_macros.rs
index 14dc76fd..a285740e 100644
--- a/rust/tvm-ffi-macros/src/object_macros.rs
+++ b/rust/tvm-ffi-macros/src/object_macros.rs
@@ -31,6 +31,11 @@ pub fn derive_object(input: proc_macro::TokenStream) ->
TokenStream {
let type_key = get_attr(&derive_input, "type_key")
.map(attr_to_str)
.expect("Expect #[type_key = \"<my_type_key>\"] attribute");
+ let type_final = match get_attr(&derive_input, "type_final") {
+ Some(attr) if matches!(attr.parse_meta(), Ok(syn::Meta::Path(_))) =>
true,
+ Some(_) => panic!("Expect #[type_final] attribute"),
+ None => false,
+ };
// type index can be optional
// for now we make it required for static index
@@ -70,6 +75,20 @@ pub fn derive_object(input: proc_macro::TokenStream) ->
TokenStream {
};
// search for field name base and derive the base type
// we expect base always to be the first field
+ let final_parent_check = match &derive_input.data {
+ syn::Data::Struct(s) => s.fields.iter().next().and_then(|f| {
+ let base_ty = f.ty.clone();
+ Some(quote! {
+ const _: () = {
+ ::core::assert!(
+ !<#base_ty as
#tvm_ffi_crate::object::ObjectCore>::TYPE_FINAL,
+ "an object type cannot derive from a final parent"
+ );
+ };
+ })
+ }),
+ _ => panic!("First field must be `<base_name>: <ObjectCoreType>`"),
+ };
let base_def_tokens = match &derive_input.data {
syn::Data::Struct(s) => s.fields.iter().next().and_then(|f| {
let (base_id, base_ty) = (f.ident.clone()?, f.ty.clone());
@@ -94,8 +113,11 @@ pub fn derive_object(input: proc_macro::TokenStream) ->
TokenStream {
};
let expanded = quote! {
+ #final_parent_check
+
unsafe impl #tvm_ffi_crate::object::ObjectCore for #struct_name {
const TYPE_KEY: &'static str = #type_key;
+ const TYPE_FINAL: bool = #type_final;
#type_index_tokens
@@ -146,6 +168,19 @@ pub fn derive_object_ref(input: proc_macro::TokenStream)
-> TokenStream {
// implement AnyCompatible for #struct_name
unsafe impl #tvm_ffi_crate::type_traits::AnyCompatible for
#struct_name {
+ const MATCH_ANY_EXACT: bool = {
+ type ContainerType =
+ <#struct_name as
#tvm_ffi_crate::object::ObjectRefCore>::ContainerType;
+ <ContainerType as
#tvm_ffi_crate::object::ObjectCore>::TYPE_FINAL
+ };
+
+ #[inline]
+ fn match_any_exact_type_index() -> i32 {
+ type ContainerType = <#struct_name as
#tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ <ContainerType as
#tvm_ffi_crate::object::ObjectCore>::type_index()
+ }
+
fn type_str() -> std::string::String {
type ContainerType = <#struct_name as
#tvm_ffi_crate::object::ObjectRefCore>
::ContainerType;
diff --git a/rust/tvm-ffi/src/extra/module.rs b/rust/tvm-ffi/src/extra/module.rs
index 82630dfb..84cc4780 100644
--- a/rust/tvm-ffi/src/extra/module.rs
+++ b/rust/tvm-ffi/src/extra/module.rs
@@ -31,6 +31,7 @@ use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
#[derive(Object)]
#[type_key = "ffi.Module"]
#[type_index(TypeIndex::kTVMFFIModule)]
+#[type_final]
pub struct ModuleObj {
object: Object,
}
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index 1ba502b5..12492308 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -26,6 +26,8 @@ pub mod extra;
pub mod function;
pub mod function_internal;
pub mod macros;
+#[doc(hidden)]
+pub mod match_any_internal;
pub mod object;
pub mod optional;
pub mod string;
diff --git a/rust/tvm-ffi/src/match_any_internal.rs
b/rust/tvm-ffi/src/match_any_internal.rs
new file mode 100644
index 00000000..ef3e8c7e
--- /dev/null
+++ b/rust/tvm-ffi/src/match_any_internal.rs
@@ -0,0 +1,149 @@
+/*
+ * 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::any::TypeId;
+use std::marker::PhantomData;
+
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+
+use crate::AnyCompatible;
+
+/// A call-site-local arm number generated by `match_any!`.
+#[doc(hidden)]
+pub type ArmId = u32;
+
+/// Internal map from runtime object type indices to call-site-local arm
numbers.
+///
+/// The table stores only [`ArmId`] values. Arm bodies remain native branches
in
+/// the macro expansion and are never stored as functions or closures.
+#[doc(hidden)]
+pub struct LeafLookupTable {
+ pattern_list_id: TypeId,
+ arm_by_type_index: Box<[(i32, ArmId)]>,
+}
+
+impl LeafLookupTable {
+ #[doc(hidden)]
+ pub fn build(pattern_list_id: TypeId, entries: &[(i32, ArmId)]) -> Self {
+ assert!(
+ !entries.is_empty(),
+ "match_any! leaf lookup requires at least one arm"
+ );
+ for &(type_index, _) in entries {
+ assert!(
+ type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32,
+ "match_any! leaf pattern returned a non-object type index"
+ );
+ }
+
+ let mut arm_by_type_index = entries.to_vec();
+ // Stable sorting plus deduplication preserves the first arm for a
+ // duplicate runtime type index.
+ arm_by_type_index.sort_by_key(|&(type_index, _)| type_index);
+ arm_by_type_index.dedup_by_key(|entry| entry.0);
+
+ Self {
+ pattern_list_id,
+ arm_by_type_index: arm_by_type_index.into_boxed_slice(),
+ }
+ }
+
+ #[doc(hidden)]
+ #[inline(always)]
+ pub fn lookup(&self, pattern_list_id: TypeId, type_index: i32) ->
Result<Option<ArmId>, ()> {
+ if self.pattern_list_id != pattern_list_id {
+ return Err(());
+ }
+
+ Ok(self
+ .arm_by_type_index
+ .binary_search_by_key(&type_index, |&(entry_type_index, _)|
entry_type_index)
+ .ok()
+ .map(|offset| self.arm_by_type_index[offset].1))
+ }
+}
+
+/// Type-level list used to collect leaf object-pattern metadata.
+#[doc(hidden)]
+pub trait LeafPatternList: 'static {
+ #[doc(hidden)]
+ const ALL_EXACT_LEAF: bool;
+
+ #[doc(hidden)]
+ fn fill_leaf_type_indices(out: &mut [i32]);
+}
+
+impl LeafPatternList for () {
+ const ALL_EXACT_LEAF: bool = true;
+
+ fn fill_leaf_type_indices(out: &mut [i32]) {
+ debug_assert!(out.is_empty());
+ }
+}
+
+impl<Head, Tail> LeafPatternList for (Head, Tail)
+where
+ Head: AnyCompatible + 'static,
+ Tail: LeafPatternList,
+{
+ const ALL_EXACT_LEAF: bool = Head::MATCH_ANY_EXACT && Tail::ALL_EXACT_LEAF;
+
+ fn fill_leaf_type_indices(out: &mut [i32]) {
+ let (head, tail) = out
+ .split_first_mut()
+ .expect("match_any! pattern metadata length mismatch");
+ *head = Head::match_any_exact_type_index();
+ Tail::fill_leaf_type_indices(tail);
+ }
+}
+
+/// Probe used by the macro to retain ordered matching for non-leaf patterns.
+#[doc(hidden)]
+pub struct LeafPatternProbe<T>(PhantomData<fn() -> T>);
+
+impl<T> LeafPatternProbe<T> {
+ #[doc(hidden)]
+ pub const fn new() -> Self {
+ Self(PhantomData)
+ }
+}
+
+/// Leaf metadata with an autoref fallback for non-`AnyCompatible` patterns.
+#[doc(hidden)]
+pub trait LeafPatternMetadata {
+ #[doc(hidden)]
+ fn leaf_pattern_list_id(&self) -> Option<TypeId> {
+ None
+ }
+
+ #[doc(hidden)]
+ fn fill_leaf_type_indices(&self, _out: &mut [i32]) {}
+}
+
+impl<T> LeafPatternMetadata for &LeafPatternProbe<T> {}
+
+impl<T: LeafPatternList> LeafPatternMetadata for LeafPatternProbe<T> {
+ fn leaf_pattern_list_id(&self) -> Option<TypeId> {
+ T::ALL_EXACT_LEAF.then(TypeId::of::<T>)
+ }
+
+ fn fill_leaf_type_indices(&self, out: &mut [i32]) {
+ T::fill_leaf_type_indices(out);
+ }
+}
diff --git a/rust/tvm-ffi/src/object.rs b/rust/tvm-ffi/src/object.rs
index 78394001..712aa35e 100644
--- a/rust/tvm-ffi/src/object.rs
+++ b/rust/tvm-ffi/src/object.rs
@@ -60,6 +60,11 @@ pub unsafe trait ObjectCore: Sized + 'static {
/// depth one greater than its parent. This value must be non-negative and
/// agree with the runtime type table entry for `Self`.
const TYPE_DEPTH: i32;
+ /// Whether every instance of this type has exactly `Self::type_index()`.
+ ///
+ /// A final type has no separately registered object-system subtype.
+ #[doc(hidden)]
+ const TYPE_FINAL: bool = false;
// return the type index of the object
fn type_index() -> i32;
/// Return the object header
diff --git a/rust/tvm-ffi/src/type_traits.rs b/rust/tvm-ffi/src/type_traits.rs
index d39da4b4..9e7d0871 100644
--- a/rust/tvm-ffi/src/type_traits.rs
+++ b/rust/tvm-ffi/src/type_traits.rs
@@ -25,6 +25,23 @@ use tvm_ffi_sys::{TVMFFIAny, TVMFFIGetTypeInfo};
/// Trait to enable a value to be compatible with Any
/// Enables TryFrom/Into AnyView/Any
pub unsafe trait AnyCompatible: Sized {
+ /// Whether the `AnyView -> Self` conversion matches one exact leaf type.
+ ///
+ /// If this is `true`, [`Self::match_any_exact_type_index`] must always
+ /// return the same index for the process lifetime, and conversion from an
+ /// object-backed `AnyView` must succeed if and only if its runtime type
+ /// index equals that index. Matching must not depend on any other part of
+ /// the value. `match_any!` may use this contract to skip conversions for
+ /// all other arms.
+ #[doc(hidden)]
+ const MATCH_ANY_EXACT: bool = false;
+
+ /// Return the sole runtime object type index accepted by this conversion.
+ #[doc(hidden)]
+ fn match_any_exact_type_index() -> i32 {
+ unreachable!("MATCH_ANY_EXACT is false")
+ }
+
/// the value to copy to TVMFFIAny
unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny);
/// consume the value to move to Any
diff --git a/rust/tvm-ffi/tests/test_match_any.rs
b/rust/tvm-ffi/tests/test_match_any.rs
index de40b482..ab831f79 100644
--- a/rust/tvm-ffi/tests/test_match_any.rs
+++ b/rust/tvm-ffi/tests/test_match_any.rs
@@ -17,7 +17,10 @@
* under the License.
*/
-use tvm_ffi::{match_any, Any, AnyView, Array, Map, Shape, Tensor};
+use std::any::TypeId;
+
+use tvm_ffi::match_any_internal::{ArmId, LeafLookupTable, LeafPatternMetadata,
LeafPatternProbe};
+use tvm_ffi::{match_any, Any, AnyView, Array, Function, Map, Module, Shape,
Tensor, TypeIndex};
#[test]
fn matches_concrete_object_containers_in_source_order() {
@@ -59,3 +62,87 @@ fn matches_concrete_object_containers_in_source_order() {
};
assert_eq!(matched_view, ("tensor", 2));
}
+
+#[test]
+fn parameterized_containers_keep_ordered_conversion() {
+ let array = [1.5_f64, 2.5].into_iter().collect::<Array<f64>>();
+ let selected = match_any! {
+ Any::from(array) {
+ Array::<i64>(_) => "integer array",
+ Tensor(_) => "tensor",
+ Shape(_) => "shape",
+ Array::<f64>(_) => "float array",
+ _ => "unsupported",
+ }
+ };
+
+ assert_eq!(selected, "float array");
+}
+
+#[test]
+fn leaf_lookup_keeps_the_first_arm() {
+ fn classify(value: Any) -> usize {
+ match_any! {
+ value {
+ Module(_) => 0,
+ Module(_) => 1,
+ _ => 2,
+ }
+ }
+ }
+
+ let module: Module = Function::get_global("ffi.SystemLib")
+ .unwrap()
+ .call_tuple_with_len::<0, _>(())
+ .unwrap()
+ .try_into()
+ .unwrap();
+ assert_eq!(classify(Any::from(module)), 0);
+ assert_eq!(classify(Any::from(Array::<i64>::default())), 2);
+}
+
+#[test]
+fn lookup_table_maps_runtime_indices_to_local_arm_ids() {
+ const ARM_0: ArmId = 0;
+ const ARM_1: ArmId = 1;
+ const ARM_2: ArmId = 2;
+ let pattern_list_id = TypeId::of::<(i32, i64, f32)>();
+ let table = LeafLookupTable::build(pattern_list_id, &[(73, ARM_0), (73,
ARM_1), (75, ARM_2)]);
+
+ assert_eq!(table.lookup(pattern_list_id, 73), Ok(Some(ARM_0)));
+ assert_eq!(table.lookup(pattern_list_id, 72), Ok(None));
+ assert_eq!(table.lookup(pattern_list_id, 74), Ok(None));
+ assert_eq!(table.lookup(pattern_list_id, 75), Ok(Some(ARM_2)));
+ assert_eq!(table.lookup(pattern_list_id, 76), Ok(None));
+}
+
+#[test]
+fn a_generic_pattern_list_cannot_reuse_another_lists_table() {
+ let pattern_list_id = TypeId::of::<(i32, i64)>();
+ let table = LeafLookupTable::build(pattern_list_id, &[(73, 0), (75, 1)]);
+
+ assert_eq!(table.lookup(TypeId::of::<(u8, u16)>(), 73), Err(()));
+}
+
+#[test]
+fn metadata_only_accepts_exact_leaf_patterns() {
+ type Leaf = (Module, ());
+ let leaf = LeafPatternProbe::<Leaf>::new();
+ let mut type_indices = [0; 1];
+ assert!((&leaf).leaf_pattern_list_id().is_some());
+ (&leaf).fill_leaf_type_indices(&mut type_indices);
+ assert!(type_indices[0] >= TypeIndex::kTVMFFIStaticObjectBegin as i32);
+
+ type Parameterized = (Array<i64>, ());
+ let parameterized = LeafPatternProbe::<Parameterized>::new();
+ assert!((¶meterized).leaf_pattern_list_id().is_none());
+
+ type NonFinal = (Tensor, ());
+ let non_final = LeafPatternProbe::<NonFinal>::new();
+ assert!((&non_final).leaf_pattern_list_id().is_none());
+
+ struct NoAnyCompatibleMetadata;
+ type Custom = (NoAnyCompatibleMetadata, ());
+ let custom = LeafPatternProbe::<Custom>::new();
+ assert!((&custom).leaf_pattern_list_id().is_none());
+}
