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

rust/tvm-ffi-macros/src/lib.rs

@@ -20,9 +20,17 @@
 use proc_macro::TokenStream;
 use proc_macro_error::proc_macro_error;
 
+mod match_any;
 mod object_macros;
 mod utils;
 
+/// Match a TVM-FFI value against typed patterns in source order.
+#[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,146 @@
+/*
+ * 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 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 = input
+        .arms
+        .into_iter()
+        .rev()
+        .fold(quote!({ #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::option::Option::Some(#binding) if #guard)
+            } else {
+                quote!(::core::option::Option::Some(#binding))
+            };
+
+            quote! {
+                match <#matcher as #tvm_ffi::AnyPattern>::try_match(#view) {
+                    #matched => { #body },
+                    #rejected => {
+                        ::core::mem::drop(#rejected);
+                        #next
+                    },
+                }
+            }
+        });
+
+    quote! {
+        {
+            let #source = &(#scrutinee);
+            let #view = #tvm_ffi::AsAnyView::as_any_view(#source);
+            #dispatch
+        }
+    }
+}

rust/tvm-ffi-macros/src/utils.rs

@@ -23,7 +23,9 @@ use std::env;
 /// Get the tvm-rt crate name
 /// \return The tvm-rt crate name
 pub(crate) fn get_tvm_ffi_crate() -> TokenStream {
-    if env::var("CARGO_PKG_NAME").unwrap() == "tvm-ffi" {
+    if env::var("CARGO_PKG_NAME").unwrap() == "tvm-ffi"
+        && env::var("CARGO_CRATE_NAME").unwrap() == "tvm_ffi"
+    {
         quote!(crate)
     } else {
         quote!(tvm_ffi)

rust/tvm-ffi/src/lib.rs

@@ -26,6 +26,7 @@ pub mod extra;
 pub mod function;
 pub mod function_internal;
 pub mod macros;
+pub mod match_any;
 pub mod object;
 pub mod optional;
 pub mod string;
@@ -45,10 +46,12 @@ pub use crate::error::{
 };
 pub use crate::extra::module::Module;
 pub use crate::function::Function;
+pub use crate::match_any::{AnyPattern, AsAnyView};
 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_sys::TVMFFITypeIndex as TypeIndex;
 pub use tvm_ffi_sys::{

rust/tvm-ffi/src/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.
+ */
+
+//! Runtime traits used by [`crate::match_any!`].
+
+use crate::{Any, AnyCompatible, AnyView};
+
+/// A value that can expose its contents as an erased [`AnyView`].
+///
+/// [`Any`], [`AnyView`], and every [`AnyCompatible`] value implement this
+/// trait.
+pub trait AsAnyView {
+    /// Borrow this value as an [`AnyView`].
+    fn as_any_view(&self) -> AnyView<'_>;
+}
+
+impl<T: AnyCompatible> AsAnyView for T {
+    #[inline]
+    fn as_any_view(&self) -> AnyView<'_> {
+        AnyView::from(self)
+    }
+}
+
+impl AsAnyView for Any {
+    #[inline]
+    fn as_any_view(&self) -> AnyView<'_> {
+        AnyView::from(self)
+    }
+}
+
+impl AsAnyView for AnyView<'_> {
+    #[inline]
+    fn as_any_view(&self) -> AnyView<'_> {
+        *self
+    }
+}
+
+/// A typed pattern accepted by [`crate::match_any!`].
+pub trait AnyPattern {
+    /// The binding produced by a successful match.
+    type Bound<'a>;
+
+    /// Try to match an erased value.
+    fn try_match<'a>(value: AnyView<'a>) -> Option<Self::Bound<'a>>;
+}