From e25e3a8f72cdca005aeeb40817c484ded8024c0a Mon Sep 17 00:00:00 2001
From: Robin Freyler <robin.freyler@gmail.com>
Date: Wed, 27 Nov 2024 17:17:23 +0100
Subject: [PATCH] Improve C-API `prefix_symbol` attribute macro (#1317)

* restructure prefix_symbol attribute macro impl

* prefix_symbols: check that attributes is empty

* no longer attribute fixed wasmi APIs

* no longer prefix non-Wasm spec functions

* clean-up macro implementation

* apply clippy suggestions
---
 crates/c_api/macro/lib.rs | 72 +++++++++++++++++++++++++--------------
 crates/c_api/src/store.rs |  6 ----
 2 files changed, 46 insertions(+), 32 deletions(-)

diff --git a/crates/c_api/macro/lib.rs b/crates/c_api/macro/lib.rs
index 075e0ca66e..d2903748ca 100644
--- a/crates/c_api/macro/lib.rs
+++ b/crates/c_api/macro/lib.rs
@@ -159,38 +159,58 @@ pub fn declare_ref(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     .into()
 }
 
+macro_rules! bail {
+    ($message:literal) => {{
+        return ::core::result::Result::Err(Error($message.into()));
+    }};
+}
+
+/// An error with its error message.
+struct Error(String);
+
+impl Error {
+    /// Converts the [`Error`] into a `compile_error!` token stream.
+    fn to_compile_error(&self) -> proc_macro::TokenStream {
+        let message = &self.0;
+        quote! { ::core::compile_error!(#message) }.into()
+    }
+}
+
+/// Applied on Rust `fn` items from the Wasm spec.
+///
+/// Annotates the given function with `#[export_name = $func_name]`
+/// where `$func_name` is the name of the given function.
 #[proc_macro_attribute]
 pub fn prefix_symbol(
-    _attributes: proc_macro::TokenStream,
+    attributes: proc_macro::TokenStream,
     input: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
-    let mut stream = TokenStream::from(input.clone()).into_iter();
-
-    let mut fn_name = None;
-
-    while let Some(token) = stream.next() {
-        match token {
-            TokenTree::Ident(ref ident) if *ident == "fn" => {
-                if let Some(TokenTree::Ident(ident_name)) = stream.next() {
-                    fn_name = Some(ident_name.to_string());
-                    break;
-                }
-            }
-            _ => continue,
-        }
+    match prefix_symbol_impl(attributes.into(), input.into()) {
+        Ok(result) => result.into(),
+        Err(error) => error.to_compile_error(),
     }
+}
 
-    if fn_name.is_none() {
-        panic!("expected a valid Rust function definition, but it does not appear in: {input:?}");
+fn prefix_symbol_impl(attributes: TokenStream, input: TokenStream) -> Result<TokenStream, Error> {
+    if !attributes.is_empty() {
+        bail!("err(prefix_symbol): attributes must be empty")
     }
-
-    let prefixed_fn_name = format!("wasmi_{}", fn_name.unwrap());
-
-    let mut attr: proc_macro::TokenStream = quote! {
-        #[export_name = #prefixed_fn_name]
+    let mut stream = input.clone().into_iter();
+    let fn_token = stream.find(|tt| matches!(tt, TokenTree::Ident(ref ident) if *ident == "fn"));
+    if fn_token.is_none() {
+        bail!("can only apply on `fn` items")
     }
-    .into();
-    attr.extend(input);
-
-    attr
+    let Some(TokenTree::Ident(fn_ident)) = stream.next() else {
+        bail!("function name must follow `fn` keyword")
+    };
+    let fn_name = fn_ident.to_string();
+    if !fn_name.starts_with("wasm_") {
+        // No prefix needed since the function is not a part of the Wasm spec.
+        return Ok(input);
+    }
+    let prefixed_fn_name = format!("wasmi_{}", fn_name);
+    Ok(quote! {
+        #[export_name = #prefixed_fn_name]
+        #input
+    })
 }
diff --git a/crates/c_api/src/store.rs b/crates/c_api/src/store.rs
index 21293e93d9..0b1f68d3bf 100644
--- a/crates/c_api/src/store.rs
+++ b/crates/c_api/src/store.rs
@@ -98,7 +98,6 @@ pub struct WasmiStoreData {
 ///
 /// Wraps [`<wasmi::Store<()>>::new`](wasmi::Store::new).
 #[no_mangle]
-#[cfg_attr(feature = "prefix-symbols", wasmi_c_api_macros::prefix_symbol)]
 pub extern "C" fn wasmi_store_new(
     engine: &wasm_engine_t,
     data: *mut ffi::c_void,
@@ -122,7 +121,6 @@ pub extern "C" fn wasmi_store_new(
 ///
 /// It is the callers responsibility to provide a valid `self`.
 #[no_mangle]
-#[cfg_attr(feature = "prefix-symbols", wasmi_c_api_macros::prefix_symbol)]
 pub extern "C" fn wasmi_store_context(
     store: &mut wasmi_store_t,
 ) -> StoreContextMut<'_, WasmiStoreData> {
@@ -131,7 +129,6 @@ pub extern "C" fn wasmi_store_context(
 
 /// Returns a pointer to the foreign data of the Wasmi store context.
 #[no_mangle]
-#[cfg_attr(feature = "prefix-symbols", wasmi_c_api_macros::prefix_symbol)]
 pub extern "C" fn wasmi_context_get_data(
     store: StoreContext<'_, WasmiStoreData>,
 ) -> *mut ffi::c_void {
@@ -140,7 +137,6 @@ pub extern "C" fn wasmi_context_get_data(
 
 /// Sets the foreign data of the Wasmi store context.
 #[no_mangle]
-#[cfg_attr(feature = "prefix-symbols", wasmi_c_api_macros::prefix_symbol)]
 pub extern "C" fn wasmi_context_set_data(
     mut store: StoreContextMut<'_, WasmiStoreData>,
     data: *mut ffi::c_void,
@@ -156,7 +152,6 @@ pub extern "C" fn wasmi_context_set_data(
 ///
 /// If [`Store::get_fuel`] errors.
 #[no_mangle]
-#[cfg_attr(feature = "prefix-symbols", wasmi_c_api_macros::prefix_symbol)]
 pub extern "C" fn wasmi_context_get_fuel(
     store: StoreContext<'_, WasmiStoreData>,
     fuel: &mut u64,
@@ -174,7 +169,6 @@ pub extern "C" fn wasmi_context_get_fuel(
 ///
 /// If [`Store::set_fuel`] errors.
 #[no_mangle]
-#[cfg_attr(feature = "prefix-symbols", wasmi_c_api_macros::prefix_symbol)]
 pub extern "C" fn wasmi_context_set_fuel(
     mut store: StoreContextMut<'_, WasmiStoreData>,
     fuel: u64,