hyperlight_common/flatbuffer_wrappers/
util.rs

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