Skip to main content

hyperlight_common/flatbuffer_wrappers/
util.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2025 The Hyperlight Authors.
3
4use alloc::vec::Vec;
5
6use flatbuffers::FlatBufferBuilder;
7
8use crate::flatbuffer_wrappers::function_types::ParameterValue;
9use crate::flatbuffers::hyperlight::generated::{
10    FunctionCallResult as FbFunctionCallResult, FunctionCallResultArgs as FbFunctionCallResultArgs,
11    FunctionCallResultType as FbFunctionCallResultType, ReturnValue as FbReturnValue,
12    ReturnValueBox, ReturnValueBoxArgs, hlbool as Fbhlbool, hlboolArgs as FbhlboolArgs,
13    hldouble as Fbhldouble, hldoubleArgs as FbhldoubleArgs, hlfloat as Fbhlfloat,
14    hlfloatArgs as FbhlfloatArgs, hlint as Fbhlint, hlintArgs as FbhlintArgs, hllong as Fbhllong,
15    hllongArgs as FbhllongArgs, hlsizeprefixedbuffer as Fbhlsizeprefixedbuffer,
16    hlsizeprefixedbufferArgs as FbhlsizeprefixedbufferArgs, hlstring as Fbhlstring,
17    hlstringArgs as FbhlstringArgs, hluint as Fbhluint, hluintArgs as FbhluintArgs,
18    hlulong as Fbhlulong, hlulongArgs as FbhlulongArgs, hlvoid as Fbhlvoid,
19    hlvoidArgs as FbhlvoidArgs,
20};
21
22/// Flatbuffer-encodes the given value
23pub fn get_flatbuffer_result<T: FlatbufferSerializable>(val: T) -> Vec<u8> {
24    let mut builder = FlatBufferBuilder::new();
25    let res = T::serialize(&val, &mut builder);
26    let result_offset = FbFunctionCallResult::create(&mut builder, &res);
27
28    builder.finish_size_prefixed(result_offset, None);
29
30    builder.finished_data().to_vec()
31}
32
33pub trait FlatbufferSerializable {
34    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs;
35}
36
37// Implementations for basic types below
38
39impl FlatbufferSerializable for () {
40    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
41        let void_off = Fbhlvoid::create(builder, &FbhlvoidArgs {});
42        let rv_box = ReturnValueBox::create(
43            builder,
44            &ReturnValueBoxArgs {
45                value_type: FbReturnValue::hlvoid,
46                value: Some(void_off.as_union_value()),
47            },
48        );
49        FbFunctionCallResultArgs {
50            result_type: FbFunctionCallResultType::ReturnValueBox,
51            result: Some(rv_box.as_union_value()),
52        }
53    }
54}
55
56impl FlatbufferSerializable for &str {
57    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
58        let string_offset = builder.create_string(self);
59        let str_off = Fbhlstring::create(
60            builder,
61            &FbhlstringArgs {
62                value: Some(string_offset),
63            },
64        );
65        let rv_box = ReturnValueBox::create(
66            builder,
67            &ReturnValueBoxArgs {
68                value_type: FbReturnValue::hlstring,
69                value: Some(str_off.as_union_value()),
70            },
71        );
72        FbFunctionCallResultArgs {
73            result_type: FbFunctionCallResultType::ReturnValueBox,
74            result: Some(rv_box.as_union_value()),
75        }
76    }
77}
78
79impl FlatbufferSerializable for &[u8] {
80    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
81        let vec_off = builder.create_vector(self);
82        let buf_off = Fbhlsizeprefixedbuffer::create(
83            builder,
84            &FbhlsizeprefixedbufferArgs {
85                size: self.len() as i32,
86                value: Some(vec_off),
87            },
88        );
89        let rv_box = ReturnValueBox::create(
90            builder,
91            &ReturnValueBoxArgs {
92                value_type: FbReturnValue::hlsizeprefixedbuffer,
93                value: Some(buf_off.as_union_value()),
94            },
95        );
96        FbFunctionCallResultArgs {
97            result_type: FbFunctionCallResultType::ReturnValueBox,
98            result: Some(rv_box.as_union_value()),
99        }
100    }
101}
102
103impl FlatbufferSerializable for f32 {
104    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
105        let off = Fbhlfloat::create(builder, &FbhlfloatArgs { value: *self });
106        let rv_box = ReturnValueBox::create(
107            builder,
108            &ReturnValueBoxArgs {
109                value_type: FbReturnValue::hlfloat,
110                value: Some(off.as_union_value()),
111            },
112        );
113        FbFunctionCallResultArgs {
114            result_type: FbFunctionCallResultType::ReturnValueBox,
115            result: Some(rv_box.as_union_value()),
116        }
117    }
118}
119
120impl FlatbufferSerializable for f64 {
121    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
122        let off = Fbhldouble::create(builder, &FbhldoubleArgs { value: *self });
123        let rv_box = ReturnValueBox::create(
124            builder,
125            &ReturnValueBoxArgs {
126                value_type: FbReturnValue::hldouble,
127                value: Some(off.as_union_value()),
128            },
129        );
130        FbFunctionCallResultArgs {
131            result_type: FbFunctionCallResultType::ReturnValueBox,
132            result: Some(rv_box.as_union_value()),
133        }
134    }
135}
136
137impl FlatbufferSerializable for i32 {
138    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
139        let off = Fbhlint::create(builder, &FbhlintArgs { value: *self });
140        let rv_box = ReturnValueBox::create(
141            builder,
142            &ReturnValueBoxArgs {
143                value_type: FbReturnValue::hlint,
144                value: Some(off.as_union_value()),
145            },
146        );
147        FbFunctionCallResultArgs {
148            result_type: FbFunctionCallResultType::ReturnValueBox,
149            result: Some(rv_box.as_union_value()),
150        }
151    }
152}
153
154impl FlatbufferSerializable for i64 {
155    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
156        let off = Fbhllong::create(builder, &FbhllongArgs { value: *self });
157        let rv_box = ReturnValueBox::create(
158            builder,
159            &ReturnValueBoxArgs {
160                value_type: FbReturnValue::hllong,
161                value: Some(off.as_union_value()),
162            },
163        );
164        FbFunctionCallResultArgs {
165            result_type: FbFunctionCallResultType::ReturnValueBox,
166            result: Some(rv_box.as_union_value()),
167        }
168    }
169}
170
171impl FlatbufferSerializable for u32 {
172    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
173        let off = Fbhluint::create(builder, &FbhluintArgs { value: *self });
174        let rv_box = ReturnValueBox::create(
175            builder,
176            &ReturnValueBoxArgs {
177                value_type: FbReturnValue::hluint,
178                value: Some(off.as_union_value()),
179            },
180        );
181        FbFunctionCallResultArgs {
182            result_type: FbFunctionCallResultType::ReturnValueBox,
183            result: Some(rv_box.as_union_value()),
184        }
185    }
186}
187
188impl FlatbufferSerializable for u64 {
189    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
190        let off = Fbhlulong::create(builder, &FbhlulongArgs { value: *self });
191        let rv_box = ReturnValueBox::create(
192            builder,
193            &ReturnValueBoxArgs {
194                value_type: FbReturnValue::hlulong,
195                value: Some(off.as_union_value()),
196            },
197        );
198        FbFunctionCallResultArgs {
199            result_type: FbFunctionCallResultType::ReturnValueBox,
200            result: Some(rv_box.as_union_value()),
201        }
202    }
203}
204
205impl FlatbufferSerializable for bool {
206    fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs {
207        let off = Fbhlbool::create(builder, &FbhlboolArgs { value: *self });
208        let rv_box = ReturnValueBox::create(
209            builder,
210            &ReturnValueBoxArgs {
211                value_type: FbReturnValue::hlbool,
212                value: Some(off.as_union_value()),
213            },
214        );
215        FbFunctionCallResultArgs {
216            result_type: FbFunctionCallResultType::ReturnValueBox,
217            result: Some(rv_box.as_union_value()),
218        }
219    }
220}
221
222/// Estimates the required buffer capacity for encoding a FunctionCall with the given parameters.
223/// This helps avoid reallocation during FlatBuffer encoding when passing large slices and strings.
224///
225/// The function aims to be lightweight and fast and run in O(1) as long as the number of parameters is limited
226/// (which it is since hyperlight only currently supports up to 12).
227///
228/// Note: This estimates the capacity needed for the inner vec inside a FlatBufferBuilder. It does not
229/// necessarily match the size of the final encoded buffer. The estimation always rounds up to the
230/// nearest power of two to match FlatBufferBuilder's allocation strategy.
231///
232/// The estimations are numbers used are empirically derived based on the tests below and vaguely based
233/// on https://flatbuffers.dev/internals/ and https://github.com/dvidelabs/flatcc/blob/f064cefb2034d1e7407407ce32a6085c322212a7/doc/binary-format.md#flatbuffers-binary-format
234#[inline] // allow cross-crate inlining (for hyperlight-host calls)
235pub fn estimate_flatbuffer_capacity(function_name: &str, args: &[ParameterValue]) -> usize {
236    let mut estimated_capacity = 20;
237
238    // Function name overhead
239    estimated_capacity += function_name.len() + 12;
240
241    // Parameters vector overhead
242    estimated_capacity += 12 + args.len() * 6;
243
244    // Per-parameter overhead
245    for arg in args {
246        estimated_capacity += 16; // Base parameter structure
247        estimated_capacity += match arg {
248            ParameterValue::String(s) => s.len() + 20,
249            ParameterValue::VecBytes(v) => v.len() + 20,
250            ParameterValue::Int(_) | ParameterValue::UInt(_) => 16,
251            ParameterValue::Long(_) | ParameterValue::ULong(_) => 20,
252            ParameterValue::Float(_) => 16,
253            ParameterValue::Double(_) => 20,
254            ParameterValue::Bool(_) => 12,
255        };
256    }
257
258    // match how vec grows
259    estimated_capacity.next_power_of_two()
260}
261
262#[cfg(test)]
263mod tests {
264    use alloc::string::ToString;
265    use alloc::vec;
266    use alloc::vec::Vec;
267
268    use super::*;
269    use crate::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
270    use crate::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
271
272    /// Helper function to check that estimation is within reasonable bounds (±25%)
273    fn assert_estimation_accuracy(
274        function_name: &str,
275        args: Vec<ParameterValue>,
276        call_type: FunctionCallType,
277        return_type: ReturnType,
278    ) {
279        let estimated = estimate_flatbuffer_capacity(function_name, &args);
280
281        let fc = FunctionCall::new(
282            function_name.to_string(),
283            Some(args),
284            call_type.clone(),
285            return_type,
286        );
287        // Important that this FlatBufferBuilder is created with capacity 0 so it grows to its needed capacity
288        let mut builder = FlatBufferBuilder::new();
289        let _buffer = fc.encode(&mut builder);
290        let actual = builder.collapse().0.capacity();
291
292        let lower_bound = (actual as f64 * 0.75) as usize;
293        let upper_bound = (actual as f64 * 1.25) as usize;
294
295        assert!(
296            estimated >= lower_bound && estimated <= upper_bound,
297            "Estimation {} outside bounds [{}, {}] for actual size {} (function: {}, call_type: {:?}, return_type: {:?})",
298            estimated,
299            lower_bound,
300            upper_bound,
301            actual,
302            function_name,
303            call_type,
304            return_type
305        );
306    }
307
308    #[test]
309    fn test_estimate_no_parameters() {
310        assert_estimation_accuracy(
311            "simple_function",
312            vec![],
313            FunctionCallType::Guest,
314            ReturnType::Void,
315        );
316    }
317
318    #[test]
319    fn test_estimate_single_int_parameter() {
320        assert_estimation_accuracy(
321            "add_one",
322            vec![ParameterValue::Int(42)],
323            FunctionCallType::Guest,
324            ReturnType::Int,
325        );
326    }
327
328    #[test]
329    fn test_estimate_multiple_scalar_parameters() {
330        assert_estimation_accuracy(
331            "calculate",
332            vec![
333                ParameterValue::Int(10),
334                ParameterValue::UInt(20),
335                ParameterValue::Long(30),
336                ParameterValue::ULong(40),
337                ParameterValue::Float(1.5),
338                ParameterValue::Double(2.5),
339                ParameterValue::Bool(true),
340            ],
341            FunctionCallType::Guest,
342            ReturnType::Double,
343        );
344    }
345
346    #[test]
347    fn test_estimate_string_parameters() {
348        assert_estimation_accuracy(
349            "process_strings",
350            vec![
351                ParameterValue::String("hello".to_string()),
352                ParameterValue::String("world".to_string()),
353                ParameterValue::String("this is a longer string for testing".to_string()),
354            ],
355            FunctionCallType::Host,
356            ReturnType::String,
357        );
358    }
359
360    #[test]
361    fn test_estimate_very_long_string() {
362        let long_string = "a".repeat(1000);
363        assert_estimation_accuracy(
364            "process_long_string",
365            vec![ParameterValue::String(long_string)],
366            FunctionCallType::Guest,
367            ReturnType::String,
368        );
369    }
370
371    #[test]
372    fn test_estimate_vector_parameters() {
373        assert_estimation_accuracy(
374            "process_vectors",
375            vec![
376                ParameterValue::VecBytes(vec![1, 2, 3, 4, 5]),
377                ParameterValue::VecBytes(vec![]),
378                ParameterValue::VecBytes(vec![0; 100]),
379            ],
380            FunctionCallType::Host,
381            ReturnType::VecBytes,
382        );
383    }
384
385    #[test]
386    fn test_estimate_mixed_parameters() {
387        assert_estimation_accuracy(
388            "complex_function",
389            vec![
390                ParameterValue::String("test".to_string()),
391                ParameterValue::Int(42),
392                ParameterValue::VecBytes(vec![1, 2, 3, 4, 5]),
393                ParameterValue::Bool(true),
394                ParameterValue::Double(553.14159),
395                ParameterValue::String("another string".to_string()),
396                ParameterValue::Long(9223372036854775807),
397            ],
398            FunctionCallType::Guest,
399            ReturnType::VecBytes,
400        );
401    }
402
403    #[test]
404    fn test_estimate_large_function_name() {
405        let long_name = "very_long_function_name_that_exceeds_normal_lengths_for_testing_purposes";
406        assert_estimation_accuracy(
407            long_name,
408            vec![ParameterValue::Int(1)],
409            FunctionCallType::Host,
410            ReturnType::Long,
411        );
412    }
413
414    #[test]
415    fn test_estimate_large_vector() {
416        let large_vector = vec![42u8; 10000];
417        assert_estimation_accuracy(
418            "process_large_data",
419            vec![ParameterValue::VecBytes(large_vector)],
420            FunctionCallType::Guest,
421            ReturnType::Bool,
422        );
423    }
424
425    #[test]
426    fn test_estimate_all_parameter_types() {
427        assert_estimation_accuracy(
428            "comprehensive_test",
429            vec![
430                ParameterValue::Int(i32::MIN),
431                ParameterValue::UInt(u32::MAX),
432                ParameterValue::Long(i64::MIN),
433                ParameterValue::ULong(u64::MAX),
434                ParameterValue::Float(f32::MIN),
435                ParameterValue::Double(f64::MAX),
436                ParameterValue::Bool(false),
437                ParameterValue::String("test string".to_string()),
438                ParameterValue::VecBytes(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
439            ],
440            FunctionCallType::Host,
441            ReturnType::ULong,
442        );
443    }
444
445    #[test]
446    fn test_different_function_call_types() {
447        assert_estimation_accuracy(
448            "guest_function",
449            vec![ParameterValue::String("guest call".to_string())],
450            FunctionCallType::Guest,
451            ReturnType::String,
452        );
453
454        assert_estimation_accuracy(
455            "host_function",
456            vec![ParameterValue::String("host call".to_string())],
457            FunctionCallType::Host,
458            ReturnType::String,
459        );
460    }
461
462    #[test]
463    fn test_different_return_types() {
464        let args = vec![
465            ParameterValue::Int(42),
466            ParameterValue::String("test".to_string()),
467        ];
468
469        let void_est = estimate_flatbuffer_capacity("test_void", &args);
470        let int_est = estimate_flatbuffer_capacity("test_int", &args);
471        let string_est = estimate_flatbuffer_capacity("test_string", &args);
472
473        assert!((void_est as i32 - int_est as i32).abs() < 10);
474        assert!((int_est as i32 - string_est as i32).abs() < 10);
475
476        assert_estimation_accuracy(
477            "test_void",
478            args.clone(),
479            FunctionCallType::Guest,
480            ReturnType::Void,
481        );
482        assert_estimation_accuracy(
483            "test_int",
484            args.clone(),
485            FunctionCallType::Guest,
486            ReturnType::Int,
487        );
488        assert_estimation_accuracy(
489            "test_string",
490            args,
491            FunctionCallType::Guest,
492            ReturnType::String,
493        );
494    }
495
496    #[test]
497    fn test_estimate_many_large_vectors_and_strings() {
498        assert_estimation_accuracy(
499            "process_bulk_data",
500            vec![
501                ParameterValue::String("Large string data: ".to_string() + &"x".repeat(2000)),
502                ParameterValue::VecBytes(vec![1u8; 5000]),
503                ParameterValue::String(
504                    "Another large string with lots of content ".to_string() + &"y".repeat(3000),
505                ),
506                ParameterValue::VecBytes(vec![255u8; 7500]),
507                ParameterValue::String(
508                    "Third massive string parameter ".to_string() + &"z".repeat(1500),
509                ),
510                ParameterValue::VecBytes(vec![128u8; 10000]),
511                ParameterValue::Int(42),
512                ParameterValue::String("Final large string ".to_string() + &"a".repeat(4000)),
513                ParameterValue::VecBytes(vec![64u8; 2500]),
514                ParameterValue::Bool(true),
515            ],
516            FunctionCallType::Host,
517            ReturnType::VecBytes,
518        );
519    }
520
521    #[test]
522    fn test_estimate_twenty_parameters() {
523        assert_estimation_accuracy(
524            "function_with_many_parameters",
525            vec![
526                ParameterValue::Int(1),
527                ParameterValue::String("param2".to_string()),
528                ParameterValue::Bool(true),
529                ParameterValue::Float(3213.14),
530                ParameterValue::VecBytes(vec![1, 2, 3]),
531                ParameterValue::Long(1000000),
532                ParameterValue::Double(322.718),
533                ParameterValue::UInt(42),
534                ParameterValue::String("param9".to_string()),
535                ParameterValue::Bool(false),
536                ParameterValue::ULong(9999999999),
537                ParameterValue::VecBytes(vec![4, 5, 6, 7, 8]),
538                ParameterValue::Int(-100),
539                ParameterValue::Float(1.414),
540                ParameterValue::String("param15".to_string()),
541                ParameterValue::Double(1.732),
542                ParameterValue::Bool(true),
543                ParameterValue::VecBytes(vec![9, 10]),
544                ParameterValue::Long(-5000000),
545                ParameterValue::UInt(12345),
546            ],
547            FunctionCallType::Guest,
548            ReturnType::Int,
549        );
550    }
551
552    /// When running under Miri, this test takes too long due to large allocations
553    #[test]
554    #[cfg_attr(miri, ignore)]
555    fn test_estimate_megabyte_parameters() {
556        assert_estimation_accuracy(
557            "process_megabyte_data",
558            vec![
559                ParameterValue::String("MB String 1: ".to_string() + &"x".repeat(1_048_576)), // 1MB string
560                ParameterValue::VecBytes(vec![42u8; 2_097_152]), // 2MB vector
561                ParameterValue::String("MB String 2: ".to_string() + &"y".repeat(1_572_864)), // 1.5MB string
562                ParameterValue::VecBytes(vec![128u8; 3_145_728]), // 3MB vector
563                ParameterValue::String("MB String 3: ".to_string() + &"z".repeat(2_097_152)), // 2MB string
564            ],
565            FunctionCallType::Host,
566            ReturnType::VecBytes,
567        );
568    }
569}