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 135c2e8b [Rust][Feat] Add match_any! for object-backed values (#682)
135c2e8b is described below
commit 135c2e8bbf51a353db1c26a75404e03d03e802cb
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Jul 23 00:37:24 2026 -0400
[Rust][Feat] Add match_any! for object-backed values (#682)
This PR adds Rust match_any!, a Phase 1 macro that evaluates an
Any-compatible scrutinee once, converts it to AnyView, and tries
object-backed target conversions in source order using the existing
standard TryInto/TryFrom implementations. The first successful
conversion whose optional guard passes selects its arm; a failed
conversion or guard continues to the next arm, while a required final _
handles non-object or unmatched values. This implementation
intentionally does not add type-index dispatch tables, caches, custom
pattern protocols, or heuristic dispatch.
```rust
match_any! {
value {
Tensor(tensor)
if tensor.shape().len() == 2 => ("matrix",
tensor.shape().len()),
Tensor(tensor) => ("tensor", tensor.shape().len()),
Shape(shape) => ("shape", shape.len()),
Array::<i64>(array) => ("array", array.len()),
_ => ("unsupported", 0),
}
}
```
---
rust/tvm-ffi-macros/src/lib.rs | 13 +++
rust/tvm-ffi-macros/src/match_any.rs | 161 +++++++++++++++++++++++++++++++++++
rust/tvm-ffi-macros/src/utils.rs | 5 +-
rust/tvm-ffi/src/any.rs | 7 ++
rust/tvm-ffi/src/lib.rs | 1 +
rust/tvm-ffi/tests/test_match_any.rs | 61 +++++++++++++
6 files changed, 245 insertions(+), 3 deletions(-)
diff --git a/rust/tvm-ffi-macros/src/lib.rs b/rust/tvm-ffi-macros/src/lib.rs
index 64fe3f18..aeffad1f 100644
--- a/rust/tvm-ffi-macros/src/lib.rs
+++ b/rust/tvm-ffi-macros/src/lib.rs
@@ -20,9 +20,22 @@
use proc_macro::TokenStream;
use proc_macro_error::proc_macro_error;
+mod match_any;
mod object_macros;
mod utils;
+/// Match object-backed values carried by an Any-compatible scrutinee.
+///
+/// The scrutinee may be an owned object handle, `Any`, or `AnyView`. Convert
an
+/// already-borrowed object handle to `AnyView` before invoking the macro.
+///
+/// Non-object values skip the typed patterns and use the `_` fallback.
+#[proc_macro_error]
+#[proc_macro]
+pub fn match_any(input: TokenStream) -> TokenStream {
+ match_any::expand(input)
+}
+
#[proc_macro_error]
#[proc_macro_derive(Object, attributes(type_key, type_index))]
pub fn derive_object(input: TokenStream) -> TokenStream {
diff --git a/rust/tvm-ffi-macros/src/match_any.rs
b/rust/tvm-ffi-macros/src/match_any.rs
new file mode 100644
index 00000000..14839210
--- /dev/null
+++ b/rust/tvm-ffi-macros/src/match_any.rs
@@ -0,0 +1,161 @@
+/*
+ * 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_macro2::{Ident, Span, TokenStream};
+use quote::quote;
+use syn::parse::{Parse, ParseStream};
+use syn::{braced, parenthesized, Expr, Pat, Path, Result, Token};
+
+use crate::utils::get_tvm_ffi_crate;
+
+struct MatchAnyInput {
+ scrutinee: Expr,
+ arms: Vec<TypedArm>,
+ fallback: Expr,
+}
+
+struct TypedArm {
+ matcher: Path,
+ binding: Pat,
+ guard: Option<Expr>,
+ body: Expr,
+}
+
+impl Parse for MatchAnyInput {
+ fn parse(input: ParseStream<'_>) -> Result<Self> {
+ let scrutinee = input.call(Expr::parse_without_eager_brace)?;
+ let content;
+ braced!(content in input);
+
+ let mut arms = Vec::new();
+ let mut fallback = None;
+ while !content.is_empty() {
+ if fallback.is_some() {
+ return Err(content.error("the `_` fallback must be the final
arm"));
+ }
+
+ if content.peek(Token![_]) {
+ content.parse::<Token![_]>()?;
+ if content.peek(Token![if]) {
+ return Err(content.error("the `_` fallback cannot have a
guard"));
+ }
+ content.parse::<Token![=>]>()?;
+ fallback = Some(content.parse::<Expr>()?);
+ } else {
+ let matcher = content.parse::<Path>()?;
+ let binding_content;
+ parenthesized!(binding_content in content);
+ let binding = binding_content.parse::<Pat>()?;
+ if !binding_content.is_empty() {
+ return Err(binding_content.error("expected one binding
pattern"));
+ }
+ let guard = if content.peek(Token![if]) {
+ content.parse::<Token![if]>()?;
+ Some(content.parse::<Expr>()?)
+ } else {
+ None
+ };
+ content.parse::<Token![=>]>()?;
+ let body = content.parse::<Expr>()?;
+ arms.push(TypedArm {
+ matcher,
+ binding,
+ guard,
+ body,
+ });
+ }
+
+ if content.peek(Token![,]) {
+ content.parse::<Token![,]>()?;
+ } else if !content.is_empty() {
+ return Err(content.error("expected `,` between match_any!
arms"));
+ }
+ }
+
+ let fallback = fallback
+ .ok_or_else(|| content.error("match_any! requires a final `_`
fallback arm"))?;
+ Ok(Self {
+ scrutinee,
+ arms,
+ fallback,
+ })
+ }
+}
+
+pub fn expand(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+ let input = syn::parse_macro_input!(input as MatchAnyInput);
+ expand_match_any(input).into()
+}
+
+fn expand_match_any(input: MatchAnyInput) -> TokenStream {
+ let tvm_ffi = get_tvm_ffi_crate();
+ 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))
+ };
+
+ quote! {
+ match ::core::convert::TryInto::<#matcher>::try_into(#view) {
+ #matched => { #body },
+ #rejected => {
+ ::core::mem::drop(#rejected);
+ #next
+ },
+ }
+ }
+ });
+
+ 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
+ }
+ }
+ }
+}
diff --git a/rust/tvm-ffi-macros/src/utils.rs b/rust/tvm-ffi-macros/src/utils.rs
index da86534f..b1ebd5f1 100644
--- a/rust/tvm-ffi-macros/src/utils.rs
+++ b/rust/tvm-ffi-macros/src/utils.rs
@@ -20,10 +20,9 @@ use proc_macro2::TokenStream;
use quote::quote;
use std::env;
-/// Get the tvm-rt crate name
-/// \return The tvm-rt crate name
+/// Return the path used to reference the `tvm-ffi` crate in generated code.
pub(crate) fn get_tvm_ffi_crate() -> TokenStream {
- if env::var("CARGO_PKG_NAME").unwrap() == "tvm-ffi" {
+ if env::var("CARGO_CRATE_NAME").unwrap() == "tvm_ffi" {
quote!(crate)
} else {
quote!(tvm_ffi)
diff --git a/rust/tvm-ffi/src/any.rs b/rust/tvm-ffi/src/any.rs
index ecf8b9ea..0c4c8476 100644
--- a/rust/tvm-ffi/src/any.rs
+++ b/rust/tvm-ffi/src/any.rs
@@ -103,6 +103,13 @@ impl<'a, T: AnyCompatible> From<&'a T> for AnyView<'a> {
}
}
+impl<'a> From<&AnyView<'a>> for AnyView<'a> {
+ #[inline]
+ fn from(value: &AnyView<'a>) -> Self {
+ *value
+ }
+}
+
impl Default for AnyView<'_> {
fn default() -> Self {
Self::new()
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index 72b8d160..d5e2b2e8 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -49,6 +49,7 @@ pub use crate::object::{Object, ObjectArc, ObjectCore,
ObjectCoreWithExtraItems,
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_sys::TVMFFITypeIndex as TypeIndex;
pub use tvm_ffi_sys::{
diff --git a/rust/tvm-ffi/tests/test_match_any.rs
b/rust/tvm-ffi/tests/test_match_any.rs
new file mode 100644
index 00000000..de40b482
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_match_any.rs
@@ -0,0 +1,61 @@
+/*
+ * 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::{match_any, Any, AnyView, Array, Map, Shape, Tensor};
+
+#[test]
+fn matches_concrete_object_containers_in_source_order() {
+ fn classify(expr: Any) -> (&'static str, usize) {
+ match_any! {
+ expr {
+ Tensor(tensor)
+ if tensor.shape().len() == 2 => ("matrix",
tensor.shape().len()),
+ Tensor(tensor) => ("tensor", tensor.shape().len()),
+ Shape(shape) => ("shape", shape.len()),
+ Array::<i64>(array) => ("array", array.len()),
+ _ => ("unsupported", 0),
+ }
+ }
+ }
+
+ let matrix = Tensor::from_slice(&[0_f32; 6], &[2, 3]).unwrap();
+ let volume = Tensor::from_slice(&[0_f32; 24], &[2, 3, 4]).unwrap();
+ let shape = Shape::from([2_i64, 3, 4, 5]);
+ let array = [1_i64, 2, 3].into_iter().collect::<Array<i64>>();
+
+ assert_eq!(classify(Any::from(matrix)), ("matrix", 2));
+ assert_eq!(classify(Any::from(volume)), ("tensor", 3));
+ assert_eq!(classify(Any::from(shape)), ("shape", 4));
+ assert_eq!(classify(Any::from(array)), ("array", 3));
+ assert_eq!(
+ classify(Any::from(Map::<i64, i64>::default())),
+ ("unsupported", 0)
+ );
+ assert_eq!(classify(Any::from(1_i64)), ("unsupported", 0));
+
+ let tensor = Tensor::from_slice(&[0_f32; 6], &[2, 3]).unwrap();
+ let view = AnyView::from(&tensor);
+ let matched_view = match_any! {
+ view {
+ Tensor(tensor) => ("tensor", tensor.shape().len()),
+ _ => ("unsupported", 0),
+ }
+ };
+ assert_eq!(matched_view, ("tensor", 2));
+}
