[Rust][Feat] Add match_any! for object-backed values

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 {

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
+            }
+        }
+    }
+}

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)

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()

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::{

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));
+}