1 use crate::leb128::{self, largest_max_leb128_len};
2 use crate::serialize::{Decodable, Decoder, Encodable, Encoder};
4 use std::io::{self, Write};
5 use std::mem::MaybeUninit;
9 // -----------------------------------------------------------------------------
11 // -----------------------------------------------------------------------------
13 pub struct MemEncoder {
18 pub fn new() -> MemEncoder {
19 MemEncoder { data: vec![] }
23 pub fn position(&self) -> usize {
27 pub fn finish(self) -> Vec<u8> {
32 macro_rules! write_leb128 {
33 ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
34 const MAX_ENCODED_LEN: usize = $crate::leb128::max_leb128_len::<$int_ty>();
35 let old_len = $enc.data.len();
37 if MAX_ENCODED_LEN > $enc.data.capacity() - old_len {
38 $enc.data.reserve(MAX_ENCODED_LEN);
41 // SAFETY: The above check and `reserve` ensures that there is enough
42 // room to write the encoded value to the vector's internal buffer.
44 let buf = &mut *($enc.data.as_mut_ptr().add(old_len)
45 as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN]);
46 let encoded = leb128::$fun(buf, $value);
47 $enc.data.set_len(old_len + encoded.len());
52 /// A byte that [cannot occur in UTF8 sequences][utf8]. Used to mark the end of a string.
53 /// This way we can skip validation and still be relatively sure that deserialization
54 /// did not desynchronize.
56 /// [utf8]: https://en.wikipedia.org/w/index.php?title=UTF-8&oldid=1058865525#Codepage_layout
57 const STR_SENTINEL: u8 = 0xC1;
59 impl Encoder for MemEncoder {
61 fn emit_usize(&mut self, v: usize) {
62 write_leb128!(self, v, usize, write_usize_leb128)
66 fn emit_u128(&mut self, v: u128) {
67 write_leb128!(self, v, u128, write_u128_leb128);
71 fn emit_u64(&mut self, v: u64) {
72 write_leb128!(self, v, u64, write_u64_leb128);
76 fn emit_u32(&mut self, v: u32) {
77 write_leb128!(self, v, u32, write_u32_leb128);
81 fn emit_u16(&mut self, v: u16) {
82 self.data.extend_from_slice(&v.to_le_bytes());
86 fn emit_u8(&mut self, v: u8) {
91 fn emit_isize(&mut self, v: isize) {
92 write_leb128!(self, v, isize, write_isize_leb128)
96 fn emit_i128(&mut self, v: i128) {
97 write_leb128!(self, v, i128, write_i128_leb128)
101 fn emit_i64(&mut self, v: i64) {
102 write_leb128!(self, v, i64, write_i64_leb128)
106 fn emit_i32(&mut self, v: i32) {
107 write_leb128!(self, v, i32, write_i32_leb128)
111 fn emit_i16(&mut self, v: i16) {
112 self.data.extend_from_slice(&v.to_le_bytes());
116 fn emit_i8(&mut self, v: i8) {
117 self.emit_u8(v as u8);
121 fn emit_bool(&mut self, v: bool) {
122 self.emit_u8(if v { 1 } else { 0 });
126 fn emit_f64(&mut self, v: f64) {
127 let as_u64: u64 = v.to_bits();
128 self.emit_u64(as_u64);
132 fn emit_f32(&mut self, v: f32) {
133 let as_u32: u32 = v.to_bits();
134 self.emit_u32(as_u32);
138 fn emit_char(&mut self, v: char) {
139 self.emit_u32(v as u32);
143 fn emit_str(&mut self, v: &str) {
144 self.emit_usize(v.len());
145 self.emit_raw_bytes(v.as_bytes());
146 self.emit_u8(STR_SENTINEL);
150 fn emit_raw_bytes(&mut self, s: &[u8]) {
151 self.data.extend_from_slice(s);
155 pub type FileEncodeResult = Result<usize, io::Error>;
157 /// `FileEncoder` encodes data to file via fixed-size buffer.
159 /// When encoding large amounts of data to a file, using `FileEncoder` may be
160 /// preferred over using `MemEncoder` to encode to a `Vec`, and then writing the
161 /// `Vec` to file, as the latter uses as much memory as there is encoded data,
162 /// while the former uses the fixed amount of memory allocated to the buffer.
163 /// `FileEncoder` also has the advantage of not needing to reallocate as data
164 /// is appended to it, but the disadvantage of requiring more error handling,
165 /// which has some runtime overhead.
166 pub struct FileEncoder {
167 /// The input buffer. For adequate performance, we need more control over
168 /// buffering than `BufWriter` offers. If `BufWriter` ever offers a raw
169 /// buffer access API, we can use it, and remove `buf` and `buffered`.
170 buf: Box<[MaybeUninit<u8>]>,
174 // This is used to implement delayed error handling, as described in the
175 // comment on `trait Encoder`.
176 res: Result<(), io::Error>,
180 pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
181 const DEFAULT_BUF_SIZE: usize = 8192;
182 FileEncoder::with_capacity(path, DEFAULT_BUF_SIZE)
185 pub fn with_capacity<P: AsRef<Path>>(path: P, capacity: usize) -> io::Result<Self> {
186 // Require capacity at least as large as the largest LEB128 encoding
187 // here, so that we don't have to check or handle this on every write.
188 assert!(capacity >= largest_max_leb128_len());
190 // Require capacity small enough such that some capacity checks can be
191 // done using guaranteed non-overflowing add rather than sub, which
192 // shaves an instruction off those code paths (on x86 at least).
193 assert!(capacity <= usize::MAX - largest_max_leb128_len());
195 // Create the file for reading and writing, because some encoders do both
196 // (e.g. the metadata encoder when -Zmeta-stats is enabled)
197 let file = File::options().read(true).write(true).create(true).truncate(true).open(path)?;
200 buf: Box::new_uninit_slice(capacity),
209 pub fn position(&self) -> usize {
210 // Tracking position this way instead of having a `self.position` field
211 // means that we don't have to update the position on every write call.
212 self.flushed + self.buffered
215 pub fn flush(&mut self) {
216 // This is basically a copy of `BufWriter::flush`. If `BufWriter` ever
217 // offers a raw buffer access API, we can use it, and remove this.
219 /// Helper struct to ensure the buffer is updated after all the writes
220 /// are complete. It tracks the number of written bytes and drains them
221 /// all from the front of the buffer when dropped.
222 struct BufGuard<'a> {
223 buffer: &'a mut [u8],
224 encoder_buffered: &'a mut usize,
225 encoder_flushed: &'a mut usize,
229 impl<'a> BufGuard<'a> {
231 buffer: &'a mut [u8],
232 encoder_buffered: &'a mut usize,
233 encoder_flushed: &'a mut usize,
235 assert_eq!(buffer.len(), *encoder_buffered);
236 Self { buffer, encoder_buffered, encoder_flushed, flushed: 0 }
239 /// The unwritten part of the buffer
240 fn remaining(&self) -> &[u8] {
241 &self.buffer[self.flushed..]
244 /// Flag some bytes as removed from the front of the buffer
245 fn consume(&mut self, amt: usize) {
249 /// true if all of the bytes have been written
250 fn done(&self) -> bool {
251 self.flushed >= *self.encoder_buffered
255 impl Drop for BufGuard<'_> {
257 if self.flushed > 0 {
259 *self.encoder_flushed += *self.encoder_buffered;
260 *self.encoder_buffered = 0;
262 self.buffer.copy_within(self.flushed.., 0);
263 *self.encoder_flushed += self.flushed;
264 *self.encoder_buffered -= self.flushed;
270 // If we've already had an error, do nothing. It'll get reported after
271 // `finish` is called.
272 if self.res.is_err() {
276 let mut guard = BufGuard::new(
277 unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[..self.buffered]) },
282 while !guard.done() {
283 match self.file.write(guard.remaining()) {
285 self.res = Err(io::Error::new(
286 io::ErrorKind::WriteZero,
287 "failed to write the buffered data",
291 Ok(n) => guard.consume(n),
292 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
301 pub fn file(&self) -> &File {
306 fn capacity(&self) -> usize {
311 fn write_one(&mut self, value: u8) {
312 // We ensure this during `FileEncoder` construction.
313 debug_assert!(self.capacity() >= 1);
315 let mut buffered = self.buffered;
317 if std::intrinsics::unlikely(buffered >= self.capacity()) {
322 // SAFETY: The above check and `flush` ensures that there is enough
323 // room to write the input to the buffer.
325 *MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered) = value;
328 self.buffered = buffered + 1;
332 fn write_all(&mut self, buf: &[u8]) {
333 let capacity = self.capacity();
334 let buf_len = buf.len();
336 if std::intrinsics::likely(buf_len <= capacity) {
337 let mut buffered = self.buffered;
339 if std::intrinsics::unlikely(buf_len > capacity - buffered) {
344 // SAFETY: The above check and `flush` ensures that there is enough
345 // room to write the input to the buffer.
347 let src = buf.as_ptr();
348 let dst = MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered);
349 ptr::copy_nonoverlapping(src, dst, buf_len);
352 self.buffered = buffered + buf_len;
354 self.write_all_unbuffered(buf);
358 fn write_all_unbuffered(&mut self, mut buf: &[u8]) {
359 // If we've already had an error, do nothing. It'll get reported after
360 // `finish` is called.
361 if self.res.is_err() {
365 if self.buffered > 0 {
369 // This is basically a copy of `Write::write_all` but also updates our
370 // `self.flushed`. It's necessary because `Write::write_all` does not
371 // return the number of bytes written when an error is encountered, and
372 // without that, we cannot accurately update `self.flushed` on error.
373 while !buf.is_empty() {
374 match self.file.write(buf) {
376 self.res = Err(io::Error::new(
377 io::ErrorKind::WriteZero,
378 "failed to write whole buffer",
386 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
395 pub fn finish(mut self) -> Result<usize, io::Error> {
398 let res = std::mem::replace(&mut self.res, Ok(()));
399 res.map(|()| self.position())
403 impl Drop for FileEncoder {
405 // Likely to be a no-op, because `finish` should have been called and
406 // it also flushes. But do it just in case.
407 let _result = self.flush();
411 macro_rules! file_encoder_write_leb128 {
412 ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
413 const MAX_ENCODED_LEN: usize = $crate::leb128::max_leb128_len::<$int_ty>();
415 // We ensure this during `FileEncoder` construction.
416 debug_assert!($enc.capacity() >= MAX_ENCODED_LEN);
418 let mut buffered = $enc.buffered;
420 // This can't overflow. See assertion in `FileEncoder::with_capacity`.
421 if std::intrinsics::unlikely(buffered + MAX_ENCODED_LEN > $enc.capacity()) {
426 // SAFETY: The above check and flush ensures that there is enough
427 // room to write the encoded value to the buffer.
429 &mut *($enc.buf.as_mut_ptr().add(buffered) as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN])
432 let encoded = leb128::$fun(buf, $value);
433 $enc.buffered = buffered + encoded.len();
437 impl Encoder for FileEncoder {
439 fn emit_usize(&mut self, v: usize) {
440 file_encoder_write_leb128!(self, v, usize, write_usize_leb128)
444 fn emit_u128(&mut self, v: u128) {
445 file_encoder_write_leb128!(self, v, u128, write_u128_leb128)
449 fn emit_u64(&mut self, v: u64) {
450 file_encoder_write_leb128!(self, v, u64, write_u64_leb128)
454 fn emit_u32(&mut self, v: u32) {
455 file_encoder_write_leb128!(self, v, u32, write_u32_leb128)
459 fn emit_u16(&mut self, v: u16) {
460 self.write_all(&v.to_le_bytes());
464 fn emit_u8(&mut self, v: u8) {
469 fn emit_isize(&mut self, v: isize) {
470 file_encoder_write_leb128!(self, v, isize, write_isize_leb128)
474 fn emit_i128(&mut self, v: i128) {
475 file_encoder_write_leb128!(self, v, i128, write_i128_leb128)
479 fn emit_i64(&mut self, v: i64) {
480 file_encoder_write_leb128!(self, v, i64, write_i64_leb128)
484 fn emit_i32(&mut self, v: i32) {
485 file_encoder_write_leb128!(self, v, i32, write_i32_leb128)
489 fn emit_i16(&mut self, v: i16) {
490 self.write_all(&v.to_le_bytes());
494 fn emit_i8(&mut self, v: i8) {
495 self.emit_u8(v as u8);
499 fn emit_bool(&mut self, v: bool) {
500 self.emit_u8(if v { 1 } else { 0 });
504 fn emit_f64(&mut self, v: f64) {
505 let as_u64: u64 = v.to_bits();
506 self.emit_u64(as_u64);
510 fn emit_f32(&mut self, v: f32) {
511 let as_u32: u32 = v.to_bits();
512 self.emit_u32(as_u32);
516 fn emit_char(&mut self, v: char) {
517 self.emit_u32(v as u32);
521 fn emit_str(&mut self, v: &str) {
522 self.emit_usize(v.len());
523 self.emit_raw_bytes(v.as_bytes());
524 self.emit_u8(STR_SENTINEL);
528 fn emit_raw_bytes(&mut self, s: &[u8]) {
533 // -----------------------------------------------------------------------------
535 // -----------------------------------------------------------------------------
537 pub struct MemDecoder<'a> {
542 impl<'a> MemDecoder<'a> {
544 pub fn new(data: &'a [u8], position: usize) -> MemDecoder<'a> {
545 MemDecoder { data, position }
549 pub fn position(&self) -> usize {
554 pub fn set_position(&mut self, pos: usize) {
559 pub fn advance(&mut self, bytes: usize) {
560 self.position += bytes;
564 macro_rules! read_leb128 {
565 ($dec:expr, $fun:ident) => {{ leb128::$fun($dec.data, &mut $dec.position) }};
568 impl<'a> Decoder for MemDecoder<'a> {
570 fn read_u128(&mut self) -> u128 {
571 read_leb128!(self, read_u128_leb128)
575 fn read_u64(&mut self) -> u64 {
576 read_leb128!(self, read_u64_leb128)
580 fn read_u32(&mut self) -> u32 {
581 read_leb128!(self, read_u32_leb128)
585 fn read_u16(&mut self) -> u16 {
586 let bytes = [self.data[self.position], self.data[self.position + 1]];
587 let value = u16::from_le_bytes(bytes);
593 fn read_u8(&mut self) -> u8 {
594 let value = self.data[self.position];
600 fn read_usize(&mut self) -> usize {
601 read_leb128!(self, read_usize_leb128)
605 fn read_i128(&mut self) -> i128 {
606 read_leb128!(self, read_i128_leb128)
610 fn read_i64(&mut self) -> i64 {
611 read_leb128!(self, read_i64_leb128)
615 fn read_i32(&mut self) -> i32 {
616 read_leb128!(self, read_i32_leb128)
620 fn read_i16(&mut self) -> i16 {
621 let bytes = [self.data[self.position], self.data[self.position + 1]];
622 let value = i16::from_le_bytes(bytes);
628 fn read_i8(&mut self) -> i8 {
629 let value = self.data[self.position];
635 fn read_isize(&mut self) -> isize {
636 read_leb128!(self, read_isize_leb128)
640 fn read_bool(&mut self) -> bool {
641 let value = self.read_u8();
646 fn read_f64(&mut self) -> f64 {
647 let bits = self.read_u64();
652 fn read_f32(&mut self) -> f32 {
653 let bits = self.read_u32();
658 fn read_char(&mut self) -> char {
659 let bits = self.read_u32();
660 std::char::from_u32(bits).unwrap()
664 fn read_str(&mut self) -> &'a str {
665 let len = self.read_usize();
666 let sentinel = self.data[self.position + len];
667 assert!(sentinel == STR_SENTINEL);
669 std::str::from_utf8_unchecked(&self.data[self.position..self.position + len])
671 self.position += len + 1;
676 fn read_raw_bytes(&mut self, bytes: usize) -> &'a [u8] {
677 let start = self.position;
678 self.position += bytes;
679 &self.data[start..self.position]
683 // Specializations for contiguous byte sequences follow. The default implementations for slices
684 // encode and decode each element individually. This isn't necessary for `u8` slices when using
685 // opaque encoders and decoders, because each `u8` is unchanged by encoding and decoding.
686 // Therefore, we can use more efficient implementations that process the entire sequence at once.
688 // Specialize encoding byte slices. This specialization also applies to encoding `Vec<u8>`s, etc.,
689 // since the default implementations call `encode` on their slices internally.
690 impl Encodable<MemEncoder> for [u8] {
691 fn encode(&self, e: &mut MemEncoder) {
692 Encoder::emit_usize(e, self.len());
693 e.emit_raw_bytes(self);
697 impl Encodable<FileEncoder> for [u8] {
698 fn encode(&self, e: &mut FileEncoder) {
699 Encoder::emit_usize(e, self.len());
700 e.emit_raw_bytes(self);
704 // Specialize decoding `Vec<u8>`. This specialization also applies to decoding `Box<[u8]>`s, etc.,
705 // since the default implementations call `decode` to produce a `Vec<u8>` internally.
706 impl<'a> Decodable<MemDecoder<'a>> for Vec<u8> {
707 fn decode(d: &mut MemDecoder<'a>) -> Self {
708 let len = Decoder::read_usize(d);
709 d.read_raw_bytes(len).to_owned()
713 /// An integer that will always encode to 8 bytes.
714 pub struct IntEncodedWithFixedSize(pub u64);
716 impl IntEncodedWithFixedSize {
717 pub const ENCODED_SIZE: usize = 8;
720 impl Encodable<MemEncoder> for IntEncodedWithFixedSize {
722 fn encode(&self, e: &mut MemEncoder) {
723 let _start_pos = e.position();
724 e.emit_raw_bytes(&self.0.to_le_bytes());
725 let _end_pos = e.position();
726 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
730 impl Encodable<FileEncoder> for IntEncodedWithFixedSize {
732 fn encode(&self, e: &mut FileEncoder) {
733 let _start_pos = e.position();
734 e.emit_raw_bytes(&self.0.to_le_bytes());
735 let _end_pos = e.position();
736 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
740 impl<'a> Decodable<MemDecoder<'a>> for IntEncodedWithFixedSize {
742 fn decode(decoder: &mut MemDecoder<'a>) -> IntEncodedWithFixedSize {
743 let _start_pos = decoder.position();
744 let bytes = decoder.read_raw_bytes(IntEncodedWithFixedSize::ENCODED_SIZE);
745 let value = u64::from_le_bytes(bytes.try_into().unwrap());
746 let _end_pos = decoder.position();
747 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
749 IntEncodedWithFixedSize(value)