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

ra_ap_rustc_serialize/
opaque.rs
