rustc_serialize/opaque/
mem_encoder.rs

1use super::IntEncodedWithFixedSize;
2use crate::{Encodable, Encoder, leb128};
3
4pub struct MemEncoder {
5    pub data: Vec<u8>,
6}
7
8impl MemEncoder {
9    pub fn new() -> MemEncoder {
10        MemEncoder { data: vec![] }
11    }
12
13    #[inline]
14    pub fn position(&self) -> usize {
15        self.data.len()
16    }
17
18    pub fn finish(self) -> Vec<u8> {
19        self.data
20    }
21
22    /// Write up to `N` bytes to this encoder.
23    ///
24    /// This function can be used to avoid the overhead of calling memcpy for writes that
25    /// have runtime-variable length, but are small and have a small fixed upper bound.
26    ///
27    /// This can be used to do in-place encoding as is done for leb128 (without this function
28    /// we would need to write to a temporary buffer then memcpy into the encoder), and it can
29    /// also be used to implement the varint scheme we use for rmeta and dep graph encoding,
30    /// where we only want to encode the first few bytes of an integer. Note that common
31    /// architectures support fixed-size writes up to 8 bytes with one instruction, so while this
32    /// does in some sense do wasted work, we come out ahead.
33    #[inline]
34    pub fn write_with<const N: usize>(&mut self, visitor: impl FnOnce(&mut [u8; N]) -> usize) {
35        self.data.reserve(N);
36
37        let old_len = self.data.len();
38
39        // SAFETY: The above `reserve` ensures that there is enough
40        // room to write the encoded value to the vector's internal buffer.
41        // The memory is also initialized as 0.
42        let buf = unsafe {
43            let buf = self.data.as_mut_ptr().add(old_len) as *mut [u8; N];
44            *buf = [0; N];
45            &mut *buf
46        };
47        let written = visitor(buf);
48        if written > N {
49            Self::panic_invalid_write::<N>(written);
50        }
51        unsafe { self.data.set_len(old_len + written) };
52    }
53
54    #[cold]
55    #[inline(never)]
56    fn panic_invalid_write<const N: usize>(written: usize) {
57        panic!("MemEncoder::write_with::<{N}> cannot be used to write {written} bytes");
58    }
59
60    /// Helper for calls where [`MemEncoder::write_with`] always writes the whole array.
61    #[inline]
62    pub fn write_array<const N: usize>(&mut self, buf: [u8; N]) {
63        self.write_with(|dest| {
64            *dest = buf;
65            N
66        })
67    }
68}
69
70macro_rules! write_leb128 {
71    ($this_fn:ident, $int_ty:ty, $write_leb_fn:ident) => {
72        #[inline]
73        fn $this_fn(&mut self, v: $int_ty) {
74            self.write_with(|buf| leb128::$write_leb_fn(buf, v))
75        }
76    };
77}
78
79impl Encoder for MemEncoder {
80    write_leb128!(emit_usize, usize, write_usize_leb128);
81    write_leb128!(emit_u128, u128, write_u128_leb128);
82    write_leb128!(emit_u64, u64, write_u64_leb128);
83    write_leb128!(emit_u32, u32, write_u32_leb128);
84
85    #[inline]
86    fn emit_u16(&mut self, v: u16) {
87        self.write_array(v.to_le_bytes());
88    }
89
90    #[inline]
91    fn emit_u8(&mut self, v: u8) {
92        self.write_array([v]);
93    }
94
95    write_leb128!(emit_isize, isize, write_isize_leb128);
96    write_leb128!(emit_i128, i128, write_i128_leb128);
97    write_leb128!(emit_i64, i64, write_i64_leb128);
98    write_leb128!(emit_i32, i32, write_i32_leb128);
99
100    #[inline]
101    fn emit_i16(&mut self, v: i16) {
102        self.write_array(v.to_le_bytes());
103    }
104
105    #[inline]
106    fn emit_raw_bytes(&mut self, s: &[u8]) {
107        self.data.extend_from_slice(s);
108    }
109}
110
111// Specialize encoding byte slices. This specialization also applies to encoding `Vec<u8>`s, etc.,
112// since the default implementations call `encode` on their slices internally.
113impl Encodable<MemEncoder> for [u8] {
114    fn encode(&self, e: &mut MemEncoder) {
115        Encoder::emit_usize(e, self.len());
116        e.emit_raw_bytes(self);
117    }
118}
119
120impl Encodable<MemEncoder> for IntEncodedWithFixedSize {
121    #[inline]
122    fn encode(&self, e: &mut MemEncoder) {
123        let start_pos = e.position();
124        e.write_array(self.0.to_le_bytes());
125        let end_pos = e.position();
126        debug_assert_eq!((end_pos - start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
127    }
128}