1 use crate::leb128::{self, max_leb128_len};
2 use crate::serialize::{self, Encoder as _};
4 use std::convert::TryInto;
6 use std::io::{self, Write};
7 use std::mem::MaybeUninit;
11 // -----------------------------------------------------------------------------
13 // -----------------------------------------------------------------------------
15 pub type EncodeResult = Result<(), !>;
22 pub fn new(data: Vec<u8>) -> Encoder {
26 pub fn into_inner(self) -> Vec<u8> {
31 pub fn position(&self) -> usize {
36 macro_rules! write_leb128 {
37 ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
38 const MAX_ENCODED_LEN: usize = max_leb128_len!($int_ty);
39 let old_len = $enc.data.len();
41 if MAX_ENCODED_LEN > $enc.data.capacity() - old_len {
42 $enc.data.reserve(MAX_ENCODED_LEN);
45 // SAFETY: The above check and `reserve` ensures that there is enough
46 // room to write the encoded value to the vector's internal buffer.
48 let buf = &mut *($enc.data.as_mut_ptr().add(old_len)
49 as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN]);
50 let encoded = leb128::$fun(buf, $value);
51 $enc.data.set_len(old_len + encoded.len());
58 impl serialize::Encoder for Encoder {
62 fn emit_unit(&mut self) -> EncodeResult {
67 fn emit_usize(&mut self, v: usize) -> EncodeResult {
68 write_leb128!(self, v, usize, write_usize_leb128)
72 fn emit_u128(&mut self, v: u128) -> EncodeResult {
73 write_leb128!(self, v, u128, write_u128_leb128)
77 fn emit_u64(&mut self, v: u64) -> EncodeResult {
78 write_leb128!(self, v, u64, write_u64_leb128)
82 fn emit_u32(&mut self, v: u32) -> EncodeResult {
83 write_leb128!(self, v, u32, write_u32_leb128)
87 fn emit_u16(&mut self, v: u16) -> EncodeResult {
88 write_leb128!(self, v, u16, write_u16_leb128)
92 fn emit_u8(&mut self, v: u8) -> EncodeResult {
98 fn emit_isize(&mut self, v: isize) -> EncodeResult {
99 write_leb128!(self, v, isize, write_isize_leb128)
103 fn emit_i128(&mut self, v: i128) -> EncodeResult {
104 write_leb128!(self, v, i128, write_i128_leb128)
108 fn emit_i64(&mut self, v: i64) -> EncodeResult {
109 write_leb128!(self, v, i64, write_i64_leb128)
113 fn emit_i32(&mut self, v: i32) -> EncodeResult {
114 write_leb128!(self, v, i32, write_i32_leb128)
118 fn emit_i16(&mut self, v: i16) -> EncodeResult {
119 write_leb128!(self, v, i16, write_i16_leb128)
123 fn emit_i8(&mut self, v: i8) -> EncodeResult {
124 let as_u8: u8 = unsafe { std::mem::transmute(v) };
129 fn emit_bool(&mut self, v: bool) -> EncodeResult {
130 self.emit_u8(if v { 1 } else { 0 })
134 fn emit_f64(&mut self, v: f64) -> EncodeResult {
135 let as_u64: u64 = v.to_bits();
136 self.emit_u64(as_u64)
140 fn emit_f32(&mut self, v: f32) -> EncodeResult {
141 let as_u32: u32 = v.to_bits();
142 self.emit_u32(as_u32)
146 fn emit_char(&mut self, v: char) -> EncodeResult {
147 self.emit_u32(v as u32)
151 fn emit_str(&mut self, v: &str) -> EncodeResult {
152 self.emit_usize(v.len())?;
153 self.emit_raw_bytes(v.as_bytes())
157 fn emit_raw_bytes(&mut self, s: &[u8]) -> EncodeResult {
158 self.data.extend_from_slice(s);
163 pub type FileEncodeResult = Result<(), io::Error>;
165 // `FileEncoder` encodes data to file via fixed-size buffer.
167 // When encoding large amounts of data to a file, using `FileEncoder` may be
168 // preferred over using `Encoder` to encode to a `Vec`, and then writing the
169 // `Vec` to file, as the latter uses as much memory as there is encoded data,
170 // while the former uses the fixed amount of memory allocated to the buffer.
171 // `FileEncoder` also has the advantage of not needing to reallocate as data
172 // is appended to it, but the disadvantage of requiring more error handling,
173 // which has some runtime overhead.
174 pub struct FileEncoder {
175 // The input buffer. For adequate performance, we need more control over
176 // buffering than `BufWriter` offers. If `BufWriter` ever offers a raw
177 // buffer access API, we can use it, and remove `buf` and `buffered`.
178 buf: Box<[MaybeUninit<u8>]>,
185 pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
186 const DEFAULT_BUF_SIZE: usize = 8192;
187 FileEncoder::with_capacity(path, DEFAULT_BUF_SIZE)
190 pub fn with_capacity<P: AsRef<Path>>(path: P, capacity: usize) -> io::Result<Self> {
191 // Require capacity at least as large as the largest LEB128 encoding
192 // here, so that we don't have to check or handle this on every write.
193 assert!(capacity >= max_leb128_len());
195 // Require capacity small enough such that some capacity checks can be
196 // done using guaranteed non-overflowing add rather than sub, which
197 // shaves an instruction off those code paths (on x86 at least).
198 assert!(capacity <= usize::MAX - max_leb128_len());
200 let file = File::create(path)?;
202 Ok(FileEncoder { buf: Box::new_uninit_slice(capacity), buffered: 0, flushed: 0, file })
206 pub fn position(&self) -> usize {
207 // Tracking position this way instead of having a `self.position` field
208 // means that we don't have to update the position on every write call.
209 self.flushed + self.buffered
212 pub fn flush(&mut self) -> FileEncodeResult {
213 // This is basically a copy of `BufWriter::flush`. If `BufWriter` ever
214 // offers a raw buffer access API, we can use it, and remove this.
216 /// Helper struct to ensure the buffer is updated after all the writes
217 /// are complete. It tracks the number of written bytes and drains them
218 /// all from the front of the buffer when dropped.
219 struct BufGuard<'a> {
220 buffer: &'a mut [u8],
221 encoder_buffered: &'a mut usize,
222 encoder_flushed: &'a mut usize,
226 impl<'a> BufGuard<'a> {
228 buffer: &'a mut [u8],
229 encoder_buffered: &'a mut usize,
230 encoder_flushed: &'a mut usize,
232 assert_eq!(buffer.len(), *encoder_buffered);
233 Self { buffer, encoder_buffered, encoder_flushed, flushed: 0 }
236 /// The unwritten part of the buffer
237 fn remaining(&self) -> &[u8] {
238 &self.buffer[self.flushed..]
241 /// Flag some bytes as removed from the front of the buffer
242 fn consume(&mut self, amt: usize) {
246 /// true if all of the bytes have been written
247 fn done(&self) -> bool {
248 self.flushed >= *self.encoder_buffered
252 impl Drop for BufGuard<'_> {
254 if self.flushed > 0 {
256 *self.encoder_flushed += *self.encoder_buffered;
257 *self.encoder_buffered = 0;
259 self.buffer.copy_within(self.flushed.., 0);
260 *self.encoder_flushed += self.flushed;
261 *self.encoder_buffered -= self.flushed;
267 let mut guard = BufGuard::new(
268 unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[..self.buffered]) },
273 while !guard.done() {
274 match self.file.write(guard.remaining()) {
276 return Err(io::Error::new(
277 io::ErrorKind::WriteZero,
278 "failed to write the buffered data",
281 Ok(n) => guard.consume(n),
282 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
283 Err(e) => return Err(e),
291 fn capacity(&self) -> usize {
296 fn write_one(&mut self, value: u8) -> FileEncodeResult {
297 // We ensure this during `FileEncoder` construction.
298 debug_assert!(self.capacity() >= 1);
300 let mut buffered = self.buffered;
302 if std::intrinsics::unlikely(buffered >= self.capacity()) {
307 // SAFETY: The above check and `flush` ensures that there is enough
308 // room to write the input to the buffer.
310 *MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered) = value;
313 self.buffered = buffered + 1;
319 fn write_all(&mut self, buf: &[u8]) -> FileEncodeResult {
320 let capacity = self.capacity();
321 let buf_len = buf.len();
323 if std::intrinsics::likely(buf_len <= capacity) {
324 let mut buffered = self.buffered;
326 if std::intrinsics::unlikely(buf_len > capacity - buffered) {
331 // SAFETY: The above check and `flush` ensures that there is enough
332 // room to write the input to the buffer.
334 let src = buf.as_ptr();
335 let dst = MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered);
336 ptr::copy_nonoverlapping(src, dst, buf_len);
339 self.buffered = buffered + buf_len;
343 self.write_all_unbuffered(buf)
347 fn write_all_unbuffered(&mut self, mut buf: &[u8]) -> FileEncodeResult {
348 if self.buffered > 0 {
352 // This is basically a copy of `Write::write_all` but also updates our
353 // `self.flushed`. It's necessary because `Write::write_all` does not
354 // return the number of bytes written when an error is encountered, and
355 // without that, we cannot accurately update `self.flushed` on error.
356 while !buf.is_empty() {
357 match self.file.write(buf) {
359 return Err(io::Error::new(
360 io::ErrorKind::WriteZero,
361 "failed to write whole buffer",
368 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
369 Err(e) => return Err(e),
377 impl Drop for FileEncoder {
379 let _result = self.flush();
383 macro_rules! file_encoder_write_leb128 {
384 ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
385 const MAX_ENCODED_LEN: usize = max_leb128_len!($int_ty);
387 // We ensure this during `FileEncoder` construction.
388 debug_assert!($enc.capacity() >= MAX_ENCODED_LEN);
390 let mut buffered = $enc.buffered;
392 // This can't overflow. See assertion in `FileEncoder::with_capacity`.
393 if std::intrinsics::unlikely(buffered + MAX_ENCODED_LEN > $enc.capacity()) {
398 // SAFETY: The above check and flush ensures that there is enough
399 // room to write the encoded value to the buffer.
401 &mut *($enc.buf.as_mut_ptr().add(buffered) as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN])
404 let encoded = leb128::$fun(buf, $value);
405 $enc.buffered = buffered + encoded.len();
411 impl serialize::Encoder for FileEncoder {
412 type Error = io::Error;
415 fn emit_unit(&mut self) -> FileEncodeResult {
420 fn emit_usize(&mut self, v: usize) -> FileEncodeResult {
421 file_encoder_write_leb128!(self, v, usize, write_usize_leb128)
425 fn emit_u128(&mut self, v: u128) -> FileEncodeResult {
426 file_encoder_write_leb128!(self, v, u128, write_u128_leb128)
430 fn emit_u64(&mut self, v: u64) -> FileEncodeResult {
431 file_encoder_write_leb128!(self, v, u64, write_u64_leb128)
435 fn emit_u32(&mut self, v: u32) -> FileEncodeResult {
436 file_encoder_write_leb128!(self, v, u32, write_u32_leb128)
440 fn emit_u16(&mut self, v: u16) -> FileEncodeResult {
441 file_encoder_write_leb128!(self, v, u16, write_u16_leb128)
445 fn emit_u8(&mut self, v: u8) -> FileEncodeResult {
450 fn emit_isize(&mut self, v: isize) -> FileEncodeResult {
451 file_encoder_write_leb128!(self, v, isize, write_isize_leb128)
455 fn emit_i128(&mut self, v: i128) -> FileEncodeResult {
456 file_encoder_write_leb128!(self, v, i128, write_i128_leb128)
460 fn emit_i64(&mut self, v: i64) -> FileEncodeResult {
461 file_encoder_write_leb128!(self, v, i64, write_i64_leb128)
465 fn emit_i32(&mut self, v: i32) -> FileEncodeResult {
466 file_encoder_write_leb128!(self, v, i32, write_i32_leb128)
470 fn emit_i16(&mut self, v: i16) -> FileEncodeResult {
471 file_encoder_write_leb128!(self, v, i16, write_i16_leb128)
475 fn emit_i8(&mut self, v: i8) -> FileEncodeResult {
476 let as_u8: u8 = unsafe { std::mem::transmute(v) };
481 fn emit_bool(&mut self, v: bool) -> FileEncodeResult {
482 self.emit_u8(if v { 1 } else { 0 })
486 fn emit_f64(&mut self, v: f64) -> FileEncodeResult {
487 let as_u64: u64 = v.to_bits();
488 self.emit_u64(as_u64)
492 fn emit_f32(&mut self, v: f32) -> FileEncodeResult {
493 let as_u32: u32 = v.to_bits();
494 self.emit_u32(as_u32)
498 fn emit_char(&mut self, v: char) -> FileEncodeResult {
499 self.emit_u32(v as u32)
503 fn emit_str(&mut self, v: &str) -> FileEncodeResult {
504 self.emit_usize(v.len())?;
505 self.emit_raw_bytes(v.as_bytes())
509 fn emit_raw_bytes(&mut self, s: &[u8]) -> FileEncodeResult {
514 // -----------------------------------------------------------------------------
516 // -----------------------------------------------------------------------------
518 pub struct Decoder<'a> {
523 impl<'a> Decoder<'a> {
525 pub fn new(data: &'a [u8], position: usize) -> Decoder<'a> {
526 Decoder { data, position }
530 pub fn position(&self) -> usize {
535 pub fn set_position(&mut self, pos: usize) {
540 pub fn advance(&mut self, bytes: usize) {
541 self.position += bytes;
545 pub fn read_raw_bytes(&mut self, bytes: usize) -> &'a [u8] {
546 let start = self.position;
547 self.position += bytes;
548 &self.data[start..self.position]
552 macro_rules! read_leb128 {
553 ($dec:expr, $fun:ident) => {{
554 let (value, bytes_read) = leb128::$fun(&$dec.data[$dec.position..]);
555 $dec.position += bytes_read;
560 impl<'a> serialize::Decoder for Decoder<'a> {
564 fn read_nil(&mut self) -> Result<(), Self::Error> {
569 fn read_u128(&mut self) -> Result<u128, Self::Error> {
570 read_leb128!(self, read_u128_leb128)
574 fn read_u64(&mut self) -> Result<u64, Self::Error> {
575 read_leb128!(self, read_u64_leb128)
579 fn read_u32(&mut self) -> Result<u32, Self::Error> {
580 read_leb128!(self, read_u32_leb128)
584 fn read_u16(&mut self) -> Result<u16, Self::Error> {
585 read_leb128!(self, read_u16_leb128)
589 fn read_u8(&mut self) -> Result<u8, Self::Error> {
590 let value = self.data[self.position];
596 fn read_usize(&mut self) -> Result<usize, Self::Error> {
597 read_leb128!(self, read_usize_leb128)
601 fn read_i128(&mut self) -> Result<i128, Self::Error> {
602 read_leb128!(self, read_i128_leb128)
606 fn read_i64(&mut self) -> Result<i64, Self::Error> {
607 read_leb128!(self, read_i64_leb128)
611 fn read_i32(&mut self) -> Result<i32, Self::Error> {
612 read_leb128!(self, read_i32_leb128)
616 fn read_i16(&mut self) -> Result<i16, Self::Error> {
617 read_leb128!(self, read_i16_leb128)
621 fn read_i8(&mut self) -> Result<i8, Self::Error> {
622 let as_u8 = self.data[self.position];
624 unsafe { Ok(::std::mem::transmute(as_u8)) }
628 fn read_isize(&mut self) -> Result<isize, Self::Error> {
629 read_leb128!(self, read_isize_leb128)
633 fn read_bool(&mut self) -> Result<bool, Self::Error> {
634 let value = self.read_u8()?;
639 fn read_f64(&mut self) -> Result<f64, Self::Error> {
640 let bits = self.read_u64()?;
641 Ok(f64::from_bits(bits))
645 fn read_f32(&mut self) -> Result<f32, Self::Error> {
646 let bits = self.read_u32()?;
647 Ok(f32::from_bits(bits))
651 fn read_char(&mut self) -> Result<char, Self::Error> {
652 let bits = self.read_u32()?;
653 Ok(std::char::from_u32(bits).unwrap())
657 fn read_str(&mut self) -> Result<Cow<'_, str>, Self::Error> {
658 let len = self.read_usize()?;
659 let s = std::str::from_utf8(&self.data[self.position..self.position + len]).unwrap();
660 self.position += len;
665 fn error(&mut self, err: &str) -> Self::Error {
670 fn read_raw_bytes_into(&mut self, s: &mut [u8]) -> Result<(), String> {
671 let start = self.position;
672 self.position += s.len();
673 s.copy_from_slice(&self.data[start..self.position]);
678 // Specializations for contiguous byte sequences follow. The default implementations for slices
679 // encode and decode each element individually. This isn't necessary for `u8` slices when using
680 // opaque encoders and decoders, because each `u8` is unchanged by encoding and decoding.
681 // Therefore, we can use more efficient implementations that process the entire sequence at once.
683 // Specialize encoding byte slices. This specialization also applies to encoding `Vec<u8>`s, etc.,
684 // since the default implementations call `encode` on their slices internally.
685 impl serialize::Encodable<Encoder> for [u8] {
686 fn encode(&self, e: &mut Encoder) -> EncodeResult {
687 serialize::Encoder::emit_usize(e, self.len())?;
688 e.emit_raw_bytes(self)
692 impl serialize::Encodable<FileEncoder> for [u8] {
693 fn encode(&self, e: &mut FileEncoder) -> FileEncodeResult {
694 serialize::Encoder::emit_usize(e, self.len())?;
695 e.emit_raw_bytes(self)
699 // Specialize decoding `Vec<u8>`. This specialization also applies to decoding `Box<[u8]>`s, etc.,
700 // since the default implementations call `decode` to produce a `Vec<u8>` internally.
701 impl<'a> serialize::Decodable<Decoder<'a>> for Vec<u8> {
702 fn decode(d: &mut Decoder<'a>) -> Result<Self, String> {
703 let len = serialize::Decoder::read_usize(d)?;
704 Ok(d.read_raw_bytes(len).to_owned())
708 // An integer that will always encode to 8 bytes.
709 pub struct IntEncodedWithFixedSize(pub u64);
711 impl IntEncodedWithFixedSize {
712 pub const ENCODED_SIZE: usize = 8;
715 impl serialize::Encodable<Encoder> for IntEncodedWithFixedSize {
717 fn encode(&self, e: &mut Encoder) -> EncodeResult {
718 let _start_pos = e.position();
719 e.emit_raw_bytes(&self.0.to_le_bytes())?;
720 let _end_pos = e.position();
721 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
726 impl serialize::Encodable<FileEncoder> for IntEncodedWithFixedSize {
728 fn encode(&self, e: &mut FileEncoder) -> FileEncodeResult {
729 let _start_pos = e.position();
730 e.emit_raw_bytes(&self.0.to_le_bytes())?;
731 let _end_pos = e.position();
732 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
737 impl<'a> serialize::Decodable<Decoder<'a>> for IntEncodedWithFixedSize {
739 fn decode(decoder: &mut Decoder<'a>) -> Result<IntEncodedWithFixedSize, String> {
740 let _start_pos = decoder.position();
741 let bytes = decoder.read_raw_bytes(IntEncodedWithFixedSize::ENCODED_SIZE);
742 let _end_pos = decoder.position();
743 debug_assert_eq!((_end_pos - _start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
745 let value = u64::from_le_bytes(bytes.try_into().unwrap());
746 Ok(IntEncodedWithFixedSize(value))