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

rustc_serialize/
opaque.rs
