fbthrift_git/
bufext.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::io::Cursor;
18
19use bufsize::SizeCounter;
20use bytes::buf::Chain;
21use bytes::Buf;
22use bytes::BufMut;
23use bytes::Bytes;
24use bytes::BytesMut;
25
26use crate::varint;
27
28pub trait BufExt: Buf {
29    /// Copy `len` Bytes from this (and advance the position), or reuse them from the underlying
30    /// buffer if possible.
31    fn copy_or_reuse_bytes(&mut self, len: usize) -> Bytes {
32        // Default is to just copy.
33        self.copy_to_bytes(len)
34    }
35}
36
37impl BufExt for Bytes {}
38
39impl BufExt for Cursor<Bytes> {
40    // We can get a reference to the underlying Bytes here, and reuse that.
41    fn copy_or_reuse_bytes(&mut self, len: usize) -> Bytes {
42        let pos = self.position() as usize;
43        let end = pos + len;
44        // Panics if len is too large (same as Bytes)
45        let bytes = self.get_ref().slice(pos..end);
46        self.set_position(end as u64);
47        bytes
48    }
49}
50
51impl<T: AsRef<[u8]> + ?Sized> BufExt for Cursor<&T> {}
52
53impl<T: BufExt, U: BufExt> BufExt for Chain<T, U> {}
54
55pub trait BufMutExt: BufMut {
56    type Final: Send + 'static;
57
58    fn put_varint_u64(&mut self, v: u64)
59    where
60        Self: Sized,
61    {
62        varint::write_u64(self, v)
63    }
64
65    fn put_varint_i64(&mut self, v: i64)
66    where
67        Self: Sized,
68    {
69        varint::write_u64(self, varint::zigzag(v))
70    }
71
72    fn finalize(self) -> Self::Final;
73}
74
75impl BufMutExt for BytesMut {
76    type Final = Bytes;
77
78    fn finalize(self) -> Self::Final {
79        self.freeze()
80    }
81}
82
83impl BufMutExt for SizeCounter {
84    type Final = usize;
85
86    #[inline]
87    fn put_varint_u64(&mut self, v: u64) {
88        self.put_uint(v, varint::u64_len(v));
89    }
90
91    #[inline]
92    fn put_varint_i64(&mut self, v: i64) {
93        self.put_int(v, varint::u64_len(varint::zigzag(v)));
94    }
95
96    #[inline]
97    fn finalize(self) -> Self::Final {
98        self.size()
99    }
100}
101// new type so we can impl From for Bytes vs other things that are AsRef<[u8]>
102// Not implemented for AsRef<[u8]> as Bytes implements that
103pub struct DeserializeSource<B: BufExt>(pub(crate) B);
104
105impl<B: BufExt> DeserializeSource<B> {
106    pub fn new(b: B) -> Self {
107        DeserializeSource(b)
108    }
109}
110
111// These types will use a copying cursor
112macro_rules! impl_deser_as_ref_u8 {
113    ( $($t:ty),* ) => {
114        $(
115            impl<'a> From<&'a $t> for DeserializeSource<Cursor<&'a [u8]>> {
116                fn from(from: &'a $t) -> Self {
117                    let data: &[u8] = from.as_ref();
118                    Self(Cursor::new(data))
119                }
120            }
121        )*
122    }
123}
124
125impl_deser_as_ref_u8!([u8], Vec<u8>, String, str);
126
127// These types take ownership without copying
128// Have to explicitly do the Into types as well due to no upstream From<&Bytes> for Bytes.
129macro_rules! impl_deser_into_bytes {
130    ( $($t:ty),* ) => {
131        $(
132            impl From<$t> for DeserializeSource<Cursor<Bytes>> {
133                fn from(from: $t) -> Self {
134                    Self(Cursor::new(from.into()))
135                }
136            }
137        )*
138    }
139}
140
141impl_deser_into_bytes!(Bytes, Vec<u8>, String);
142
143// Special case for &Bytes that is not covered in upstream crates From defs
144impl From<&Bytes> for DeserializeSource<Cursor<Bytes>> {
145    fn from(from: &Bytes) -> Self {
146        // ok to clone Bytes, it just increments ref count
147        Self(Cursor::new(from.clone()))
148    }
149}
fbthrift_git/bufext.rs

fbthrift_git/
bufext.rs
