bitcode_lightyear_patch/
code.rs

1use crate::buffer::BufferTrait;
2use crate::encoding::{Encoding, Fixed};
3use crate::read::Read;
4use crate::write::Write;
5use crate::Result;
6
7pub(crate) fn encode_internal<'a>(
8    buffer: &'a mut impl BufferTrait,
9    t: &(impl Encode + ?Sized),
10) -> Result<&'a [u8]> {
11    let mut writer = buffer.start_write();
12    t.encode(Fixed, &mut writer)?;
13    Ok(buffer.finish_write(writer))
14}
15
16pub(crate) fn decode_internal<B: BufferTrait, T: Decode>(
17    buffer: &mut B,
18    bytes: &[u8],
19) -> Result<T> {
20    let (mut reader, context) = buffer.start_read(bytes);
21    let decode_result = T::decode(Fixed, &mut reader);
22    B::finish_read_with_result(reader, context, decode_result)
23}
24
25/// A type which can be encoded to bytes with [`encode`][`crate::encode`].
26///
27/// Must use `#[derive(Encode)]` to implement.
28/// ```
29/// #[derive(bitcode::Encode)]
30/// // If your struct contains itself you must annotate it with `#[bitcode(recursive)]`.
31/// // This disables certain speed optimizations that aren't possible on recursive types.
32/// struct MyStruct {
33///     a: u32,
34///     b: bool,
35///     // If you want to use serde::Serialize on a field instead of bitcode::Encode.
36///     #[cfg(feature = "serde")]
37///     #[bitcode(with_serde)]
38///     c: String,
39/// }
40/// ```
41pub trait Encode {
42    // The minimum and maximum number of bits a type can encode as. For now these are only valid if
43    // the encoding is fixed. Before using them make sure the encoding passed to encode is fixed.
44    // TODO make these const functions that take an encoding (once const fn is available in traits).
45    #[doc(hidden)]
46    const ENCODE_MIN: usize;
47
48    // If max is lower than the actual max, we may not encode all the bits.
49    #[doc(hidden)]
50    const ENCODE_MAX: usize;
51
52    #[doc(hidden)]
53    fn encode(&self, encoding: impl Encoding, writer: &mut impl Write) -> Result<()>;
54}
55
56/// A type which can be decoded from bytes with [`decode`][`crate::decode`].
57///
58/// Must use `#[derive(Decode)]` to implement.
59/// ```
60/// #[derive(bitcode::Decode)]
61/// // If your struct contains itself you must annotate it with `#[bitcode(recursive)]`.
62/// // This disables certain speed optimizations that aren't possible on recursive types.
63/// struct MyStruct {
64///     a: u32,
65///     b: bool,
66///     // If you want to use serde::Deserialize on a field instead of bitcode::Decode.
67///     #[cfg(feature = "serde")]
68///     #[bitcode(with_serde)]
69///     c: String,
70/// }
71/// ```
72pub trait Decode: Sized {
73    // Copy of Encode constants. See Encode for details.
74    // If min is higher than the actual min, we may get EOFs.
75    #[doc(hidden)]
76    const DECODE_MIN: usize;
77
78    #[doc(hidden)]
79    const DECODE_MAX: usize;
80
81    #[doc(hidden)]
82    fn decode(encoding: impl Encoding, reader: &mut impl Read) -> Result<Self>;
83}
84
85/// A macro that facilitates writing to a RegisterWriter when encoding multiple values less than 64 bits.
86/// This can dramatically speed operations like encoding a tuple of 8 bytes.
87///
88/// Once you call `optimized_enc!()`, you must call `end_enc!()` at the end to flush the remaining bits.
89/// If the execution path diverges it must never converge or this won't optimize well.
90#[doc(hidden)]
91#[macro_export]
92macro_rules! optimized_enc {
93    ($encoding:ident, $writer:ident) => {
94        let mut buf = $crate::__private::RegisterWriter::new($writer);
95        #[allow(unused_mut)]
96        let mut i: usize = 0;
97        #[allow(unused)]
98        let no_encoding_upstream = $encoding.is_fixed();
99
100        // Call on each field (that doesn't get it's encoding overridden in the derive macro).
101        #[allow(unused)]
102        macro_rules! enc {
103            ($t:expr, $T:ty) => {
104                // ENCODE_MAX is only accurate if there isn't any encoding upstream.
105                // Downstream encodings make ENCODE_MAX = usize::MAX in derive macro.
106                if <$T as $crate::__private::Encode>::ENCODE_MAX.saturating_add(i) <= 64
107                    && no_encoding_upstream
108                {
109                    <$T as $crate::__private::Encode>::encode(&$t, $encoding, &mut buf.inner)?;
110                } else {
111                    if i != 0 {
112                        buf.flush();
113                    }
114
115                    if <$T as $crate::__private::Encode>::ENCODE_MAX < 64 && no_encoding_upstream {
116                        <$T as $crate::__private::Encode>::encode(&$t, $encoding, &mut buf.inner)?;
117                    } else {
118                        <$T as $crate::__private::Encode>::encode(&$t, $encoding, buf.writer)?;
119                    }
120                }
121
122                i = if <$T as $crate::__private::Encode>::ENCODE_MAX.saturating_add(i) <= 64
123                    && no_encoding_upstream
124                {
125                    <$T as $crate::__private::Encode>::ENCODE_MAX + i
126                } else {
127                    if <$T as $crate::__private::Encode>::ENCODE_MAX < 64 && no_encoding_upstream {
128                        <$T as $crate::__private::Encode>::ENCODE_MAX
129                    } else {
130                        0
131                    }
132                };
133            };
134        }
135
136        // Call to flush the contents of the RegisterWriter and get the inner writer.
137        macro_rules! flush {
138            () => {{
139                if i != 0 {
140                    buf.flush();
141                }
142                i = 0;
143                &mut *buf.writer
144            }};
145        }
146
147        // Call to encode an enum variant. Faster than flush!().write_bits($variant, $bits).
148        // Must be fist call after optimized_enc!.
149        #[allow(unused)]
150        macro_rules! enc_variant {
151            ($variant:literal, $bits:literal) => {
152                debug_assert!(i == 0);
153                buf.inner.write_bits($variant, $bits);
154                i = $bits + i;
155            };
156        }
157
158        // Call once done encoding.
159        macro_rules! end_enc {
160            () => {
161                let _ = flush!();
162                let _ = i;
163                #[allow(clippy::drop_non_drop)]
164                drop(buf);
165            };
166        }
167    };
168}
169pub use optimized_enc;
170
171// These benchmarks ensure that optimized_enc is working. They all run about 8 times faster with optimized_enc.
172#[cfg(all(test, not(miri)))]
173mod optimized_enc_tests {
174    use std::collections::{BinaryHeap, VecDeque};
175    use test::{black_box, Bencher};
176
177    type A = u8;
178    type B = u8;
179
180    #[derive(Clone, Debug, PartialEq, crate::Encode, crate::Decode)]
181    struct Foo {
182        a: A,
183        b: B,
184    }
185
186    #[bench]
187    fn bench_foo(b: &mut Bencher) {
188        let mut buffer = crate::Buffer::new();
189        let foo = Foo { a: 1, b: 2 };
190        let foo = vec![foo; 4000];
191
192        let bytes = buffer.encode(&foo).unwrap().to_vec();
193        let decoded: Vec<Foo> = buffer.decode(&bytes).unwrap();
194        assert_eq!(foo, decoded);
195
196        b.iter(|| {
197            let foo = black_box(foo.as_slice());
198            let bytes = buffer.encode(foo).unwrap();
199            black_box(bytes);
200        })
201    }
202
203    #[bench]
204    fn bench_tuple(b: &mut Bencher) {
205        let mut buffer = crate::Buffer::new();
206        let foo = vec![(0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8); 1000];
207
208        b.iter(|| {
209            let foo = black_box(foo.as_slice());
210            let bytes = buffer.encode(foo).unwrap();
211            black_box(bytes);
212        })
213    }
214
215    #[bench]
216    fn bench_array(b: &mut Bencher) {
217        let mut buffer = crate::Buffer::new();
218        let foo = vec![[0u8; 8]; 1000];
219
220        b.iter(|| {
221            let foo = black_box(foo.as_slice());
222            let bytes = buffer.encode(foo).unwrap();
223            black_box(bytes);
224        })
225    }
226
227    #[bench]
228    fn bench_bool_slice(b: &mut Bencher) {
229        let mut buffer = crate::Buffer::new();
230        let foo = vec![false; 8 * 1000];
231
232        b.iter(|| {
233            let foo = black_box(foo.as_slice());
234            let bytes = buffer.encode(foo).unwrap();
235            black_box(bytes);
236        })
237    }
238
239    #[bench]
240    fn bench_vec(b: &mut Bencher) {
241        let mut buffer = crate::Buffer::new();
242        let foo = vec![0u8; 8 * 1000];
243
244        b.iter(|| {
245            let foo = black_box(foo.as_slice());
246            let bytes = buffer.encode(foo).unwrap();
247            black_box(bytes);
248        })
249    }
250
251    #[bench]
252    fn bench_vec_deque(b: &mut Bencher) {
253        let mut buffer = crate::Buffer::new();
254        let mut foo = VecDeque::from(vec![0u8; 8000]);
255        for _ in 0..4000 {
256            // Make it not contiguous.
257            foo.pop_front().unwrap();
258            foo.push_back(1u8);
259        }
260
261        b.iter(|| {
262            let foo = black_box(&foo);
263            let bytes = buffer.encode(foo).unwrap();
264            black_box(bytes);
265        })
266    }
267
268    // BinaryHeap::encode isn't optimized yet.
269    #[bench]
270    fn bench_binary_heap(b: &mut Bencher) {
271        let mut buffer = crate::Buffer::new();
272        let foo = BinaryHeap::from_iter((0u16..8000).map(|v| v as u8));
273
274        b.iter(|| {
275            let foo = black_box(&foo);
276            let bytes = buffer.encode(foo).unwrap();
277            black_box(bytes);
278        })
279    }
280}
281
282/// A macro that facilitates reading from a RegisterReader when decoding multiple values less than 64 bits.
283/// This can dramatically speed operations like decoding a tuple of 8 bytes.
284///
285/// Once you call `optimized_dec!()`, you must call `end_dec!()` at the end to advance the reader.
286/// If the execution path diverges it must never converge or this won't optimize well.
287#[doc(hidden)]
288#[macro_export]
289macro_rules! optimized_dec {
290    ($encoding:ident, $reader:ident) => {
291        #[allow(unused_mut)]
292        let mut buf = $crate::__private::RegisterReader::new($reader);
293        #[allow(unused_mut)]
294        let mut i: usize = 0;
295        #[allow(unused)]
296        let no_encoding_upstream = $encoding.is_fixed();
297
298        // Call on each field (that doesn't get it's encoding overridden in the derive macro).
299        #[allow(unused)]
300        macro_rules! dec {
301            ($t:ident, $T:ty) => {
302                // DECODE_MAX is only accurate if there isn't any encoding upstream.
303                // Downstream encodings make DECODE_MAX = usize::MAX in derive macro.
304                let $t = if i >= <$T as $crate::__private::Decode>::DECODE_MAX
305                    && no_encoding_upstream
306                {
307                    <$T as $crate::__private::Decode>::decode($encoding, &mut buf.inner)?
308                } else {
309                    if <$T as $crate::__private::Decode>::DECODE_MAX < 64 && no_encoding_upstream {
310                        buf.refill()?;
311                        <$T as $crate::__private::Decode>::decode($encoding, &mut buf.inner)?
312                    } else {
313                        buf.advance_reader();
314                        <$T as $crate::__private::Decode>::decode($encoding, buf.reader)?
315                    }
316                };
317
318                i = if i >= <$T as $crate::__private::Decode>::DECODE_MAX && no_encoding_upstream {
319                    i - <$T as $crate::__private::Decode>::DECODE_MAX
320                } else {
321                    if <$T as $crate::__private::Decode>::DECODE_MAX < 64 && no_encoding_upstream {
322                        // Needs saturating since it's const (even though we've checked it).
323                        64usize.saturating_sub(<$T as $crate::__private::Decode>::DECODE_MAX)
324                    } else {
325                        0
326                    }
327                };
328            };
329        }
330
331        // Call to flush the contents of the RegisterReader and get the inner reader.
332        macro_rules! flush {
333            () => {{
334                let _ = i;
335                i = 0;
336                buf.advance_reader();
337                &mut *buf.reader
338            }};
339        }
340
341        // Call to peek bits to decode an enum variant. Faster than flush!().peek_bits()?.
342        // Must be fist call after optimized_dec!.
343        #[allow(unused)]
344        macro_rules! dec_variant_peek {
345            () => {{
346                buf.refill()?;
347                buf.inner.peek_bits()?
348            }};
349        }
350
351        // Call to advance bits to decode an enum variant. Faster than flush!().advance($bits)?.
352        // Must be second call after dec_variant_peek!.
353        #[allow(unused)]
354        macro_rules! dec_variant_advance {
355            ($bits:literal) => {
356                debug_assert!(i == 0);
357                buf.inner.advance($bits);
358                i = 64 - $bits;
359            };
360        }
361
362        // Call once done decoding.
363        macro_rules! end_dec {
364            () => {
365                let _ = flush!();
366                let _ = i;
367                #[allow(clippy::drop_non_drop)]
368                drop(buf);
369            };
370        }
371    };
372}
373pub use optimized_dec;
374
375// These benchmarks ensure that optimized_dec is working. They run 4-8 times faster with optimized_dec.
376#[cfg(all(test, not(miri)))]
377mod optimized_dec_tests {
378    use std::collections::{BTreeSet, BinaryHeap, VecDeque};
379    use test::{black_box, Bencher};
380
381    type A = u8;
382    type B = u8;
383
384    #[derive(Clone, Debug, PartialEq, crate::Encode, crate::Decode)]
385    #[repr(C, align(8))]
386    struct Foo {
387        a: A,
388        b: B,
389        c: A,
390        d: B,
391        e: A,
392        f: B,
393        g: A,
394        h: B,
395    }
396
397    #[bench]
398    fn bench_foo(b: &mut Bencher) {
399        let mut buffer = crate::Buffer::new();
400        let foo = Foo {
401            a: 1,
402            b: 2,
403            c: 3,
404            d: 4,
405            e: 5,
406            f: 6,
407            g: 7,
408            h: 8,
409        };
410        let foo = vec![foo; 1000];
411        type T = Vec<Foo>;
412
413        let bytes = buffer.encode(&foo).unwrap().to_vec();
414        let decoded: T = buffer.decode(&bytes).unwrap();
415        assert_eq!(foo, decoded);
416
417        b.iter(|| {
418            let bytes = black_box(bytes.as_slice());
419            black_box(buffer.decode::<T>(bytes).unwrap())
420        })
421    }
422
423    #[bench]
424    fn bench_tuple(b: &mut Bencher) {
425        let mut buffer = crate::Buffer::new();
426        let foo = vec![(0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8); 1000];
427        type T = Vec<(u8, u8, u8, u8, u8, u8, u8, u8)>;
428
429        let bytes = buffer.encode(&foo).unwrap().to_vec();
430        let decoded: T = buffer.decode(&bytes).unwrap();
431        assert_eq!(foo, decoded);
432
433        b.iter(|| {
434            let bytes = black_box(bytes.as_slice());
435            black_box(buffer.decode::<T>(bytes).unwrap())
436        })
437    }
438
439    #[bench]
440    fn bench_array(b: &mut Bencher) {
441        let mut buffer = crate::Buffer::new();
442        let foo = vec![[0u8; 8]; 1000];
443        type T = Vec<[u8; 8]>;
444
445        let bytes = buffer.encode(&foo).unwrap().to_vec();
446        let decoded: T = buffer.decode(&bytes).unwrap();
447        assert_eq!(foo, decoded);
448
449        b.iter(|| {
450            let bytes = black_box(bytes.as_slice());
451            black_box(buffer.decode::<T>(bytes).unwrap())
452        })
453    }
454
455    #[bench]
456    fn bench_vec(b: &mut Bencher) {
457        let mut buffer = crate::Buffer::new();
458        let foo = vec![0u8; 8000];
459        type T = Vec<u8>;
460
461        let bytes = buffer.encode(&foo).unwrap().to_vec();
462        let decoded: T = buffer.decode(&bytes).unwrap();
463        assert_eq!(foo, decoded);
464
465        b.iter(|| {
466            let bytes = black_box(bytes.as_slice());
467            black_box(buffer.decode::<T>(bytes).unwrap())
468        })
469    }
470
471    #[bench]
472    fn bench_vec_deque(b: &mut Bencher) {
473        let mut buffer = crate::Buffer::new();
474        let mut foo = VecDeque::from(vec![0u8; 8000]);
475        for _ in 0..4000 {
476            // Make it not contiguous.
477            foo.pop_front().unwrap();
478            foo.push_back(1u8);
479        }
480        type T = VecDeque<u8>;
481
482        let bytes = buffer.encode(&foo).unwrap().to_vec();
483        let decoded: T = buffer.decode(&bytes).unwrap();
484        assert_eq!(foo, decoded);
485
486        b.iter(|| {
487            let bytes = black_box(bytes.as_slice());
488            black_box(buffer.decode::<T>(bytes).unwrap())
489        })
490    }
491
492    #[bench]
493    fn bench_binary_heap(b: &mut Bencher) {
494        let mut buffer = crate::Buffer::new();
495        let foo = BinaryHeap::from_iter((0u16..8000).map(|v| v as u8));
496        type T = BinaryHeap<u8>;
497
498        let bytes = buffer.encode(&foo).unwrap().to_vec();
499        let decoded: T = buffer.decode(&bytes).unwrap();
500
501        // Binary heaps can't be compared directly.
502        assert_eq!(
503            BTreeSet::from_iter(foo.iter().copied()),
504            BTreeSet::from_iter(decoded.iter().copied())
505        );
506
507        b.iter(|| {
508            let bytes = black_box(bytes.as_slice());
509            black_box(buffer.decode::<T>(bytes).unwrap())
510        })
511    }
512}