rustc_serialize/
opaque.rs

1use std::fs::File;
2use std::io::{self, Write};
3use std::marker::PhantomData;
4use std::ops::Range;
5use std::path::{Path, PathBuf};
6
7// This code is very hot and uses lots of arithmetic, avoid overflow checks for performance.
8// See https://github.com/rust-lang/rust/pull/119440#issuecomment-1874255727
9use crate::int_overflow::DebugStrictAdd;
10use crate::leb128;
11use crate::serialize::{Decodable, Decoder, Encodable, Encoder};
12
13pub mod mem_encoder;
14
15// -----------------------------------------------------------------------------
16// Encoder
17// -----------------------------------------------------------------------------
18
19pub type FileEncodeResult = Result<usize, (PathBuf, io::Error)>;
20
21pub const MAGIC_END_BYTES: &[u8] = b"rust-end-file";
22
23/// The size of the buffer in `FileEncoder`.
24const BUF_SIZE: usize = 64 * 1024;
25
26/// `FileEncoder` encodes data to file via fixed-size buffer.
27///
28/// There used to be a `MemEncoder` type that encoded all the data into a
29/// `Vec`. `FileEncoder` is better because its memory use is determined by the
30/// size of the buffer, rather than the full length of the encoded data, and
31/// because it doesn't need to reallocate memory along the way.
32pub struct FileEncoder {
33    // The input buffer. For adequate performance, we need to be able to write
34    // directly to the unwritten region of the buffer, without calling copy_from_slice.
35    // Note that our buffer is always initialized so that we can do that direct access
36    // without unsafe code. Users of this type write many more than BUF_SIZE bytes, so the
37    // initialization is approximately free.
38    buf: Box<[u8; BUF_SIZE]>,
39    buffered: usize,
40    flushed: usize,
41    file: File,
42    // This is used to implement delayed error handling, as described in the
43    // comment on `trait Encoder`.
44    res: Result<(), io::Error>,
45    path: PathBuf,
46    #[cfg(debug_assertions)]
47    finished: bool,
48}
49
50impl FileEncoder {
51    pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
52        // File::create opens the file for writing only. When -Zmeta-stats is enabled, the metadata
53        // encoder rewinds the file to inspect what was written. So we need to always open the file
54        // for reading and writing.
55        let file =
56            File::options().read(true).write(true).create(true).truncate(true).open(&path)?;
57
58        Ok(FileEncoder {
59            buf: vec![0u8; BUF_SIZE].into_boxed_slice().try_into().unwrap(),
60            path: path.as_ref().into(),
61            buffered: 0,
62            flushed: 0,
63            file,
64            res: Ok(()),
65            #[cfg(debug_assertions)]
66            finished: false,
67        })
68    }
69
70    #[inline]
71    pub fn position(&self) -> usize {
72        // Tracking position this way instead of having a `self.position` field
73        // means that we only need to update `self.buffered` on a write call,
74        // as opposed to updating `self.position` and `self.buffered`.
75        self.flushed.debug_strict_add(self.buffered)
76    }
77
78    #[cold]
79    #[inline(never)]
80    pub fn flush(&mut self) {
81        #[cfg(debug_assertions)]
82        {
83            self.finished = false;
84        }
85        if self.res.is_ok() {
86            self.res = self.file.write_all(&self.buf[..self.buffered]);
87        }
88        self.flushed += self.buffered;
89        self.buffered = 0;
90    }
91
92    pub fn file(&self) -> &File {
93        &self.file
94    }
95
96    pub fn path(&self) -> &Path {
97        &self.path
98    }
99
100    #[inline]
101    fn buffer_empty(&mut self) -> &mut [u8] {
102        // SAFETY: self.buffered is inbounds as an invariant of the type
103        unsafe { self.buf.get_unchecked_mut(self.buffered..) }
104    }
105
106    #[cold]
107    #[inline(never)]
108    fn write_all_cold_path(&mut self, buf: &[u8]) {
109        self.flush();
110        if let Some(dest) = self.buf.get_mut(..buf.len()) {
111            dest.copy_from_slice(buf);
112            self.buffered += buf.len();
113        } else {
114            if self.res.is_ok() {
115                self.res = self.file.write_all(buf);
116            }
117            self.flushed += buf.len();
118        }
119    }
120
121    #[inline]
122    fn write_all(&mut self, buf: &[u8]) {
123        #[cfg(debug_assertions)]
124        {
125            self.finished = false;
126        }
127        if let Some(dest) = self.buffer_empty().get_mut(..buf.len()) {
128            dest.copy_from_slice(buf);
129            self.buffered = self.buffered.debug_strict_add(buf.len());
130        } else {
131            self.write_all_cold_path(buf);
132        }
133    }
134
135    /// Write up to `N` bytes to this encoder.
136    ///
137    /// This function can be used to avoid the overhead of calling memcpy for writes that
138    /// have runtime-variable length, but are small and have a small fixed upper bound.
139    ///
140    /// This can be used to do in-place encoding as is done for leb128 (without this function
141    /// we would need to write to a temporary buffer then memcpy into the encoder), and it can
142    /// also be used to implement the varint scheme we use for rmeta and dep graph encoding,
143    /// where we only want to encode the first few bytes of an integer. Copying in the whole
144    /// integer then only advancing the encoder state for the few bytes we care about is more
145    /// efficient than calling [`FileEncoder::write_all`], because variable-size copies are
146    /// always lowered to `memcpy`, which has overhead and contains a lot of logic we can bypass
147    /// with this function. Note that common architectures support fixed-size writes up to 8 bytes
148    /// with one instruction, so while this does in some sense do wasted work, we come out ahead.
149    #[inline]
150    pub fn write_with<const N: usize>(&mut self, visitor: impl FnOnce(&mut [u8; N]) -> usize) {
151        #[cfg(debug_assertions)]
152        {
153            self.finished = false;
154        }
155        let flush_threshold = const { BUF_SIZE.checked_sub(N).unwrap() };
156        if std::intrinsics::unlikely(self.buffered > flush_threshold) {
157            self.flush();
158        }
159        // SAFETY: We checked above that N < self.buffer_empty().len(),
160        // and if isn't, flush ensures that our empty buffer is now BUF_SIZE.
161        // We produce a post-mono error if N > BUF_SIZE.
162        let buf = unsafe { self.buffer_empty().first_chunk_mut::<N>().unwrap_unchecked() };
163        let written = visitor(buf);
164        // We have to ensure that an errant visitor cannot cause self.buffered to exceed BUF_SIZE.
165        if written > N {
166            Self::panic_invalid_write::<N>(written);
167        }
168        self.buffered = self.buffered.debug_strict_add(written);
169    }
170
171    #[cold]
172    #[inline(never)]
173    fn panic_invalid_write<const N: usize>(written: usize) {
174        panic!("FileEncoder::write_with::<{N}> cannot be used to write {written} bytes");
175    }
176
177    /// Helper for calls where [`FileEncoder::write_with`] always writes the whole array.
178    #[inline]
179    pub fn write_array<const N: usize>(&mut self, buf: [u8; N]) {
180        self.write_with(|dest| {
181            *dest = buf;
182            N
183        })
184    }
185
186    pub fn finish(&mut self) -> FileEncodeResult {
187        self.write_all(MAGIC_END_BYTES);
188        self.flush();
189        #[cfg(debug_assertions)]
190        {
191            self.finished = true;
192        }
193        match std::mem::replace(&mut self.res, Ok(())) {
194            Ok(()) => Ok(self.position()),
195            Err(e) => Err((self.path.clone(), e)),
196        }
197    }
198}
199
200#[cfg(debug_assertions)]
201impl Drop for FileEncoder {
202    fn drop(&mut self) {
203        if !std::thread::panicking() {
204            assert!(self.finished);
205        }
206    }
207}
208
209macro_rules! write_leb128 {
210    ($this_fn:ident, $int_ty:ty, $write_leb_fn:ident) => {
211        #[inline]
212        fn $this_fn(&mut self, v: $int_ty) {
213            self.write_with(|buf| leb128::$write_leb_fn(buf, v))
214        }
215    };
216}
217
218impl Encoder for FileEncoder {
219    write_leb128!(emit_usize, usize, write_usize_leb128);
220    write_leb128!(emit_u128, u128, write_u128_leb128);
221    write_leb128!(emit_u64, u64, write_u64_leb128);
222    write_leb128!(emit_u32, u32, write_u32_leb128);
223
224    #[inline]
225    fn emit_u16(&mut self, v: u16) {
226        self.write_array(v.to_le_bytes());
227    }
228
229    #[inline]
230    fn emit_u8(&mut self, v: u8) {
231        self.write_array([v]);
232    }
233
234    write_leb128!(emit_isize, isize, write_isize_leb128);
235    write_leb128!(emit_i128, i128, write_i128_leb128);
236    write_leb128!(emit_i64, i64, write_i64_leb128);
237    write_leb128!(emit_i32, i32, write_i32_leb128);
238
239    #[inline]
240    fn emit_i16(&mut self, v: i16) {
241        self.write_array(v.to_le_bytes());
242    }
243
244    #[inline]
245    fn emit_raw_bytes(&mut self, s: &[u8]) {
246        self.write_all(s);
247    }
248}
249
250// -----------------------------------------------------------------------------
251// Decoder
252// -----------------------------------------------------------------------------
253
254// Conceptually, `MemDecoder` wraps a `&[u8]` with a cursor into it that is always valid.
255// This is implemented with three pointers, two which represent the original slice and a
256// third that is our cursor.
257// It is an invariant of this type that start <= current <= end.
258// Additionally, the implementation of this type never modifies start and end.
259pub struct MemDecoder<'a> {
260    start: *const u8,
261    current: *const u8,
262    end: *const u8,
263    _marker: PhantomData<&'a u8>,
264}
265
266impl<'a> MemDecoder<'a> {
267    #[inline]
268    pub fn new(data: &'a [u8], position: usize) -> Result<MemDecoder<'a>, ()> {
269        let data = data.strip_suffix(MAGIC_END_BYTES).ok_or(())?;
270        let Range { start, end } = data.as_ptr_range();
271        Ok(MemDecoder { start, current: data[position..].as_ptr(), end, _marker: PhantomData })
272    }
273
274    #[inline]
275    pub fn split_at(&self, position: usize) -> MemDecoder<'a> {
276        assert!(position <= self.len());
277        // SAFETY: We checked above that this offset is within the original slice
278        let current = unsafe { self.start.add(position) };
279        MemDecoder { start: self.start, current, end: self.end, _marker: PhantomData }
280    }
281
282    #[inline]
283    pub fn len(&self) -> usize {
284        // SAFETY: This recovers the length of the original slice, only using members we never modify.
285        unsafe { self.end.offset_from_unsigned(self.start) }
286    }
287
288    #[inline]
289    pub fn remaining(&self) -> usize {
290        // SAFETY: This type guarantees current <= end.
291        unsafe { self.end.offset_from_unsigned(self.current) }
292    }
293
294    #[cold]
295    #[inline(never)]
296    fn decoder_exhausted() -> ! {
297        panic!("MemDecoder exhausted")
298    }
299
300    #[inline]
301    pub fn read_array<const N: usize>(&mut self) -> [u8; N] {
302        self.read_raw_bytes(N).try_into().unwrap()
303    }
304
305    /// While we could manually expose manipulation of the decoder position,
306    /// all current users of that method would need to reset the position later,
307    /// incurring the bounds check of set_position twice.
308    #[inline]
309    pub fn with_position<F, T>(&mut self, pos: usize, func: F) -> T
310    where
311        F: Fn(&mut MemDecoder<'a>) -> T,
312    {
313        struct SetOnDrop<'a, 'guarded> {
314            decoder: &'guarded mut MemDecoder<'a>,
315            current: *const u8,
316        }
317        impl Drop for SetOnDrop<'_, '_> {
318            fn drop(&mut self) {
319                self.decoder.current = self.current;
320            }
321        }
322
323        if pos >= self.len() {
324            Self::decoder_exhausted();
325        }
326        let previous = self.current;
327        // SAFETY: We just checked if this add is in-bounds above.
328        unsafe {
329            self.current = self.start.add(pos);
330        }
331        let guard = SetOnDrop { current: previous, decoder: self };
332        func(guard.decoder)
333    }
334}
335
336macro_rules! read_leb128 {
337    ($this_fn:ident, $int_ty:ty, $read_leb_fn:ident) => {
338        #[inline]
339        fn $this_fn(&mut self) -> $int_ty {
340            leb128::$read_leb_fn(self)
341        }
342    };
343}
344
345impl<'a> Decoder for MemDecoder<'a> {
346    read_leb128!(read_usize, usize, read_usize_leb128);
347    read_leb128!(read_u128, u128, read_u128_leb128);
348    read_leb128!(read_u64, u64, read_u64_leb128);
349    read_leb128!(read_u32, u32, read_u32_leb128);
350
351    #[inline]
352    fn read_u16(&mut self) -> u16 {
353        u16::from_le_bytes(self.read_array())
354    }
355
356    #[inline]
357    fn read_u8(&mut self) -> u8 {
358        if self.current == self.end {
359            Self::decoder_exhausted();
360        }
361        // SAFETY: This type guarantees current <= end, and we just checked current == end.
362        unsafe {
363            let byte = *self.current;
364            self.current = self.current.add(1);
365            byte
366        }
367    }
368
369    read_leb128!(read_isize, isize, read_isize_leb128);
370    read_leb128!(read_i128, i128, read_i128_leb128);
371    read_leb128!(read_i64, i64, read_i64_leb128);
372    read_leb128!(read_i32, i32, read_i32_leb128);
373
374    #[inline]
375    fn read_i16(&mut self) -> i16 {
376        i16::from_le_bytes(self.read_array())
377    }
378
379    #[inline]
380    fn read_raw_bytes(&mut self, bytes: usize) -> &'a [u8] {
381        if bytes > self.remaining() {
382            Self::decoder_exhausted();
383        }
384        // SAFETY: We just checked if this range is in-bounds above.
385        unsafe {
386            let slice = std::slice::from_raw_parts(self.current, bytes);
387            self.current = self.current.add(bytes);
388            slice
389        }
390    }
391
392    #[inline]
393    fn peek_byte(&self) -> u8 {
394        if self.current == self.end {
395            Self::decoder_exhausted();
396        }
397        // SAFETY: This type guarantees current is inbounds or one-past-the-end, which is end.
398        // Since we just checked current == end, the current pointer must be inbounds.
399        unsafe { *self.current }
400    }
401
402    #[inline]
403    fn position(&self) -> usize {
404        // SAFETY: This type guarantees start <= current
405        unsafe { self.current.offset_from_unsigned(self.start) }
406    }
407}
408
409// Specializations for contiguous byte sequences follow. The default implementations for slices
410// encode and decode each element individually. This isn't necessary for `u8` slices when using
411// opaque encoders and decoders, because each `u8` is unchanged by encoding and decoding.
412// Therefore, we can use more efficient implementations that process the entire sequence at once.
413
414// Specialize encoding byte slices. This specialization also applies to encoding `Vec<u8>`s, etc.,
415// since the default implementations call `encode` on their slices internally.
416impl Encodable<FileEncoder> for [u8] {
417    fn encode(&self, e: &mut FileEncoder) {
418        Encoder::emit_usize(e, self.len());
419        e.emit_raw_bytes(self);
420    }
421}
422
423// Specialize decoding `Vec<u8>`. This specialization also applies to decoding `Box<[u8]>`s, etc.,
424// since the default implementations call `decode` to produce a `Vec<u8>` internally.
425impl<'a> Decodable<MemDecoder<'a>> for Vec<u8> {
426    fn decode(d: &mut MemDecoder<'a>) -> Self {
427        let len = Decoder::read_usize(d);
428        d.read_raw_bytes(len).to_owned()
429    }
430}
431
432/// An integer that will always encode to 8 bytes.
433pub struct IntEncodedWithFixedSize(pub u64);
434
435impl IntEncodedWithFixedSize {
436    pub const ENCODED_SIZE: usize = 8;
437}
438
439impl Encodable<FileEncoder> for IntEncodedWithFixedSize {
440    #[inline]
441    fn encode(&self, e: &mut FileEncoder) {
442        let start_pos = e.position();
443        e.write_array(self.0.to_le_bytes());
444        let end_pos = e.position();
445        debug_assert_eq!((end_pos - start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
446    }
447}
448
449impl<'a> Decodable<MemDecoder<'a>> for IntEncodedWithFixedSize {
450    #[inline]
451    fn decode(decoder: &mut MemDecoder<'a>) -> IntEncodedWithFixedSize {
452        let bytes = decoder.read_array::<{ IntEncodedWithFixedSize::ENCODED_SIZE }>();
453        IntEncodedWithFixedSize(u64::from_le_bytes(bytes))
454    }
455}
456
457#[cfg(test)]
458mod tests;
rustc_serialize/opaque.rs

rustc_serialize/
opaque.rs
