Make get_flatbuffer_result generic

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_common/src/flatbuffer_wrappers/util.rs

@@ -16,141 +16,156 @@ limitations under the License.
 
 use alloc::vec::Vec;
 
-use flatbuffers::{FlatBufferBuilder, UnionWIPOffset, WIPOffset};
+use flatbuffers::FlatBufferBuilder;
 
 use crate::flatbuffers::hyperlight::generated::{
-    hldouble as Fbhldouble, hldoubleArgs as FbhldoubleArgs, hlfloat as Fbhlfloat,
-    hlfloatArgs as FbhlfloatArgs, hlint as Fbhlint, hlintArgs as FbhlintArgs, hllong as Fbhllong,
-    hllongArgs as FbhllongArgs, hlsizeprefixedbuffer as Fbhlsizeprefixedbuffer,
+    hlbool as Fbhlbool, hlboolArgs as FbhlboolArgs, hldouble as Fbhldouble,
+    hldoubleArgs as FbhldoubleArgs, hlfloat as Fbhlfloat, hlfloatArgs as FbhlfloatArgs,
+    hlint as Fbhlint, hlintArgs as FbhlintArgs, hllong as Fbhllong, hllongArgs as FbhllongArgs,
+    hlsizeprefixedbuffer as Fbhlsizeprefixedbuffer,
     hlsizeprefixedbufferArgs as FbhlsizeprefixedbufferArgs, hlstring as Fbhlstring,
     hlstringArgs as FbhlstringArgs, hluint as Fbhluint, hluintArgs as FbhluintArgs,
     hlulong as Fbhlulong, hlulongArgs as FbhlulongArgs, hlvoid as Fbhlvoid,
     hlvoidArgs as FbhlvoidArgs, FunctionCallResult as FbFunctionCallResult,
     FunctionCallResultArgs as FbFunctionCallResultArgs, ReturnValue as FbReturnValue,
 };
 
-pub fn get_flatbuffer_result_from_double(value: f64) -> Vec<u8> {
+/// Flatbuffer-encodes the given value
+pub fn get_flatbuffer_result<T: FlatbufferSerializable>(val: T) -> Vec<u8> {
     let mut builder = FlatBufferBuilder::new();
-    let hldouble = Fbhldouble::create(&mut builder, &FbhldoubleArgs { value });
+    let res = &T::serialize(&val, &mut builder);
+    let result_offset = FbFunctionCallResult::create(&mut builder, res);
 
-    let rt = FbReturnValue::hldouble;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hldouble.as_union_value());
+    builder.finish_size_prefixed(result_offset, None);
 
-    get_flatbuffer_result(&mut builder, rt, rv)
+    builder.finished_data().to_vec()
 }
 
-pub fn get_flatbuffer_result_from_float(value: f32) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-    let hlfloat = Fbhlfloat::create(&mut builder, &FbhlfloatArgs { value });
-
-    let rt = FbReturnValue::hlfloat;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hlfloat.as_union_value());
-
-    get_flatbuffer_result(&mut builder, rt, rv)
+pub trait FlatbufferSerializable {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs;
 }
 
-pub fn get_flatbuffer_result_from_int(value: i32) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-    let hlint = Fbhlint::create(&mut builder, &FbhlintArgs { value });
-
-    let rt = FbReturnValue::hlint;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hlint.as_union_value());
+/// Implementations for basic types below
 
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for () {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(Fbhlvoid::create(builder, &FbhlvoidArgs {}).as_union_value()),
+            return_value_type: FbReturnValue::hlvoid,
+        }
+    }
 }
 
-pub fn get_flatbuffer_result_from_uint(value: u32) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-    let hluint = Fbhluint::create(&mut builder, &FbhluintArgs { value });
-
-    let rt = FbReturnValue::hluint;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hluint.as_union_value());
-
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for &str {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        let string_offset = builder.create_string(self);
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhlstring::create(
+                    builder,
+                    &FbhlstringArgs {
+                        value: Some(string_offset),
+                    },
+                )
+                .as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hlstring,
+        }
+    }
 }
 
-pub fn get_flatbuffer_result_from_long(value: i64) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-    let hllong = Fbhllong::create(&mut builder, &FbhllongArgs { value });
-
-    let rt = FbReturnValue::hllong;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hllong.as_union_value());
-
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for &[u8] {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        let vec_offset = builder.create_vector(self);
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhlsizeprefixedbuffer::create(
+                    builder,
+                    &FbhlsizeprefixedbufferArgs {
+                        size_: self.len() as i32,
+                        value: Some(vec_offset),
+                    },
+                )
+                .as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hlsizeprefixedbuffer,
+        }
+    }
 }
 
-pub fn get_flatbuffer_result_from_ulong(value: u64) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-    let hlulong = Fbhlulong::create(&mut builder, &FbhlulongArgs { value });
-
-    let rt = FbReturnValue::hlulong;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hlulong.as_union_value());
-
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for f32 {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhlfloat::create(builder, &FbhlfloatArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hlfloat,
+        }
+    }
 }
 
-pub fn get_flatbuffer_result_from_void() -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-    let hlvoid = Fbhlvoid::create(&mut builder, &FbhlvoidArgs {});
-
-    let rt = FbReturnValue::hlvoid;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hlvoid.as_union_value());
-
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for f64 {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhldouble::create(builder, &FbhldoubleArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hldouble,
+        }
+    }
 }
 
-pub fn get_flatbuffer_result_from_string(value: &str) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-
-    let string_offset = builder.create_string(value);
-    let hlstring = Fbhlstring::create(
-        &mut builder,
-        &FbhlstringArgs {
-            value: Some(string_offset),
-        },
-    );
-
-    let rt = FbReturnValue::hlstring;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hlstring.as_union_value());
-
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for i32 {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhlint::create(builder, &FbhlintArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hlint,
+        }
+    }
 }
 
-pub fn get_flatbuffer_result_from_vec(data: &[u8]) -> Vec<u8> {
-    let mut builder = FlatBufferBuilder::new();
-
-    let vec_offset = builder.create_vector(data);
-
-    let hlsizeprefixedbuffer = Fbhlsizeprefixedbuffer::create(
-        &mut builder,
-        &FbhlsizeprefixedbufferArgs {
-            size_: data.len() as i32,
-            value: Some(vec_offset),
-        },
-    );
-
-    // Indicate that the return value is a size-prefixed buffer.
-    let rt = FbReturnValue::hlsizeprefixedbuffer;
-    let rv: Option<WIPOffset<UnionWIPOffset>> = Some(hlsizeprefixedbuffer.as_union_value());
-
-    // Get the FlatBuffer result.
-    get_flatbuffer_result(&mut builder, rt, rv)
+impl FlatbufferSerializable for i64 {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhllong::create(builder, &FbhllongArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hllong,
+        }
+    }
 }
 
-fn get_flatbuffer_result(
-    builder: &mut FlatBufferBuilder,
-    return_value_type: FbReturnValue,
-    return_value: Option<WIPOffset<UnionWIPOffset>>,
-) -> Vec<u8> {
-    let result_offset = FbFunctionCallResult::create(
-        builder,
-        &FbFunctionCallResultArgs {
-            return_value,
-            return_value_type,
-        },
-    );
+impl FlatbufferSerializable for u32 {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhluint::create(builder, &FbhluintArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hluint,
+        }
+    }
+}
 
-    builder.finish_size_prefixed(result_offset, None);
+impl FlatbufferSerializable for u64 {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhlulong::create(builder, &FbhlulongArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hlulong,
+        }
+    }
+}
 
-    builder.finished_data().to_vec()
+impl FlatbufferSerializable for bool {
+    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
+        FbFunctionCallResultArgs {
+            return_value: Some(
+                Fbhlbool::create(builder, &FbhlboolArgs { value: *self }).as_union_value(),
+            ),
+            return_value_type: FbReturnValue::hlbool,
+        }
+    }
 }

src/hyperlight_guest/src/host_function_call.rs

@@ -17,13 +17,13 @@ limitations under the License.
 use alloc::string::ToString;
 use alloc::vec::Vec;
 use core::arch::global_asm;
+use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     ParameterValue, ReturnType, ReturnValue,
 };
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
-use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result_from_int;
 use hyperlight_common::mem::RunMode;
 
 use crate::error::{HyperlightGuestError, Result};
@@ -195,7 +195,7 @@ pub fn print_output_as_guest_function(function_call: &FunctionCall) -> Result<Ve
             ReturnType::Int,
         )?;
         let res_i = get_host_value_return_as_int()?;
-        Ok(get_flatbuffer_result_from_int(res_i))
+        Ok(get_flatbuffer_result(res_i))
     } else {
         Err(HyperlightGuestError::new(
             ErrorCode::GuestError,

src/hyperlight_guest_capi/src/flatbuffer.rs

@@ -1,13 +1,7 @@
 use alloc::boxed::Box;
 use core::ffi::{c_char, CStr};
 
-use hyperlight_common::flatbuffer_wrappers::util::{
-    get_flatbuffer_result_from_double, get_flatbuffer_result_from_float,
-    get_flatbuffer_result_from_int, get_flatbuffer_result_from_long,
-    get_flatbuffer_result_from_string, get_flatbuffer_result_from_uint,
-    get_flatbuffer_result_from_ulong, get_flatbuffer_result_from_vec,
-    get_flatbuffer_result_from_void,
-};
+use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 use hyperlight_guest::host_function_call::{
     get_host_value_return_as_int, get_host_value_return_as_long, get_host_value_return_as_uint,
     get_host_value_return_as_ulong,
@@ -21,57 +15,57 @@ use crate::types::FfiVec;
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_Int(value: i32) -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_int(value);
+    let vec = get_flatbuffer_result(value);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_UInt(value: u32) -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_uint(value);
+    let vec = get_flatbuffer_result(value);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_Long(value: i64) -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_long(value);
+    let vec = get_flatbuffer_result(value);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_ULong(value: u64) -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_ulong(value);
+    let vec = get_flatbuffer_result(value);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_Float(value: f32) -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_float(value);
+    let vec = get_flatbuffer_result(value);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_Double(value: f64) -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_double(value);
+    let vec = get_flatbuffer_result(value);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_Void() -> Box<FfiVec> {
-    let vec = get_flatbuffer_result_from_void();
+    let vec = get_flatbuffer_result(());
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
 
 #[no_mangle]
 pub extern "C" fn hl_flatbuffer_result_from_String(value: *const c_char) -> Box<FfiVec> {
     let str = unsafe { CStr::from_ptr(value) };
-    let vec = get_flatbuffer_result_from_string(str.to_string_lossy().as_ref());
+    let vec = get_flatbuffer_result(str.to_string_lossy().as_ref());
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }
@@ -80,7 +74,7 @@ pub extern "C" fn hl_flatbuffer_result_from_String(value: *const c_char) -> Box<
 pub extern "C" fn hl_flatbuffer_result_from_Bytes(data: *const u8, len: usize) -> Box<FfiVec> {
     let slice = unsafe { core::slice::from_raw_parts(data, len) };
 
-    let vec = get_flatbuffer_result_from_vec(slice);
+    let vec = get_flatbuffer_result(slice);
 
     Box::new(unsafe { FfiVec::from_vec(vec) })
 }