Use function type traits to improve registration ergonomics

Signed-off-by: Jorge Prendes <jorge.prendes@gmail.com>

Author: jprendes

src/hyperlight_guest/Cargo.toml

@@ -14,7 +14,7 @@ Provides only the essential building blocks for interacting with the host enviro
 [dependencies]
 anyhow = { version = "1.0.99", default-features = false }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
-hyperlight-common = { workspace = true }
+hyperlight-common = { workspace = true, default-features = false }
 hyperlight-guest-tracing = { workspace = true, default-features = false }
 flatbuffers = { version= "25.2.10", default-features = false }
 

src/hyperlight_guest/src/error.rs

@@ -15,9 +15,10 @@ limitations under the License.
 */
 
 use alloc::format;
-use alloc::string::String;
+use alloc::string::{String, ToString as _};
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
+use hyperlight_common::func::Error as FuncError;
 use {anyhow, serde_json};
 
 pub type Result<T> = core::result::Result<T, HyperlightGuestError>;
@@ -51,3 +52,30 @@ impl From<serde_json::Error> for HyperlightGuestError {
         }
     }
 }
+
+impl From<FuncError> for HyperlightGuestError {
+    fn from(e: FuncError) -> Self {
+        match e {
+            FuncError::ParameterValueConversionFailure(..) => HyperlightGuestError::new(
+                ErrorCode::GuestFunctionParameterTypeMismatch,
+                e.to_string(),
+            ),
+            FuncError::ReturnValueConversionFailure(..) => HyperlightGuestError::new(
+                ErrorCode::GuestFunctionParameterTypeMismatch,
+                e.to_string(),
+            ),
+            FuncError::UnexpectedNoOfArguments(..) => HyperlightGuestError::new(
+                ErrorCode::GuestFunctionIncorrecNoOfParameters,
+                e.to_string(),
+            ),
+            FuncError::UnexpectedParameterValueType(..) => HyperlightGuestError::new(
+                ErrorCode::GuestFunctionParameterTypeMismatch,
+                e.to_string(),
+            ),
+            FuncError::UnexpectedReturnValueType(..) => HyperlightGuestError::new(
+                ErrorCode::GuestFunctionParameterTypeMismatch,
+                e.to_string(),
+            ),
+        }
+    }
+}

src/hyperlight_guest_bin/src/guest_function/definition.rs

@@ -18,8 +18,14 @@ use alloc::format;
 use alloc::string::String;
 use alloc::vec::Vec;
 
+use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall;
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterType, ReturnType};
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
+use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
+use hyperlight_common::for_each_tuple;
+use hyperlight_common::func::{
+    Function, ParameterTuple, ResultType, ReturnValue, SupportedReturnType,
+};
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 
 /// The definition of a function exposed from the guest to the host
@@ -35,6 +41,101 @@ pub struct GuestFunctionDefinition {
     pub function_pointer: usize,
 }
 
+/// Trait for functions that can be converted to a `fn(&FunctionCall) -> Result<Vec<u8>>`
+#[doc(hidden)]
+pub trait IntoGuestFunction<Output, Args>
+where
+    Self: Function<Output, Args, HyperlightGuestError>,
+    Self: Copy + 'static,
+    Output: SupportedReturnType,
+    Args: ParameterTuple,
+{
+    #[doc(hidden)]
+    const ASSERT_ZERO_SIZED: ();
+
+    /// Convert the function into a `fn(&FunctionCall) -> Result<Vec<u8>>`
+    fn into_guest_function(self) -> fn(&FunctionCall) -> Result<Vec<u8>>;
+}
+
+/// Trait for functions that can be converted to a `GuestFunctionDefinition`
+pub trait AsGuestFunctionDefinition<Output, Args>
+where
+    Self: Function<Output, Args, HyperlightGuestError>,
+    Self: IntoGuestFunction<Output, Args>,
+    Output: SupportedReturnType,
+    Args: ParameterTuple,
+{
+    /// Get the `GuestFunctionDefinition` for this function
+    fn as_guest_function_definition(&self, name: impl Into<String>) -> GuestFunctionDefinition;
+}
+
+fn into_flatbuffer_result(value: ReturnValue) -> Vec<u8> {
+    match value {
+        ReturnValue::Void(()) => get_flatbuffer_result(()),
+        ReturnValue::Int(i) => get_flatbuffer_result(i),
+        ReturnValue::UInt(u) => get_flatbuffer_result(u),
+        ReturnValue::Long(l) => get_flatbuffer_result(l),
+        ReturnValue::ULong(ul) => get_flatbuffer_result(ul),
+        ReturnValue::Float(f) => get_flatbuffer_result(f),
+        ReturnValue::Double(d) => get_flatbuffer_result(d),
+        ReturnValue::Bool(b) => get_flatbuffer_result(b),
+        ReturnValue::String(s) => get_flatbuffer_result(s.as_str()),
+        ReturnValue::VecBytes(v) => get_flatbuffer_result(v.as_slice()),
+    }
+}
+
+macro_rules! impl_host_function {
+    ([$N:expr] ($($p:ident: $P:ident),*)) => {
+        impl<F, R, $($P),*> IntoGuestFunction<R::ReturnType, ($($P,)*)> for F
+        where
+            F: Fn($($P),*) -> R,
+            F: Function<R::ReturnType, ($($P,)*), HyperlightGuestError>,
+            F: Copy + 'static, // Copy implies that F has no Drop impl
+            ($($P,)*): ParameterTuple,
+            R: ResultType<HyperlightGuestError>,
+        {
+            #[doc(hidden)]
+            const ASSERT_ZERO_SIZED: () = const {
+                assert!(core::mem::size_of::<Self>() == 0)
+            };
+
+            fn into_guest_function(self) -> fn(&FunctionCall) -> Result<Vec<u8>> {
+                |fc: &FunctionCall| {
+                    // SAFETY: This is safe because of the safety comment on the function.
+                    let this = unsafe { core::mem::zeroed::<F>() };
+                    let params = fc.parameters.clone().unwrap_or_default();
+                    let params = <($($P,)*) as ParameterTuple>::from_value(params)?;
+                    let result = Function::<R::ReturnType, ($($P,)*), HyperlightGuestError>::call(&this, params)?;
+                    Ok(into_flatbuffer_result(result.into_value()))
+                }
+            }
+        }
+    };
+}
+
+impl<F, Args, Output> AsGuestFunctionDefinition<Output, Args> for F
+where
+    F: IntoGuestFunction<Output, Args>,
+    Args: ParameterTuple,
+    Output: SupportedReturnType,
+{
+    fn as_guest_function_definition(&self, name: impl Into<String>) -> GuestFunctionDefinition {
+        let parameter_types = Args::TYPE.to_vec();
+        let return_type = Output::TYPE;
+        let function_pointer = self.into_guest_function();
+        let function_pointer = function_pointer as usize;
+
+        GuestFunctionDefinition {
+            function_name: name.into(),
+            parameter_types,
+            return_type,
+            function_pointer,
+        }
+    }
+}
+
+for_each_tuple!(impl_host_function);
+
 impl GuestFunctionDefinition {
     /// Create a new `GuestFunctionDefinition`.
     pub fn new(
@@ -51,6 +152,19 @@ impl GuestFunctionDefinition {
         }
     }
 
+    /// Create a new `GuestFunctionDefinition` from a function that implements
+    /// `AsGuestFunctionDefinition`.
+    pub fn from_fn<Output, Args>(
+        function_name: String,
+        function: impl AsGuestFunctionDefinition<Output, Args>,
+    ) -> Self
+    where
+        Args: ParameterTuple,
+        Output: SupportedReturnType,
+    {
+        function.as_guest_function_definition(function_name)
+    }
+
     /// Verify that `self` has same signature as the provided `parameter_types`.
     pub fn verify_parameters(&self, parameter_types: &[ParameterType]) -> Result<()> {
         // Verify that the function does not have more than `MAX_PARAMETERS` parameters.

src/hyperlight_guest_bin/src/guest_function/register.rs

@@ -17,8 +17,11 @@ limitations under the License.
 use alloc::collections::BTreeMap;
 use alloc::string::String;
 
+use hyperlight_common::func::{ParameterTuple, SupportedReturnType};
+
 use super::definition::GuestFunctionDefinition;
 use crate::REGISTERED_GUEST_FUNCTIONS;
+use crate::guest_function::definition::AsGuestFunctionDefinition;
 
 /// Represents the functions that the guest exposes to the host.
 #[derive(Debug, Default, Clone)]
@@ -47,6 +50,18 @@ impl GuestFunctionRegister {
             .insert(guest_function.function_name.clone(), guest_function)
     }
 
+    pub fn register_fn<Output, Args>(
+        &mut self,
+        name: impl Into<String>,
+        f: impl AsGuestFunctionDefinition<Output, Args>,
+    ) where
+        Args: ParameterTuple,
+        Output: SupportedReturnType,
+    {
+        let gfd = f.as_guest_function_definition(name);
+        self.register(gfd);
+    }
+
     /// Gets a `GuestFunctionDefinition` by its `name` field.
     pub fn get(&self, function_name: &str) -> Option<&GuestFunctionDefinition> {
         self.guest_functions.get(function_name)
@@ -62,3 +77,19 @@ pub fn register_function(function_definition: GuestFunctionDefinition) {
         gfd.register(function_definition);
     }
 }
+
+pub fn register_fn<Output, Args>(
+    name: impl Into<String>,
+    f: impl AsGuestFunctionDefinition<Output, Args>,
+) where
+    Args: ParameterTuple,
+    Output: SupportedReturnType,
+{
+    unsafe {
+        // This is currently safe, because we are single threaded, but we
+        // should find a better way to do this, see issue #808
+        #[allow(static_mut_refs)]
+        let gfd = &mut REGISTERED_GUEST_FUNCTIONS;
+        gfd.register_fn(name, f);
+    }
+}

src/hyperlight_guest_bin/src/host_comm.rs

@@ -26,6 +26,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
 use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
+use hyperlight_common::func::{ParameterTuple, SupportedReturnType};
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 
 const BUFFER_SIZE: usize = 1000;
@@ -45,6 +46,13 @@ where
     handle.call_host_function::<T>(function_name, parameters, return_type)
 }
 
+pub fn call_host<T>(function_name: impl AsRef<str>, args: impl ParameterTuple) -> Result<T>
+where
+    T: SupportedReturnType + TryFrom<ReturnValue>,
+{
+    call_host_function::<T>(function_name.as_ref(), Some(args.into_value()), T::TYPE)
+}
+
 pub fn call_host_function_without_returning_result(
     function_name: &str,
     parameters: Option<Vec<ParameterValue>>,