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 a7eafccd [Perf][Rust] Optimize match_any! ordered conversion and
exact-leaf dispatch (#692)
a7eafccd is described below
commit a7eafccdf2f75df35d0e7819cad7c975d63f4091
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Jul 30 21:02:35 2026 -0400
[Perf][Rust] Optimize match_any! ordered conversion and exact-leaf dispatch
(#692)
This PR improves both Rust `match_any!` dispatch paths without changing
how matching works.
## Ordered dispatch
When an arm does not match, the internal conversion now returns
`Err(())` instead of creating a `TypeError`. This makes failed checks
cheaper while keeping the public `TryFrom<AnyView>` API unchanged.
Custom `TryInto` matchers are still supported and continue to run in
source order.
## Exact-leaf dispatch
For a match with many exact final object types, the macro builds one
`TypeIndex → ArmId` table through `OnceLock`.
The table uses the smallest pattern `TypeIndex` as its starting point:
```text
base = minimum pattern TypeIndex
arm_ids[pattern_type_index - base] = source ArmId
```
Each call then selects an arm directly with:
```text
arm_ids[runtime_type_index - base]
```
This is an O(1) lookup. Only the selected object handle is created, and
its type check is not repeated.
The direct table is used when there are at least 20 typed arms, every
pattern matches one exact final runtime type, there are no guards, and
bindings are simple names or `_`. Smaller matches, guarded arms,
parent-type patterns, parameterized containers, and custom matchers
continue using ordered dispatch.
The input value is still evaluated once. Duplicate types still select
the first arm, and unmatched or non-object values still use the final
`_` fallback. A `TypeId` check prevents the same static table from being
incorrectly shared by different generic pattern lists.
---
rust/tvm-ffi-macros/src/match_any.rs | 96 +++++++---
rust/tvm-ffi/src/any.rs | 25 +++
rust/tvm-ffi/src/match_any_internal.rs | 131 +++++++++++---
rust/tvm-ffi/tests/test_match_any.rs | 90 +++++----
rust/tvm-ffi/tests/test_match_any_ast_lookup.rs | 231 ++++++++++++++++++++++++
5 files changed, 483 insertions(+), 90 deletions(-)
diff --git a/rust/tvm-ffi-macros/src/match_any.rs
b/rust/tvm-ffi-macros/src/match_any.rs
index 51a479bc..e9574fd8 100644
--- a/rust/tvm-ffi-macros/src/match_any.rs
+++ b/rust/tvm-ffi-macros/src/match_any.rs
@@ -24,10 +24,8 @@ 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;
+// Avoid call-site table setup for small matches.
+const MIN_LOOKUP_TABLE_ARMS: usize = 20;
struct MatchAnyInput {
scrutinee: Expr,
@@ -136,7 +134,7 @@ fn expand_ordered_match(
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 dispatch = expand_ordered_dispatch(arms, fallback, &view, &rejected);
+ let dispatch = expand_ordered_dispatch(tvm_ffi, arms, fallback, &view,
&rejected);
quote! {
{
@@ -161,23 +159,32 @@ fn expand_ordered_match(
}
fn expand_ordered_dispatch(
+ tvm_ffi: &TokenStream,
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 })
- }
- })
+ expand_ordered_try_into_chain(
+ tvm_ffi,
+ 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>(
+ tvm_ffi: &TokenStream,
arms: &[TypedArm],
fallback: TokenStream,
view: &Ident,
@@ -193,9 +200,10 @@ where
.fold(fallback, |next, (arm_id, arm)| {
let matcher = &arm.matcher;
let matched = matched_arm(arm_id, arm);
+ let conversion = expand_pattern_conversion(tvm_ffi, matcher, view);
quote! {
- match ::core::convert::TryInto::<#matcher>::try_into(#view) {
+ match #conversion {
#matched,
#rejected => {
::core::mem::drop(#rejected);
@@ -206,12 +214,48 @@ where
})
}
-fn expand_leaf_table_lookup(
+fn expand_pattern_conversion(tvm_ffi: &TokenStream, matcher: &Path, view:
&Ident) -> TokenStream {
+ let span = Span::mixed_site();
+ let probe = Ident::new("__tvm_ffi_match_any_conversion_probe", span);
+ let converted = Ident::new("__tvm_ffi_match_any_pattern_conversion", span);
+
+ quote! {
+ {
+ use #tvm_ffi::match_any_internal::PatternConversion as _;
+
+ let #probe =
+
#tvm_ffi::match_any_internal::PatternConversionProbe::<#matcher>::new();
+ let #converted: ::core::result::Result<#matcher, ()> =
+ (&#probe).try_convert(#view);
+ #converted
+ }
+ }
+}
+
+fn expand_exact_pattern_conversion(
tvm_ffi: &TokenStream,
- arms: &[TypedArm],
- arm_constants: &[Ident],
+ matcher: &Path,
view: &Ident,
) -> TokenStream {
+ let span = Span::mixed_site();
+ let probe = Ident::new("__tvm_ffi_match_any_conversion_probe", span);
+ let converted = Ident::new("__tvm_ffi_match_any_pattern_conversion", span);
+
+ quote! {
+ {
+ use #tvm_ffi::match_any_internal::PatternConversion as _;
+
+ let #probe =
+
#tvm_ffi::match_any_internal::PatternConversionProbe::<#matcher>::new();
+ let #converted: ::core::result::Result<#matcher, ()> = unsafe {
+ (&#probe).try_convert_after_exact_match(#view)
+ };
+ #converted
+ }
+ }
+}
+
+fn expand_leaf_table_lookup(tvm_ffi: &TokenStream, arms: &[TypedArm], 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);
@@ -219,10 +263,6 @@ fn expand_leaf_table_lookup(
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)
@@ -245,7 +285,7 @@ fn expand_leaf_table_lookup(
(&#probe).fill_leaf_type_indices(&mut #type_indices);
#tvm_ffi::match_any_internal::LeafLookupTable::build(
#pattern_list_id,
- &[#(#lookup_entries),*],
+ &#type_indices,
)
});
#table.lookup(#pattern_list_id, #view.type_index())
@@ -259,6 +299,7 @@ fn expand_leaf_table_lookup(
}
fn expand_direct_leaf_selection(
+ tvm_ffi: &TokenStream,
arms: &[TypedArm],
arm_constants: &[Ident],
arm_variants: &[Ident],
@@ -272,10 +313,11 @@ fn expand_direct_leaf_selection(
let matcher = &arm.matcher;
let variant = &arm_variants[arm_id];
let arm_constant = &arm_constants[arm_id];
+ let conversion = expand_exact_pattern_conversion(tvm_ffi, matcher,
view);
quote! {
#arm_constant => {
- match ::core::convert::TryInto::<#matcher>::try_into(#view) {
+ match #conversion {
::core::result::Result::Ok(#selected_value) => {
#selected_enum::#variant(#selected_value)
}
@@ -364,16 +406,15 @@ fn expand_leaf_lookup_match(
#arm_id as #tvm_ffi::match_any_internal::ArmId;
}
});
- let lookup_arm_id = expand_leaf_table_lookup(tvm_ffi, arms,
&arm_constants, &view);
-
+ let lookup_arm_id = expand_leaf_table_lookup(tvm_ffi, arms, &view);
let ordered_selection = expand_ordered_try_into_chain(
+ tvm_ffi,
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)
@@ -383,6 +424,7 @@ fn expand_leaf_lookup_match(
);
let direct_selection = expand_direct_leaf_selection(
+ tvm_ffi,
arms,
&arm_constants,
&arm_variants,
diff --git a/rust/tvm-ffi/src/any.rs b/rust/tvm-ffi/src/any.rs
index 0c4c8476..47379f47 100644
--- a/rust/tvm-ffi/src/any.rs
+++ b/rust/tvm-ffi/src/any.rs
@@ -232,6 +232,31 @@ impl Drop for Any {
}
}
+/// Convert an [`AnyView`] without constructing a diagnostic error on mismatch.
+#[inline]
+pub(crate) fn try_cast_from_any_view<T>(value: &AnyView<'_>) -> Result<T, ()>
+where
+ T: AnyCompatible,
+{
+ unsafe { T::try_cast_from_any_view(&value.data) }
+}
+
+/// Copy a value after exact-leaf lookup has established compatibility.
+///
+/// # Safety
+///
+/// `T::MATCH_ANY_EXACT` must be true and `value.type_index()` must equal
+/// `T::match_any_exact_type_index()`.
+#[inline(always)]
+pub(crate) unsafe fn copy_from_any_view_after_check<T>(value: &AnyView<'_>) ->
T
+where
+ T: AnyCompatible,
+{
+ debug_assert!(T::MATCH_ANY_EXACT);
+ debug_assert!(T::check_any_strict(&value.data));
+ T::copy_from_any_view_after_check(&value.data)
+}
+
// convert Any ref to AnyView
impl<'a> From<&'a Any> for AnyView<'a> {
#[inline]
diff --git a/rust/tvm-ffi/src/match_any_internal.rs
b/rust/tvm-ffi/src/match_any_internal.rs
index ef3e8c7e..8c785426 100644
--- a/rust/tvm-ffi/src/match_any_internal.rs
+++ b/rust/tvm-ffi/src/match_any_internal.rs
@@ -22,45 +22,125 @@ use std::marker::PhantomData;
use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
-use crate::AnyCompatible;
+use crate::{AnyCompatible, AnyView};
/// 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.
+/// Conversion adapter used by `match_any!` typed arms.
+#[doc(hidden)]
+pub struct PatternConversionProbe<T>(PhantomData<fn() -> T>);
+
+impl<T> PatternConversionProbe<T> {
+ #[doc(hidden)]
+ pub const fn new() -> Self {
+ Self(PhantomData)
+ }
+}
+
+/// Prefer the lightweight `AnyCompatible` conversion while retaining a
+/// `TryInto` fallback for custom matcher types.
+#[doc(hidden)]
+pub trait PatternConversion<'a, T> {
+ #[doc(hidden)]
+ fn try_convert(&self, view: AnyView<'a>) -> Result<T, ()>;
+
+ /// Convert after exact-leaf lookup has established that `view` matches
+ /// `T`. Custom `TryInto` matchers retain their checked conversion.
+ ///
+ /// # Safety
+ ///
+ /// The caller must establish the `AnyCompatible::MATCH_ANY_EXACT`
+ /// contract before using the unchecked implementation.
+ #[doc(hidden)]
+ unsafe fn try_convert_after_exact_match(&self, view: AnyView<'a>) ->
Result<T, ()> {
+ self.try_convert(view)
+ }
+}
+
+impl<'a, T: AnyCompatible> PatternConversion<'a, T> for
PatternConversionProbe<T> {
+ #[inline(always)]
+ fn try_convert(&self, view: AnyView<'a>) -> Result<T, ()> {
+ if T::MATCH_ANY_EXACT {
+ if view.type_index() == T::match_any_exact_type_index() {
+ Ok(unsafe {
crate::any::copy_from_any_view_after_check::<T>(&view) })
+ } else {
+ Err(())
+ }
+ } else {
+ crate::any::try_cast_from_any_view::<T>(&view)
+ }
+ }
+
+ #[inline(always)]
+ unsafe fn try_convert_after_exact_match(&self, view: AnyView<'a>) ->
Result<T, ()> {
+ Ok(crate::any::copy_from_any_view_after_check::<T>(&view))
+ }
+}
+
+impl<'a, T> PatternConversion<'a, T> for &PatternConversionProbe<T>
+where
+ AnyView<'a>: TryInto<T>,
+{
+ #[inline(always)]
+ fn try_convert(&self, view: AnyView<'a>) -> Result<T, ()> {
+ view.try_into().map_err(|_| ())
+ }
+}
+
+/// Runtime type-index lookup for one `match_any!` call site.
///
-/// The table stores only [`ArmId`] values. Arm bodies remain native branches
in
-/// the macro expansion and are never stored as functions or closures.
+/// The smallest pattern `TypeIndex` is used as the base of a direct table.
+/// A successful lookup returns the matching position as a call-site-local
+/// [`ArmId`]. 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)]>,
+ base: i32,
+ arm_ids: Box<[ArmId]>,
}
+const NO_ARM_ID: ArmId = ArmId::MAX;
+
impl LeafLookupTable {
#[doc(hidden)]
- pub fn build(pattern_list_id: TypeId, entries: &[(i32, ArmId)]) -> Self {
+ #[cold]
+ #[inline(never)]
+ pub fn build(pattern_list_id: TypeId, type_indices: &[i32]) -> Self {
assert!(
- !entries.is_empty(),
+ !type_indices.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"
- );
+ assert!(
+ type_indices
+ .iter()
+ .all(|&type_index| type_index >=
TypeIndex::kTVMFFIStaticObjectBegin as i32),
+ "match_any! leaf pattern returned a non-object type index"
+ );
+ let min_type_index = *type_indices.iter().min().unwrap();
+ let max_type_index = *type_indices.iter().max().unwrap();
+ let span = usize::try_from(i64::from(max_type_index) -
i64::from(min_type_index) + 1)
+ .expect("match_any! leaf pattern returned an invalid type-index
span");
+ assert!(
+ type_indices.len() <= NO_ARM_ID as usize,
+ "match_any! has too many exact-leaf arms"
+ );
+ let mut arm_ids = vec![NO_ARM_ID; span];
+ for (arm_id, &type_index) in type_indices.iter().enumerate() {
+ let offset = (type_index - min_type_index) as usize;
+ let slot = &mut arm_ids[offset];
+ // Keep the first source arm for duplicate runtime indices.
+ if *slot == NO_ARM_ID {
+ *slot = arm_id as ArmId;
+ }
}
- 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(),
+ base: min_type_index,
+ arm_ids: arm_ids.into_boxed_slice(),
}
}
@@ -70,12 +150,9 @@ impl LeafLookupTable {
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))
+ let offset = type_index.wrapping_sub(self.base) as usize;
+ let arm_id = self.arm_ids.get(offset).copied().unwrap_or(NO_ARM_ID);
+ Ok((arm_id != NO_ARM_ID).then_some(arm_id))
}
}
@@ -92,6 +169,7 @@ pub trait LeafPatternList: 'static {
impl LeafPatternList for () {
const ALL_EXACT_LEAF: bool = true;
+ #[inline(always)]
fn fill_leaf_type_indices(out: &mut [i32]) {
debug_assert!(out.is_empty());
}
@@ -104,6 +182,7 @@ where
{
const ALL_EXACT_LEAF: bool = Head::MATCH_ANY_EXACT && Tail::ALL_EXACT_LEAF;
+ #[inline(always)]
fn fill_leaf_type_indices(out: &mut [i32]) {
let (head, tail) = out
.split_first_mut()
@@ -139,10 +218,12 @@ pub trait LeafPatternMetadata {
impl<T> LeafPatternMetadata for &LeafPatternProbe<T> {}
impl<T: LeafPatternList> LeafPatternMetadata for LeafPatternProbe<T> {
+ #[inline(always)]
fn leaf_pattern_list_id(&self) -> Option<TypeId> {
T::ALL_EXACT_LEAF.then(TypeId::of::<T>)
}
+ #[inline(always)]
fn fill_leaf_type_indices(&self, out: &mut [i32]) {
T::fill_leaf_type_indices(out);
}
diff --git a/rust/tvm-ffi/tests/test_match_any.rs
b/rust/tvm-ffi/tests/test_match_any.rs
index ab831f79..40fdd90b 100644
--- a/rust/tvm-ffi/tests/test_match_any.rs
+++ b/rust/tvm-ffi/tests/test_match_any.rs
@@ -22,6 +22,16 @@ 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};
+struct CustomModuleMatcher;
+
+impl<'a> TryFrom<AnyView<'a>> for CustomModuleMatcher {
+ type Error = &'static str;
+
+ fn try_from(value: AnyView<'a>) -> Result<Self, Self::Error> {
+ value.try_as::<Module>().map(|_| Self).ok_or("not a Module")
+ }
+}
+
#[test]
fn matches_concrete_object_containers_in_source_order() {
fn classify(expr: Any) -> (&'static str, usize) {
@@ -64,29 +74,12 @@ fn matches_concrete_object_containers_in_source_order() {
}
#[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 {
+fn custom_try_into_matcher_keeps_ordered_compatibility() {
+ fn matches(value: AnyView<'_>) -> bool {
match_any! {
value {
- Module(_) => 0,
- Module(_) => 1,
- _ => 2,
+ CustomModuleMatcher(_) => true,
+ _ => false,
}
}
}
@@ -97,31 +90,52 @@ fn leaf_lookup_keeps_the_first_arm() {
.unwrap()
.try_into()
.unwrap();
- assert_eq!(classify(Any::from(module)), 0);
- assert_eq!(classify(Any::from(Array::<i64>::default())), 2);
+ assert!(matches(AnyView::from(&module)));
+ assert!(!matches(AnyView::from(&Shape::from([1_i64, 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)]);
+fn parameterized_containers_use_complete_conversion_semantics() {
+ let array = [1.5_f64, 2.5].into_iter().collect::<Array<f64>>();
+ // Both patterns have the same runtime Array TypeIndex, so matching must
+ // inspect the element types in source order.
+ let selected = match_any! {
+ Any::from(array) {
+ Array::<i64>(_) => "integer array",
+ Array::<f64>(_) => "float array",
+ _ => "unsupported",
+ }
+ };
- 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));
+ assert_eq!(selected, "float array");
}
#[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)]);
+fn direct_lookup_table_maps_runtime_indices_to_local_arm_ids() {
+ const ARM_0: ArmId = 0;
+ const ARM_2: ArmId = 2;
+ let pattern_list_id = TypeId::of::<(i32, i64, f32)>();
+ let object_begin = TypeIndex::kTVMFFIStaticObjectBegin as i32;
+ let table = LeafLookupTable::build(
+ pattern_list_id,
+ &[object_begin + 4, object_begin + 4, object_begin + 7],
+ );
- assert_eq!(table.lookup(TypeId::of::<(u8, u16)>(), 73), Err(()));
+ assert_eq!(table.lookup(pattern_list_id, object_begin + 3), Ok(None));
+ assert_eq!(
+ table.lookup(pattern_list_id, object_begin + 4),
+ Ok(Some(ARM_0))
+ );
+ assert_eq!(table.lookup(pattern_list_id, object_begin + 5), Ok(None));
+ assert_eq!(
+ table.lookup(pattern_list_id, object_begin + 7),
+ Ok(Some(ARM_2))
+ );
+ assert_eq!(table.lookup(pattern_list_id, object_begin + 8), Ok(None));
+ assert_eq!(
+ table.lookup(TypeId::of::<(u8, u16)>(), object_begin + 4),
+ Err(())
+ );
}
#[test]
diff --git a/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
b/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
new file mode 100644
index 00000000..7d69189c
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
@@ -0,0 +1,231 @@
+/*
+ * 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 tvm_ffi::derive::{Object, ObjectRef};
+use tvm_ffi::object::{Object as ObjectBase, ObjectArc, ObjectCore};
+use tvm_ffi::{match_any, Any, Array, Shape, TypeIndex};
+use tvm_ffi_sys::TVMFFIByteArray;
+
+unsafe extern "C" {
+ fn TVMFFITypeGetOrAllocIndex(
+ type_key: *const TVMFFIByteArray,
+ static_type_index: i32,
+ type_depth: i32,
+ num_child_slots: i32,
+ child_slots_can_overflow: i32,
+ parent_type_index: i32,
+ ) -> i32;
+}
+
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "testing.match_any.Expr"]
+struct ExprObj {
+ base: ObjectBase,
+}
+
+fn register_type<T: ObjectCore>(num_child_slots: i32, parent_type_index: i32)
-> i32 {
+ let type_key = unsafe { TVMFFIByteArray::from_str(T::TYPE_KEY) };
+ let type_index = unsafe {
+ TVMFFITypeGetOrAllocIndex(
+ &type_key,
+ -1,
+ T::TYPE_DEPTH,
+ num_child_slots,
+ 0,
+ parent_type_index,
+ )
+ };
+ assert!(type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32);
+ type_index
+}
+
+macro_rules! define_expr_leaves {
+ ($(($object:ident, $handle:ident, $type_key:literal)),+ $(,)?) => {
+ $(
+ #[repr(C)]
+ #[derive(Object)]
+ #[type_key = $type_key]
+ #[type_final]
+ struct $object {
+ base: ExprObj,
+ }
+
+ #[repr(C)]
+ #[derive(ObjectRef, Clone)]
+ struct $handle {
+ data: ObjectArc<$object>,
+ }
+
+ impl Default for $handle {
+ fn default() -> Self {
+ Self {
+ data: ObjectArc::new($object {
+ base: ExprObj {
+ base: ObjectBase::new(),
+ },
+ }),
+ }
+ }
+ }
+ )+
+
+ fn register_expr_types() {
+ let expr_type_index =
+ register_type::<ExprObj>(20,
TypeIndex::kTVMFFIStaticObjectBegin as i32);
+ $(
+ register_type::<$object>(0, expr_type_index);
+ )+
+ }
+ };
+}
+
+// These final nodes model the kind of downstream AST hierarchy that motivates
+// exact-leaf dispatch. This integration-test binary intentionally has one
test,
+// so its process-wide type registrations happen before any concurrent lookup.
+define_expr_leaves!(
+ (AddExprObj, AddExpr, "testing.match_any.AddExpr"),
+ (SubExprObj, SubExpr, "testing.match_any.SubExpr"),
+ (MulExprObj, MulExpr, "testing.match_any.MulExpr"),
+ (DivExprObj, DivExpr, "testing.match_any.DivExpr"),
+ (ModExprObj, ModExpr, "testing.match_any.ModExpr"),
+ (NegExprObj, NegExpr, "testing.match_any.NegExpr"),
+ (CallExprObj, CallExpr, "testing.match_any.CallExpr"),
+ (LetExprObj, LetExpr, "testing.match_any.LetExpr"),
+ (IfExprObj, IfExpr, "testing.match_any.IfExpr"),
+ (TupleExprObj, TupleExpr, "testing.match_any.TupleExpr"),
+ (
+ TupleGetItemExprObj,
+ TupleGetItemExpr,
+ "testing.match_any.TupleGetItemExpr"
+ ),
+ (CastExprObj, CastExpr, "testing.match_any.CastExpr"),
+ (LoadExprObj, LoadExpr, "testing.match_any.LoadExpr"),
+ (StoreExprObj, StoreExpr, "testing.match_any.StoreExpr"),
+ (ForExprObj, ForExpr, "testing.match_any.ForExpr"),
+ (WhileExprObj, WhileExpr, "testing.match_any.WhileExpr"),
+ (SeqExprObj, SeqExpr, "testing.match_any.SeqExpr"),
+ (ReturnExprObj, ReturnExpr, "testing.match_any.ReturnExpr"),
+ (
+ ConstantExprObj,
+ ConstantExpr,
+ "testing.match_any.ConstantExpr"
+ ),
+ (VarExprObj, VarExpr, "testing.match_any.VarExpr"),
+);
+
+#[test]
+fn dispatches_representative_ast_leaf_nodes() {
+ register_expr_types();
+
+ fn classify(value: Any) -> &'static str {
+ match_any! {
+ value {
+ AddExpr(_) => "add",
+ SubExpr(_) => "sub",
+ MulExpr(_) => "mul",
+ DivExpr(_) => "div",
+ ModExpr(_) => "mod",
+ NegExpr(_) => "neg",
+ CallExpr(_) => "call",
+ LetExpr(_) => "let",
+ IfExpr(_) => "if",
+ TupleExpr(_) => "tuple",
+ TupleGetItemExpr(_) => "tuple_get_item",
+ CastExpr(_) => "cast",
+ LoadExpr(_) => "load",
+ StoreExpr(_) => "store",
+ ForExpr(_) => "for",
+ WhileExpr(_) => "while",
+ SeqExpr(_) => "seq",
+ ReturnExpr(_) => "return",
+ ConstantExpr(_) => "constant",
+ VarExpr(_) => "var",
+ _ => "unsupported",
+ }
+ }
+ }
+
+ let cases = [
+ (Any::from(AddExpr::default()), "add"),
+ (Any::from(SubExpr::default()), "sub"),
+ (Any::from(MulExpr::default()), "mul"),
+ (Any::from(DivExpr::default()), "div"),
+ (Any::from(ModExpr::default()), "mod"),
+ (Any::from(NegExpr::default()), "neg"),
+ (Any::from(CallExpr::default()), "call"),
+ (Any::from(LetExpr::default()), "let"),
+ (Any::from(IfExpr::default()), "if"),
+ (Any::from(TupleExpr::default()), "tuple"),
+ (Any::from(TupleGetItemExpr::default()), "tuple_get_item"),
+ (Any::from(CastExpr::default()), "cast"),
+ (Any::from(LoadExpr::default()), "load"),
+ (Any::from(StoreExpr::default()), "store"),
+ (Any::from(ForExpr::default()), "for"),
+ (Any::from(WhileExpr::default()), "while"),
+ (Any::from(SeqExpr::default()), "seq"),
+ (Any::from(ReturnExpr::default()), "return"),
+ (Any::from(ConstantExpr::default()), "constant"),
+ (Any::from(VarExpr::default()), "var"),
+ ];
+ for (value, expected) in cases {
+ assert_eq!(classify(value), expected);
+ }
+
+ assert_eq!(classify(Any::from(Shape::from([1_i64, 2]))), "unsupported");
+ assert_eq!(classify(Any::from(1_i64)), "unsupported");
+
+ // This call site has enough syntactically eligible arms to consider leaf
+ // lookup, but Array<T> requires complete type conversion. The entire match
+ // must therefore retain source-ordered dispatch.
+ fn classify_mixed_patterns(value: Any) -> &'static str {
+ match_any! {
+ value {
+ AddExpr(_) => "add",
+ SubExpr(_) => "sub",
+ MulExpr(_) => "mul",
+ DivExpr(_) => "div",
+ ModExpr(_) => "mod",
+ NegExpr(_) => "neg",
+ CallExpr(_) => "call",
+ LetExpr(_) => "let",
+ IfExpr(_) => "if",
+ TupleExpr(_) => "tuple",
+ TupleGetItemExpr(_) => "tuple_get_item",
+ CastExpr(_) => "cast",
+ LoadExpr(_) => "load",
+ StoreExpr(_) => "store",
+ ForExpr(_) => "for",
+ WhileExpr(_) => "while",
+ SeqExpr(_) => "seq",
+ ReturnExpr(_) => "return",
+ Array::<i64>(_) => "integer_array",
+ Array::<f64>(_) => "float_array",
+ _ => "unsupported",
+ }
+ }
+ }
+
+ assert_eq!(
+ classify_mixed_patterns(Any::from(AddExpr::default())),
+ "add"
+ );
+ let array = [1.5_f64, 2.5].into_iter().collect::<Array<f64>>();
+ assert_eq!(classify_mixed_patterns(Any::from(array)), "float_array");
+}
