1 use crate::leb128::{self, max_leb128_len};
2 use crate::serialize::{self, Decoder as _, Encoder as _};
3 use std::convert::TryInto;
5 use std::io::{self, Write};
6 use std::mem::MaybeUninit;
10 // -----------------------------------------------------------------------------
12 // -----------------------------------------------------------------------------
14 pub type EncodeResult = Result<(), !>;
21 pub fn new(data: Vec<u8>) -> Encoder {
25 pub fn into_inner(self) -> Vec<u8> {
30 pub fn position(&self) -> usize {
35 macro_rules! write_leb128 {
36 ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
37 const MAX_ENCODED_LEN: usize = max_leb128_len!($int_ty);
38 let old_len = $enc.data.len();
40 if MAX_ENCODED_LEN > $enc.data.capacity() - old_len {
41 $enc.data.reserve(MAX_ENCODED_LEN);
44 // SAFETY: The above check and `reserve` ensures that there is enough
45 // room to write the encoded value to the vector's internal buffer.
47 let buf = &mut *($enc.data.as_mut_ptr().add(old_len)
48 as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN]);
49 let encoded = leb128::$fun(buf, $value);
50 $enc.data.set_len(old_len + encoded.len());
57 /// A byte that [cannot occur in UTF8 sequences][utf8]. Used to mark the end of a string.
58 /// This way we can skip validation and still be relatively sure that deserialization
59 /// did not desynchronize.
61 /// [utf8]: https://en.wikipedia.org/w/index.php?title=UTF-8&oldid=1058865525#Codepage_layout
62 const STR_SENTINEL: u8 = 0xC1;
64 impl serialize::Encoder for Encoder {
68 fn emit_usize(&mut self, v: usize) -> EncodeResult {
69 write_leb128!(self, v, usize, write_usize_leb128)
73 fn emit_u128(&mut self, v: u128) -> EncodeResult {
74 write_leb128!(self, v, u128, write_u128_leb128)
78 fn emit_u64(&mut self, v: u64) -> EncodeResult {
79 write_leb128!(self, v, u64, write_u64_leb128)
83 fn emit_u32(&mut self, v: u32) -> EncodeResult {
84 write_leb128!(self, v, u32, write_u32_leb128)
88 fn emit_u16(&mut self, v: u16) -> EncodeResult {
89 self.data.extend_from_slice(&v.to_le_bytes());
94 fn emit_u8(&mut self, v: u8) -> EncodeResult {
100 fn emit_isize(&mut self, v: isize) -> EncodeResult {
101 write_leb128!(self, v, isize, write_isize_leb128)
105 fn emit_i128(&mut self, v: i128) -> EncodeResult {
106 write_leb128!(self, v, i128, write_i128_leb128)
110 fn emit_i64(&mut self, v: i64) -> EncodeResult {
111 write_leb128!(self, v, i64, write_i64_leb128)
115 fn emit_i32(&mut self, v: i32) -> EncodeResult {
116 write_leb128!(self, v, i32, write_i32_leb128)
120 fn emit_i16(&mut self, v: i16) -> EncodeResult {
121 self.data.extend_from_slice(&v.to_le_bytes());
126 fn emit_i8(&mut self, v: i8) -> EncodeResult {
127 self.emit_u8(v as u8)
131 fn emit_bool(&mut self, v: bool) -> EncodeResult {
132 self.emit_u8(if v { 1 } else { 0 })
136 fn emit_f64(&mut self, v: f64) -> EncodeResult {
137 let as_u64: u64 = v.to_bits();
138 self.emit_u64(as_u64)
142 fn emit_f32(&mut self, v: f32) -> EncodeResult {
143 let as_u32: u32 = v.to_bits();
144 self.emit_u32(as_u32)
148 fn emit_char(&mut self, v: char) -> EncodeResult {
149 self.emit_u32(v as u32)
153 fn emit_str(&mut self, v: &str) -> EncodeResult {
154 self.emit_usize(v.len())?;
155 self.emit_raw_bytes(v.as_bytes())?;
156 self.emit_u8(STR_SENTINEL)
160 fn emit_raw_bytes(&mut self, s: &[u8]) -> EncodeResult {
161 self.data.extend_from_slice(s);
166 pub type FileEncodeResult = Result<(), io::Error>;
168 // `FileEncoder` encodes data to file via fixed-size buffer.
170 // When encoding large amounts of data to a file, using `FileEncoder` may be
171 // preferred over using `Encoder` to encode to a `Vec`, and then writing the
172 // `Vec` to file, as the latter uses as much memory as there is encoded data,
173 // while the former uses the fixed amount of memory allocated to the buffer.
174 // `FileEncoder` also has the advantage of not needing to reallocate as data
175 // is appended to it, but the disadvantage of requiring more error handling,
176 // which has some runtime overhead.
177 pub struct FileEncoder {
178 // The input buffer. For adequate performance, we need more control over
179 // buffering than `BufWriter` offers. If `BufWriter` ever offers a raw
180 // buffer access API, we can use it, and remove `buf` and `buffered`.
181 buf: Box<[MaybeUninit<u8>]>,
188 pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
189 const DEFAULT_BUF_SIZE: usize = 8192;
190 FileEncoder::with_capacity(path, DEFAULT_BUF_SIZE)
193 pub fn with_capacity<P: AsRef<Path>>(path: P, capacity: usize) -> io::Result<Self> {
194 // Require capacity at least as large as the largest LEB128 encoding
195 // here, so that we don't have to check or handle this on every write.
196 assert!(capacity >= max_leb128_len());
198 // Require capacity small enough such that some capacity checks can be
199 // done using guaranteed non-overflowing add rather than sub, which
200 // shaves an instruction off those code paths (on x86 at least).
201 assert!(capacity <= usize::MAX - max_leb128_len());
203 let file = File::create(path)?;
205 Ok(FileEncoder { buf: Box::new_uninit_slice(capacity), buffered: 0, flushed: 0, file })
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) -> FileEncodeResult {
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 let mut guard = BufGuard::new(
271 unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[..self.buffered]) },
276 while !guard.done() {
277 match self.file.write(guard.remaining()) {
279 return Err(io::Error::new(
280 io::ErrorKind::WriteZero,
281 "failed to write the buffered data",
284 Ok(n) => guard.consume(n),
285 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
286 Err(e) => return Err(e),
294 fn capacity(&self) -> usize {
299 fn write_one(&mut self, value: u8) -> FileEncodeResult {
300 // We ensure this during `FileEncoder` construction.
301 debug_assert!(self.capacity() >= 1);
303 let mut buffered = self.buffered;
305 if std::intrinsics::unlikely(buffered >= self.capacity()) {
310 // SAFETY: The above check and `flush` ensures that there is enough
311 // room to write the input to the buffer.
313 *MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered) = value;
316 self.buffered = buffered + 1;
322 fn write_all(&mut self, buf: &[u8]) -> FileEncodeResult {
323 let capacity = self.capacity();
324 let buf_len = buf.len();
326 if std::intrinsics::likely(buf_len <= capacity) {
327 let mut buffered = self.buffered;
329 if std::intrinsics::unlikely(buf_len > capacity - buffered) {
334 // SAFETY: The above check and `flush` ensures that there is enough
335 // room to write the input to the buffer.
337 let src = buf.as_ptr();
338 let dst = MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered);
339 ptr::copy_nonoverlapping(src, dst, buf_len);
342 self.buffered = buffered + buf_len;
346 self.write_all_unbuffered(buf)
350 fn write_all_unbuffered(&mut self, mut buf: &[u8]) -> FileEncodeResult {
351 if self.buffered > 0 {
355 // This is basically a copy of `Write::write_all` but also updates our
356 // `self.flushed`. It's necessary because `Write::write_all` does not
357 // return the number of bytes written when an error is encountered, and
358 // without that, we cannot accurately update `self.flushed` on error.
359 while !buf.is_empty() {
360 match self.file.write(buf) {
362 return Err(io::Error::new(
363 io::ErrorKind::WriteZero,
364 "failed to write whole buffer",
371 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
372 Err(e) => return Err(e),
380 impl Drop for FileEncoder {
382 let _result = self.flush();
386 macro_rules! file_encoder_write_leb128 {
387 ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
388 const MAX_ENCODED_LEN: usize = max_leb128_len!($int_ty);
390 // We ensure this during `FileEncoder` construction.
391 debug_assert!($enc.capacity() >= MAX_ENCODED_LEN);
393 let mut buffered = $enc.buffered;
395 // This can't overflow. See assertion in `FileEncoder::with_capacity`.
396 if std::intrinsics::unlikely(buffered + MAX_ENCODED_LEN > $enc.capacity()) {
401 // SAFETY: The above check and flush ensures that there is enough
402 // room to write the encoded value to the buffer.
404 &mut *($enc.buf.as_mut_ptr().add(buffered) as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN])
407 let encoded = leb128::$fun(buf, $value);
408 $enc.buffered = buffered + encoded.len();
414 impl serialize::Encoder for FileEncoder {
415 type Error = io::Error;
418 fn emit_usize(&mut self, v: usize) -> FileEncodeResult {
419 file_encoder_write_leb128!(self, v, usize, write_usize_leb128)
423 fn emit_u128(&mut self, v: u128) -> FileEncodeResult {
424 file_encoder_write_leb128!(self, v, u128, write_u128_leb128)
428 fn emit_u64(&mut self, v: u64) -> FileEncodeResult {
429 file_encoder_write_leb128!(self, v, u64, write_u64_leb128)
433 fn emit_u32(&mut self, v: u32) -> FileEncodeResult {
434 file_encoder_write_leb128!(self, v, u32, write_u32_leb128)
438 fn emit_u16(&mut self, v: u16) -> FileEncodeResult {
439 self.write_all(&v.to_le_bytes())
443 fn emit_u8(&mut self, v: u8) -> FileEncodeResult {
448 fn emit_isize(&mut self, v: isize) -> FileEncodeResult {
449 file_encoder_write_leb128!(self, v, isize, write_isize_leb128)
453 fn emit_i128(&mut self, v: i128) -> FileEncodeResult {
454 file_encoder_write_leb128!(self, v, i128, write_i128_leb128)
458 fn emit_i64(&mut self, v: i64) -> FileEncodeResult {
459 file_encoder_write_leb128!(self, v, i64, write_i64_leb128)
463 fn emit_i32(&mut self, v: i32) -> FileEncodeResult {
464 file_encoder_write_leb128!(self, v, i32, write_i32_leb128)
468 fn emit_i16(&mut self, v: i16) -> FileEncodeResult {
469 self.write_all(&v.to_le_bytes())
473 fn emit_i8(&mut self, v: i8) -> FileEncodeResult {
474 self.emit_u8(v as u8)
478 fn emit_bool(&mut self, v: bool) -> FileEncodeResult {
479 self.emit_u8(if v { 1 } else { 0 })
483 fn emit_f64(&mut self, v: f64) -> FileEncodeResult {
484 let as_u64: u64 = v.to_bits();
485 self.emit_u64(as_u64)
489 fn emit_f32(&mut self, v: f32) -> FileEncodeResult {
490 let as_u32: u32 = v.to_bits();
491 self.emit_u32(as_u32)
495 fn emit_char(&mut self, v: char) -> FileEncodeResult {
496 self.emit_u32(v as u32)
500 fn emit_str(&mut self, v: &str) -> FileEncodeResult {
501 self.emit_usize(v.len())?;
502 self.emit_raw_bytes(v.as_bytes())?;
503 self.emit_u8(STR_SENTINEL)
507 fn emit_raw_bytes(&mut self, s: &[u8]) -> FileEncodeResult {
512 // -----------------------------------------------------------------------------
514 // -----------------------------------------------------------------------------
516 pub struct Decoder<'a> {
521 impl<'a> Decoder<'a> {
523 pub fn new(data: &'a [u8], position: usize) -> Decoder<'a> {
524 Decoder { data, position }
528 pub fn position(&self) -> usize {
533 pub fn set_position(&mut self, pos: usize) {
538 pub fn advance(&mut self, bytes: usize) {
539 self.position += bytes;
543 macro_rules! read_leb128 {
544 ($dec:expr, $fun:ident) => {{ leb128::$fun($dec.data, &mut $dec.position) }};
547 impl<'a> serialize::Decoder for Decoder<'a> {
549 fn read_u128(&mut self) -> u128 {
550 read_leb128!(self, read_u128_leb128)
554 fn read_u64(&mut self) -> u64 {
555 read_leb128!(self, read_u64_leb128)
559 fn read_u32(&mut self) -> u32 {
560 read_leb128!(self, read_u32_leb128)
564 fn read_u16(&mut self) -> u16 {
565 let bytes = [self.data[self.position], self.data[self.position + 1]];
566 let value = u16::from_le_bytes(bytes);
572 fn read_u8(&mut self) -> u8 {
573 let value = self.data[self.position];
579 fn read_usize(&mut self) -> usize {
580 read_leb128!(self, read_usize_leb128)
584 fn read_i128(&mut self) -> i128 {
585 read_leb128!(self, read_i128_leb128)
589 fn read_i64(&mut self) -> i64 {
590 read_leb128!(self, read_i64_leb128)
594 fn read_i32(&mut self) -> i32 {
595 read_leb128!(self, read_i32_leb128)
599 fn read_i16(&mut self) -> i16 {
600 let bytes = [self.data[self.position], self.data[self.position + 1]];
601 let value = i16::from_le_bytes(bytes);
607 fn read_i8(&mut self) -> i8 {
608 let value = self.data[self.position];
614 fn read_isize(&mut self) -> isize {
615 read_leb128!(self, read_isize_leb128)
619 fn read_bool(&mut self) -> bool {
620 let value = self.read_u8();
625 fn read_f64(&mut self) -> f64 {
626 let bits = self.read_u64();
631 fn read_f32(&mut self) -> f32 {
632 let bits = self.read_u32();
637 fn read_char(&mut self) -> char {
638 let bits = self.read_u32();
639 std::char::from_u32(bits).unwrap()
643 fn read_str(&mut self) -> &'a str {
644 let len = self.read_usize();
645 let sentinel = self.data[self.position + len];
646 assert!(sentinel == STR_SENTINEL);
648 std::str::from_utf8_unchecked(&self.data[self.position..self.position + len])
650 self.position += len + 1;
655 fn read_raw_bytes(&mut self, bytes: usize) -> &'a [u8] {
656 let start = self.position;
657 self.position += bytes;
658 &self.data[start..self.position]
662 // Specializations for contiguous byte sequences follow. The default implementations for slices
663 // encode and decode each element individually. This isn't necessary for `u8` slices when using
664 // opaque encoders and decoders, because each `u8` is unchanged by encoding and decoding.
665 // Therefore, we can use more efficient implementations that process the entire sequence at once.
667 // Specialize encoding byte slices. This specialization also applies to encoding `Vec<u8>`s, etc.,
668 // since the default implementations call `encode` on their slices internally.
669 impl serialize::Encodable<Encoder> for [u8] {
670 fn encode(&self, e: &mut Encoder) -> EncodeResult {
671 serialize::Encoder::emit_usize(e, self.len())?;
672 e.emit_raw_bytes(self)
676 impl serialize::Encodable<FileEncoder> for [u8] {
677 fn encode(&self, e: &mut FileEncoder) -> FileEncodeResult {
678 serialize::Encoder::emit_usize(e, self.len())?;
679 e.emit_raw_bytes(self)
683 // Specialize decoding `Vec<u8>`. This specialization also applies to decoding `Box<[u8]>`s, etc.,
684 // since the default implementations call `decode` to produce a `Vec<u8>` internally.
685 impl<'a> serialize::Decodable<Decoder<'a>> for Vec<u8> {
686 fn decode(d: &mut Decoder<'a>) -> Self {
687 let len = serialize::Decoder::read_usize(d);
688 d.read_raw_bytes(len).to_owned()
692 // An integer that will always encode to 8 bytes.
693 pub struct IntEncodedWithFixedSize(pub u64);
695 impl IntEncodedWithFixedSize {
696 pub const ENCODED_SIZE: usize = 8;
699 impl serialize::Encodable<Encoder> for IntEncodedWithFixedSize {
701 fn encode(&self, e: &mut Encoder) -> EncodeResult {
702 let _start_pos = e.position();
703 e.emit_raw_bytes(&self.0.to_le_bytes())?;
704 let _end_pos = e.position();
705 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
710 impl serialize::Encodable<FileEncoder> for IntEncodedWithFixedSize {
712 fn encode(&self, e: &mut FileEncoder) -> FileEncodeResult {
713 let _start_pos = e.position();
714 e.emit_raw_bytes(&self.0.to_le_bytes())?;
715 let _end_pos = e.position();
716 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
721 impl<'a> serialize::Decodable<Decoder<'a>> for IntEncodedWithFixedSize {
723 fn decode(decoder: &mut Decoder<'a>) -> IntEncodedWithFixedSize {
724 let _start_pos = decoder.position();
725 let bytes = decoder.read_raw_bytes(IntEncodedWithFixedSize::ENCODED_SIZE);
726 let value = u64::from_le_bytes(bytes.try_into().unwrap());
727 let _end_pos = decoder.position();
728 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
730 IntEncodedWithFixedSize(value)