flatbuffers_util/
ownedfb.rs

1use flatbuffers::{Follow, InvalidFlatbuffer, Verifiable};
2
3/// Stores the owned bytes of the flatbuffer type
4/// and can access the actual type.
5pub struct OwnedFB<T> {
6    buf: Vec<u8>,
7    index: usize,
8    _phantom: std::marker::PhantomData<T>,
9}
10
11impl<T> OwnedFB<T> {
12    pub fn new<'a>(buf: &'a [u8]) -> Result<OwnedFB<T>, InvalidFlatbuffer>
13    where
14        T: Verifiable + Follow<'a> + 'a,
15    {
16        check_flatbuffer::<T>(buf, 0)?;
17        Ok(unsafe { Self::new_from_vec_unchecked(buf.to_owned(), 0) })
18    }
19
20    /// # Safety
21    /// Caller is responsible for verifying the buffer and align the type T.
22    pub unsafe fn new_from_vec_unchecked(buf: Vec<u8>, index: usize) -> Self {
23        Self {
24            buf,
25            index,
26            _phantom: std::marker::PhantomData,
27        }
28    }
29
30    /// # Safety
31    /// Caller is responsible for verifying the buffer and align the type T.
32    pub unsafe fn new_from_builder_collapse(pair: (Vec<u8>, usize)) -> Self {
33        unsafe { Self::new_from_vec_unchecked(pair.0, pair.1) }
34    }
35
36    pub fn new_from_vec(buf: Vec<u8>, index: usize) -> Result<OwnedFB<T>, InvalidFlatbuffer>
37    where
38        T: Verifiable + Follow<'static> + 'static,
39    {
40        check_flatbuffer::<T>(&buf, index)?;
41
42        Ok(unsafe { Self::new_from_vec_unchecked(buf, index) })
43    }
44
45    /// This in practice is not zero copy.
46    /// For tonic, the bytes is always shared BytesMut, so it cannot be directly turned into vec.
47    /// And in general Bytes have multiple chunks, and flatbuffer need contiguous memory.
48    /// So this will make a copy in most cases.
49    pub fn new_from_bytes(buf: bytes::Bytes) -> Result<OwnedFB<T>, InvalidFlatbuffer>
50    where
51        T: Verifiable + Follow<'static> + 'static,
52    {
53        match buf.try_into_mut() {
54            // This is zero copy if the Bytes has the full ownership of the vec.
55            Ok(bytes_mut) => Self::new_from_vec(bytes_mut.into(), 0),
56            // This will make a copy of the bytes.
57            Err(bytes) => {
58                // This happens in tonic DecodeBuf because it is a shared BytesMut.
59                // So in practice tonic forces us to make a copy here.
60                // debug_assert!(false, "zero copy failed.");
61                Self::new_from_vec(bytes.to_vec(), 0)
62            }
63        }
64    }
65
66    pub fn get_ref<'a>(&'a self) -> <T as Follow<'a>>::Inner
67    where
68        T: Follow<'a>,
69    {
70        // Safety: We have already verified the buffer in `new_owned_fb`.
71        unsafe { get_ref_flatbuffer_unchecked::<'a, T>(&self.buf, self.index) }
72    }
73
74    pub fn get_slice(&self) -> &[u8] {
75        &self.buf[self.index..]
76    }
77
78    /// This may be zero copy if the vec capacity equals to length and index is zero.
79    pub fn into_bytes(self) -> bytes::Bytes {
80        // This is zero copy if vec cap == len.
81        debug_assert_eq!(self.buf.capacity(), self.buf.len());
82        let full_bytes = bytes::Bytes::from(self.buf);
83        // adjust the offset should be zero copy.
84        full_bytes.slice(self.index..)
85    }
86}
87
88/// Generic check.
89pub fn check_flatbuffer<'a, T>(buf: &[u8], index: usize) -> Result<(), InvalidFlatbuffer>
90where
91    T: Verifiable + Follow<'a> + 'a,
92{
93    let opts = flatbuffers::VerifierOptions::default();
94    let mut v = flatbuffers::Verifier::new(&opts, buf);
95    <flatbuffers::ForwardsUOffset<T>>::run_verifier(&mut v, index)?;
96    Ok(())
97}
98
99/// # Safety
100/// Caller is responsible for verifying the buffer.
101pub unsafe fn get_ref_flatbuffer_unchecked<'a, T>(buf: &'a [u8], index: usize) -> T::Inner
102where
103    T: Follow<'a> + 'a,
104{
105    // Safety: We have already verified the buffer in `check_flatbuffer`.
106    unsafe { <flatbuffers::ForwardsUOffset<T>>::follow(buf, index) }
107}
flatbuffers_util/ownedfb.rs

flatbuffers_util/
ownedfb.rs
