hyperlight_guest_macro/

lib.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 The Hyperlight Authors.

use proc_macro::TokenStream;
use proc_macro_crate::{FoundCrate, crate_name};
use quote::quote;
use syn::parse::{Error, Parse, ParseStream, Result};
use syn::spanned::Spanned as _;
use syn::{ForeignItemFn, ItemFn, LitStr, Pat, parse_macro_input};

/// Represents the optional name argument for the guest_function and host_function macros.
enum NameArg {
    None,
    Name(LitStr),
}

impl Parse for NameArg {
    fn parse(input: ParseStream) -> Result<Self> {
        // accepts either nothing or a single string literal
        // anything else is an error
        if input.is_empty() {
            return Ok(NameArg::None);
        }
        let name: LitStr = input.parse()?;
        if !input.is_empty() {
            return Err(Error::new(input.span(), "expected a single identifier"));
        }
        Ok(NameArg::Name(name))
    }
}

/// Attribute macro to mark a function as a guest function.
/// This will register the function so that it can be called by the host.
///
/// If a name is provided as an argument, that name will be used to register the function.
/// Otherwise, the function's identifier will be used.
///
/// The function arguments must be supported parameter types, and the return type must be
/// a supported return type or a `Result<T, HyperlightGuestError>` with T being a supported
/// return type.
///
/// # Note
/// The function will be registered with the host at program initialization regardless of
/// the visibility modifier used (e.g., `pub`, `pub(crate)`, etc.).
/// This means that a private functions can be called by the host from beyond its normal
/// visibility scope.
///
/// # Example
/// ```ignore
/// use hyperlight_guest_bin::guest_function;
/// #[guest_function]
/// fn my_guest_function(arg1: i32, arg2: String) -> i32 {
///     arg1 + arg2.len() as i32
/// }
/// ```
///
/// or with a custom name:
/// ```ignore
/// use hyperlight_guest_bin::guest_function;
/// #[guest_function("custom_name")]
/// fn my_guest_function(arg1: i32, arg2: String) -> i32 {
///     arg1 + arg2.len() as i32
/// }
/// ```
///
/// or with a Result return type:
/// ```ignore
/// use hyperlight_guest_bin::guest_function;
/// use hyperlight_guest::bail;
/// #[guest_function]
/// fn my_guest_function(arg1: i32, arg2: String) -> Result<i32, HyperlightGuestError> {
///     bail!("An error occurred");
/// }
/// ```
#[proc_macro_attribute]
pub fn guest_function(attr: TokenStream, item: TokenStream) -> TokenStream {
    // Obtain the crate name for hyperlight-guest-bin
    let crate_name =
        crate_name("hyperlight-guest-bin").expect("hyperlight-guest-bin must be a dependency");
    let crate_name = match crate_name {
        FoundCrate::Itself => quote! {crate},
        FoundCrate::Name(name) => {
            let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
            quote! {::#ident}
        }
    };

    // Parse the function definition that we will be working with, and
    // early return if parsing as `ItemFn` fails.
    let fn_declaration = parse_macro_input!(item as ItemFn);

    // Obtain the name of the function being decorated.
    let ident = fn_declaration.sig.ident.clone();

    // Determine the name used to register the function, either
    // the provided name or the function's identifier.
    let exported_name = match parse_macro_input!(attr as NameArg) {
        NameArg::None => quote! { stringify!(#ident) },
        NameArg::Name(name) => quote! { #name },
    };

    // Small sanity checks to improve error messages.
    // These checks are not strictly necessary, as the generated code
    // would fail to compile anyway (due to the trait bounds of `register_fn`),
    // but they provide better feedback to the user of the macro.

    // Check that there are no receiver arguments (i.e., `self`, `&self`, `Box<Self>`, etc).
    if let Some(syn::FnArg::Receiver(arg)) = fn_declaration.sig.inputs.first() {
        return Error::new(
            arg.span(),
            "Receiver (self) argument is not allowed in guest functions",
        )
        .to_compile_error()
        .into();
    }

    // Check that the function is not async.
    if fn_declaration.sig.asyncness.is_some() {
        return Error::new(
            fn_declaration.sig.asyncness.span(),
            "Async functions are not allowed in guest functions",
        )
        .to_compile_error()
        .into();
    }

    // The generated code will replace the decorated code, so we need to
    // include the original function declaration in the output.
    let output = quote! {
        #fn_declaration

        const _: () = {
            // Add the function registration in the GUEST_FUNCTION_INIT distributed slice
            // so that it can be registered at program initialization
            #[#crate_name::__private::linkme::distributed_slice(#crate_name::__private::GUEST_FUNCTION_INIT)]
            #[linkme(crate = #crate_name::__private::linkme)]
            static REGISTRATION: fn() = || {
                #crate_name::guest_function::register::register_fn(#exported_name, #ident);
            };
        };
    };

    output.into()
}

/// Attribute macro to mark a function as the main entry point for the guest.
/// This will generate a function that is called by the host at program initialization.
///
/// # Example
/// ```ignore
/// use hyperlight_guest_bin::main;
/// #[main]
/// fn main() {
///     // do some initialization work here, e.g., initialize global state, etc.
/// }
/// ```
#[proc_macro_attribute]
pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse the function definition that we will be working with, and
    // early return if parsing as `ItemFn` fails.
    let fn_declaration = parse_macro_input!(item as ItemFn);

    // Obtain the name of the function being decorated.
    let ident = fn_declaration.sig.ident.clone();

    // The generated code will replace the decorated code, so we need to
    // include the original function declaration in the output.
    let output = quote! {
        #fn_declaration

        const _: () = {
            mod wrapper {
                #[unsafe(no_mangle)]
                pub extern "C" fn hyperlight_main() {
                    super::#ident()
                }
            }
        };
    };

    output.into()
}

/// Attribute macro to mark a function as the dispatch function for the guest.
/// This is the function that will be called by the host when a function call is made
/// to a function that is not registered with the host.
///
/// # Example
/// ```ignore
/// use hyperlight_guest_bin::dispatch;
/// use hyperlight_guest::error::Result;
/// use hyperlight_guest::bail;
/// use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall;
/// use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
/// #[dispatch]
/// fn dispatch(fc: FunctionCall) -> Result<Vec<u8>> {
///     let name = &fc.function_name;
///     if name == "greet" {
///         return Ok(get_flatbuffer_result("Hello, world!"));
///     }
///     bail!("Unknown function: {name}");
/// }
/// ```
#[proc_macro_attribute]
pub fn dispatch(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // Obtain the crate name for hyperlight-guest-bin
    let crate_name =
        crate_name("hyperlight-guest-bin").expect("hyperlight-guest-bin must be a dependency");
    let crate_name = match crate_name {
        FoundCrate::Itself => quote! {crate},
        FoundCrate::Name(name) => {
            let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
            quote! {::#ident}
        }
    };

    // Parse the function definition that we will be working with, and
    // early return if parsing as `ItemFn` fails.
    let fn_declaration = parse_macro_input!(item as ItemFn);

    // Obtain the name of the function being decorated.
    let ident = fn_declaration.sig.ident.clone();

    // The generated code will replace the decorated code, so we need to
    // include the original function declaration in the output.
    let output = quote! {
        #fn_declaration

        const _: () = {
            mod wrapper {
                use #crate_name::__private::{FunctionCall, HyperlightGuestError, Vec};
                #[unsafe(no_mangle)]
                pub fn guest_dispatch_function(function_call: FunctionCall) -> ::core::result::Result<Vec<u8>, HyperlightGuestError> {
                    super::#ident(function_call)
                }
            }
        };
    };

    output.into()
}

/// Attribute macro to mark a function as a host function.
/// This will generate a function that calls the host function with the same name.
///
/// If a name is provided as an argument, that name will be used to call the host function.
/// Otherwise, the function's identifier will be used.
///
/// The function arguments must be supported parameter types, and the return type must be
/// a supported return type or a `Result<T, HyperlightGuestError>` with T being a supported
/// return type.
///
/// # Panic
/// If the return type is not a Result, the generated function will panic if the host function
/// returns an error.
///
/// # Example
/// ```ignore
/// use hyperlight_guest_bin::host_function;
/// #[host_function]
/// fn my_host_function(arg1: i32, arg2: String) -> i32;
/// ```
///
/// or with a custom name:
/// ```ignore
/// use hyperlight_guest_bin::host_function;
/// #[host_function("custom_name")]
/// fn my_host_function(arg1: i32, arg2: String) -> i32;
/// ```
///
/// or with a Result return type:
/// ```ignore
/// use hyperlight_guest_bin::host_function;
/// use hyperlight_guest::error::HyperlightGuestError;
/// #[host_function]
/// fn my_host_function(arg1: i32, arg2: String) -> Result<i32, HyperlightGuestError>;
/// ```
#[proc_macro_attribute]
pub fn host_function(attr: TokenStream, item: TokenStream) -> TokenStream {
    // Obtain the crate name for hyperlight-guest-bin
    let crate_name =
        crate_name("hyperlight-guest-bin").expect("hyperlight-guest-bin must be a dependency");
    let crate_name = match crate_name {
        FoundCrate::Itself => quote! {crate},
        FoundCrate::Name(name) => {
            let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
            quote! {::#ident}
        }
    };

    // Parse the function declaration that we will be working with, and
    // early return if parsing as `ForeignItemFn` fails.
    // A function declaration without a body is a foreign item function, as that's what
    // you would use when declaring an FFI function.
    let fn_declaration = parse_macro_input!(item as ForeignItemFn);

    // Destructure the foreign item function to get its components.
    let ForeignItemFn {
        attrs,
        vis,
        sig,
        semi_token: _,
        modifiers: _,
    } = fn_declaration;

    // Obtain the name of the function being decorated.
    let ident = sig.ident.clone();

    // Determine the name used to call the host function, either
    // the provided name or the function's identifier.
    let exported_name = match parse_macro_input!(attr as NameArg) {
        NameArg::None => quote! { stringify!(#ident) },
        NameArg::Name(name) => quote! { #name },
    };

    // Build the list of argument identifiers to pass to the call_host function.
    // While doing that, also do some sanity checks to improve error messages.
    // These checks are not strictly necessary, as the generated code would fail
    // to compile anyway due to either:
    // * the trait bounds of `call_host`
    // * the generated code having invalid syntax
    // but they provide better feedback to the user of the macro, especially in
    // the case of invalid syntax.
    let mut args = vec![];
    for arg in sig.inputs.iter() {
        match arg {
            // Reject receiver arguments (i.e., `self`, `&self`, `Box<Self>`, etc).
            syn::FnArg::Receiver(_) => {
                return Error::new(
                    arg.span(),
                    "Receiver (self) argument is not allowed in guest functions",
                )
                .to_compile_error()
                .into();
            }
            syn::FnArg::Typed(arg) => {
                // A typed argument: `name: Type`
                // Technically, the `name` part can be any pattern, e.g., destructuring patterns
                // like `(a, b): (i32, u64)`, but we only allow simple identifiers here
                // to keep things simple.

                // Reject anything that is not a simple identifier.
                let Pat::Ident(pat) = *arg.pat.clone() else {
                    return Error::new(
                        arg.span(),
                        "Only named arguments are allowed in host functions",
                    )
                    .to_compile_error()
                    .into();
                };

                // Reject any argument with attributes, e.g., `#[cfg(feature = "gdb")] name: Type`
                if !pat.attrs.is_empty() {
                    return Error::new(
                        arg.span(),
                        "Attributes are not allowed on host function arguments",
                    )
                    .to_compile_error()
                    .into();
                }

                // Reject any argument passed by reference
                if pat.by_ref.is_some() {
                    return Error::new(
                        arg.span(),
                        "By-ref arguments are not allowed in host functions",
                    )
                    .to_compile_error()
                    .into();
                }

                // Reject any mutable argument, e.g., `mut name: Type`
                if pat.mutability.is_some() {
                    return Error::new(
                        arg.span(),
                        "Mutable arguments are not allowed in host functions",
                    )
                    .to_compile_error()
                    .into();
                }

                // Reject any sub-patterns
                if pat.subpat.is_some() {
                    return Error::new(
                        arg.span(),
                        "Sub-patterns are not allowed in host functions",
                    )
                    .to_compile_error()
                    .into();
                }

                let ident = pat.ident.clone();

                // All checks passed, add the identifier to the argument list.
                args.push(quote! { #ident });
            }
        }
    }

    // Determine the return type of the function.
    // If the return type is not specified, it is `()`.
    let ret: proc_macro2::TokenStream = match &sig.output {
        syn::ReturnType::Default => quote! { quote! { () } },
        syn::ReturnType::Type(_, ty) => {
            quote! { #ty }
        }
    };

    // Take the parts of the function declaration and generate a function definition
    // matching the provided declaration, but with a body that calls the host function.
    let output = quote! {
        #(#attrs)* #vis #sig {
            use #crate_name::__private::FromResult;
            use #crate_name::host_comm::call_host;
            <#ret as FromResult>::from_result(call_host(#exported_name, (#(#args,)*)))
        }
    };

    output.into()
}