1 use crate::leb128
::{self, max_leb128_len}
;
2 use crate::serialize
::{Decodable, Decoder, Encodable, Encoder}
;
3 use std
::convert
::TryInto
;
5 use std
::io
::{self, Write}
;
6 use std
::mem
::MaybeUninit
;
10 // -----------------------------------------------------------------------------
12 // -----------------------------------------------------------------------------
14 pub struct MemEncoder
{
19 pub fn new() -> MemEncoder
{
20 MemEncoder { data: vec![] }
24 pub fn position(&self) -> usize {
28 pub fn finish(self) -> Vec
<u8> {
33 macro_rules
! write_leb128
{
34 ($enc
:expr
, $value
:expr
, $int_ty
:ty
, $fun
:ident
) => {{
35 const MAX_ENCODED_LEN
: usize = max_leb128_len
!($int_ty
);
36 let old_len
= $enc
.data
.len();
38 if MAX_ENCODED_LEN
> $enc
.data
.capacity() - old_len
{
39 $enc
.data
.reserve(MAX_ENCODED_LEN
);
42 // SAFETY: The above check and `reserve` ensures that there is enough
43 // room to write the encoded value to the vector's internal buffer.
45 let buf
= &mut *($enc
.data
.as_mut_ptr().add(old_len
)
46 as *mut [MaybeUninit
<u8>; MAX_ENCODED_LEN
]);
47 let encoded
= leb128
::$
fun(buf
, $value
);
48 $enc
.data
.set_len(old_len
+ encoded
.len());
53 /// A byte that [cannot occur in UTF8 sequences][utf8]. Used to mark the end of a string.
54 /// This way we can skip validation and still be relatively sure that deserialization
55 /// did not desynchronize.
57 /// [utf8]: https://en.wikipedia.org/w/index.php?title=UTF-8&oldid=1058865525#Codepage_layout
58 const STR_SENTINEL
: u8 = 0xC1;
60 impl Encoder
for MemEncoder
{
62 fn emit_usize(&mut self, v
: usize) {
63 write_leb128
!(self, v
, usize, write_usize_leb128
)
67 fn emit_u128(&mut self, v
: u128
) {
68 write_leb128
!(self, v
, u128
, write_u128_leb128
);
72 fn emit_u64(&mut self, v
: u64) {
73 write_leb128
!(self, v
, u64, write_u64_leb128
);
77 fn emit_u32(&mut self, v
: u32) {
78 write_leb128
!(self, v
, u32, write_u32_leb128
);
82 fn emit_u16(&mut self, v
: u16) {
83 self.data
.extend_from_slice(&v
.to_le_bytes());
87 fn emit_u8(&mut self, v
: u8) {
92 fn emit_isize(&mut self, v
: isize) {
93 write_leb128
!(self, v
, isize, write_isize_leb128
)
97 fn emit_i128(&mut self, v
: i128
) {
98 write_leb128
!(self, v
, i128
, write_i128_leb128
)
102 fn emit_i64(&mut self, v
: i64) {
103 write_leb128
!(self, v
, i64, write_i64_leb128
)
107 fn emit_i32(&mut self, v
: i32) {
108 write_leb128
!(self, v
, i32, write_i32_leb128
)
112 fn emit_i16(&mut self, v
: i16) {
113 self.data
.extend_from_slice(&v
.to_le_bytes());
117 fn emit_i8(&mut self, v
: i8) {
118 self.emit_u8(v
as u8);
122 fn emit_bool(&mut self, v
: bool
) {
123 self.emit_u8(if v { 1 }
else { 0 }
);
127 fn emit_f64(&mut self, v
: f64) {
128 let as_u64
: u64 = v
.to_bits();
129 self.emit_u64(as_u64
);
133 fn emit_f32(&mut self, v
: f32) {
134 let as_u32
: u32 = v
.to_bits();
135 self.emit_u32(as_u32
);
139 fn emit_char(&mut self, v
: char) {
140 self.emit_u32(v
as u32);
144 fn emit_str(&mut self, v
: &str) {
145 self.emit_usize(v
.len());
146 self.emit_raw_bytes(v
.as_bytes());
147 self.emit_u8(STR_SENTINEL
);
151 fn emit_raw_bytes(&mut self, s
: &[u8]) {
152 self.data
.extend_from_slice(s
);
156 pub type FileEncodeResult
= Result
<usize, io
::Error
>;
158 // `FileEncoder` encodes data to file via fixed-size buffer.
160 // When encoding large amounts of data to a file, using `FileEncoder` may be
161 // preferred over using `MemEncoder` to encode to a `Vec`, and then writing the
162 // `Vec` to file, as the latter uses as much memory as there is encoded data,
163 // while the former uses the fixed amount of memory allocated to the buffer.
164 // `FileEncoder` also has the advantage of not needing to reallocate as data
165 // is appended to it, but the disadvantage of requiring more error handling,
166 // which has some runtime overhead.
167 pub struct FileEncoder
{
168 // The input buffer. For adequate performance, we need more control over
169 // buffering than `BufWriter` offers. If `BufWriter` ever offers a raw
170 // buffer access API, we can use it, and remove `buf` and `buffered`.
171 buf
: Box
<[MaybeUninit
<u8>]>,
175 // This is used to implement delayed error handling, as described in the
176 // comment on `trait Encoder`.
177 res
: Result
<(), io
::Error
>,
181 pub fn new
<P
: AsRef
<Path
>>(path
: P
) -> io
::Result
<Self> {
182 const DEFAULT_BUF_SIZE
: usize = 8192;
183 FileEncoder
::with_capacity(path
, DEFAULT_BUF_SIZE
)
186 pub fn with_capacity
<P
: AsRef
<Path
>>(path
: P
, capacity
: usize) -> io
::Result
<Self> {
187 // Require capacity at least as large as the largest LEB128 encoding
188 // here, so that we don't have to check or handle this on every write.
189 assert
!(capacity
>= max_leb128_len());
191 // Require capacity small enough such that some capacity checks can be
192 // done using guaranteed non-overflowing add rather than sub, which
193 // shaves an instruction off those code paths (on x86 at least).
194 assert
!(capacity
<= usize::MAX
- max_leb128_len());
196 let file
= File
::create(path
)?
;
199 buf
: Box
::new_uninit_slice(capacity
),
208 pub fn position(&self) -> usize {
209 // Tracking position this way instead of having a `self.position` field
210 // means that we don't have to update the position on every write call.
211 self.flushed
+ self.buffered
214 pub fn flush(&mut self) {
215 // This is basically a copy of `BufWriter::flush`. If `BufWriter` ever
216 // offers a raw buffer access API, we can use it, and remove this.
218 /// Helper struct to ensure the buffer is updated after all the writes
219 /// are complete. It tracks the number of written bytes and drains them
220 /// all from the front of the buffer when dropped.
221 struct BufGuard
<'a
> {
222 buffer
: &'a
mut [u8],
223 encoder_buffered
: &'a
mut usize,
224 encoder_flushed
: &'a
mut usize,
228 impl<'a
> BufGuard
<'a
> {
230 buffer
: &'a
mut [u8],
231 encoder_buffered
: &'a
mut usize,
232 encoder_flushed
: &'a
mut usize,
234 assert_eq
!(buffer
.len(), *encoder_buffered
);
235 Self { buffer, encoder_buffered, encoder_flushed, flushed: 0 }
238 /// The unwritten part of the buffer
239 fn remaining(&self) -> &[u8] {
240 &self.buffer
[self.flushed
..]
243 /// Flag some bytes as removed from the front of the buffer
244 fn consume(&mut self, amt
: usize) {
248 /// true if all of the bytes have been written
249 fn done(&self) -> bool
{
250 self.flushed
>= *self.encoder_buffered
254 impl Drop
for BufGuard
<'_
> {
256 if self.flushed
> 0 {
258 *self.encoder_flushed
+= *self.encoder_buffered
;
259 *self.encoder_buffered
= 0;
261 self.buffer
.copy_within(self.flushed
.., 0);
262 *self.encoder_flushed
+= self.flushed
;
263 *self.encoder_buffered
-= self.flushed
;
269 // If we've already had an error, do nothing. It'll get reported after
270 // `finish` is called.
271 if self.res
.is_err() {
275 let mut guard
= BufGuard
::new(
276 unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[..self.buffered]) }
,
281 while !guard
.done() {
282 match self.file
.write(guard
.remaining()) {
284 self.res
= Err(io
::Error
::new(
285 io
::ErrorKind
::WriteZero
,
286 "failed to write the buffered data",
290 Ok(n
) => guard
.consume(n
),
291 Err(ref e
) if e
.kind() == io
::ErrorKind
::Interrupted
=> {}
301 fn capacity(&self) -> usize {
306 fn write_one(&mut self, value
: u8) {
307 // We ensure this during `FileEncoder` construction.
308 debug_assert
!(self.capacity() >= 1);
310 let mut buffered
= self.buffered
;
312 if std
::intrinsics
::unlikely(buffered
>= self.capacity()) {
317 // SAFETY: The above check and `flush` ensures that there is enough
318 // room to write the input to the buffer.
320 *MaybeUninit
::slice_as_mut_ptr(&mut self.buf
).add(buffered
) = value
;
323 self.buffered
= buffered
+ 1;
327 fn write_all(&mut self, buf
: &[u8]) {
328 let capacity
= self.capacity();
329 let buf_len
= buf
.len();
331 if std
::intrinsics
::likely(buf_len
<= capacity
) {
332 let mut buffered
= self.buffered
;
334 if std
::intrinsics
::unlikely(buf_len
> capacity
- buffered
) {
339 // SAFETY: The above check and `flush` ensures that there is enough
340 // room to write the input to the buffer.
342 let src
= buf
.as_ptr();
343 let dst
= MaybeUninit
::slice_as_mut_ptr(&mut self.buf
).add(buffered
);
344 ptr
::copy_nonoverlapping(src
, dst
, buf_len
);
347 self.buffered
= buffered
+ buf_len
;
349 self.write_all_unbuffered(buf
);
353 fn write_all_unbuffered(&mut self, mut buf
: &[u8]) {
354 // If we've already had an error, do nothing. It'll get reported after
355 // `finish` is called.
356 if self.res
.is_err() {
360 if self.buffered
> 0 {
364 // This is basically a copy of `Write::write_all` but also updates our
365 // `self.flushed`. It's necessary because `Write::write_all` does not
366 // return the number of bytes written when an error is encountered, and
367 // without that, we cannot accurately update `self.flushed` on error.
368 while !buf
.is_empty() {
369 match self.file
.write(buf
) {
371 self.res
= Err(io
::Error
::new(
372 io
::ErrorKind
::WriteZero
,
373 "failed to write whole buffer",
381 Err(ref e
) if e
.kind() == io
::ErrorKind
::Interrupted
=> {}
390 pub fn finish(mut self) -> Result
<usize, io
::Error
> {
393 let res
= std
::mem
::replace(&mut self.res
, Ok(()));
394 res
.map(|()| self.position())
398 impl Drop
for FileEncoder
{
400 // Likely to be a no-op, because `finish` should have been called and
401 // it also flushes. But do it just in case.
402 let _result
= self.flush();
406 macro_rules
! file_encoder_write_leb128
{
407 ($enc
:expr
, $value
:expr
, $int_ty
:ty
, $fun
:ident
) => {{
408 const MAX_ENCODED_LEN
: usize = max_leb128_len
!($int_ty
);
410 // We ensure this during `FileEncoder` construction.
411 debug_assert
!($enc
.capacity() >= MAX_ENCODED_LEN
);
413 let mut buffered
= $enc
.buffered
;
415 // This can't overflow. See assertion in `FileEncoder::with_capacity`.
416 if std
::intrinsics
::unlikely(buffered
+ MAX_ENCODED_LEN
> $enc
.capacity()) {
421 // SAFETY: The above check and flush ensures that there is enough
422 // room to write the encoded value to the buffer.
424 &mut *($enc
.buf
.as_mut_ptr().add(buffered
) as *mut [MaybeUninit
<u8>; MAX_ENCODED_LEN
])
427 let encoded
= leb128
::$
fun(buf
, $value
);
428 $enc
.buffered
= buffered
+ encoded
.len();
432 impl Encoder
for FileEncoder
{
434 fn emit_usize(&mut self, v
: usize) {
435 file_encoder_write_leb128
!(self, v
, usize, write_usize_leb128
)
439 fn emit_u128(&mut self, v
: u128
) {
440 file_encoder_write_leb128
!(self, v
, u128
, write_u128_leb128
)
444 fn emit_u64(&mut self, v
: u64) {
445 file_encoder_write_leb128
!(self, v
, u64, write_u64_leb128
)
449 fn emit_u32(&mut self, v
: u32) {
450 file_encoder_write_leb128
!(self, v
, u32, write_u32_leb128
)
454 fn emit_u16(&mut self, v
: u16) {
455 self.write_all(&v
.to_le_bytes());
459 fn emit_u8(&mut self, v
: u8) {
464 fn emit_isize(&mut self, v
: isize) {
465 file_encoder_write_leb128
!(self, v
, isize, write_isize_leb128
)
469 fn emit_i128(&mut self, v
: i128
) {
470 file_encoder_write_leb128
!(self, v
, i128
, write_i128_leb128
)
474 fn emit_i64(&mut self, v
: i64) {
475 file_encoder_write_leb128
!(self, v
, i64, write_i64_leb128
)
479 fn emit_i32(&mut self, v
: i32) {
480 file_encoder_write_leb128
!(self, v
, i32, write_i32_leb128
)
484 fn emit_i16(&mut self, v
: i16) {
485 self.write_all(&v
.to_le_bytes());
489 fn emit_i8(&mut self, v
: i8) {
490 self.emit_u8(v
as u8);
494 fn emit_bool(&mut self, v
: bool
) {
495 self.emit_u8(if v { 1 }
else { 0 }
);
499 fn emit_f64(&mut self, v
: f64) {
500 let as_u64
: u64 = v
.to_bits();
501 self.emit_u64(as_u64
);
505 fn emit_f32(&mut self, v
: f32) {
506 let as_u32
: u32 = v
.to_bits();
507 self.emit_u32(as_u32
);
511 fn emit_char(&mut self, v
: char) {
512 self.emit_u32(v
as u32);
516 fn emit_str(&mut self, v
: &str) {
517 self.emit_usize(v
.len());
518 self.emit_raw_bytes(v
.as_bytes());
519 self.emit_u8(STR_SENTINEL
);
523 fn emit_raw_bytes(&mut self, s
: &[u8]) {
528 // -----------------------------------------------------------------------------
530 // -----------------------------------------------------------------------------
532 pub struct MemDecoder
<'a
> {
537 impl<'a
> MemDecoder
<'a
> {
539 pub fn new(data
: &'a
[u8], position
: usize) -> MemDecoder
<'a
> {
540 MemDecoder { data, position }
544 pub fn position(&self) -> usize {
549 pub fn set_position(&mut self, pos
: usize) {
554 pub fn advance(&mut self, bytes
: usize) {
555 self.position
+= bytes
;
559 macro_rules
! read_leb128
{
560 ($dec
:expr
, $fun
:ident
) => {{ leb128::$fun($dec.data, &mut $dec.position) }
};
563 impl<'a
> Decoder
for MemDecoder
<'a
> {
565 fn read_u128(&mut self) -> u128
{
566 read_leb128
!(self, read_u128_leb128
)
570 fn read_u64(&mut self) -> u64 {
571 read_leb128
!(self, read_u64_leb128
)
575 fn read_u32(&mut self) -> u32 {
576 read_leb128
!(self, read_u32_leb128
)
580 fn read_u16(&mut self) -> u16 {
581 let bytes
= [self.data
[self.position
], self.data
[self.position
+ 1]];
582 let value
= u16::from_le_bytes(bytes
);
588 fn read_u8(&mut self) -> u8 {
589 let value
= self.data
[self.position
];
595 fn read_usize(&mut self) -> usize {
596 read_leb128
!(self, read_usize_leb128
)
600 fn read_i128(&mut self) -> i128
{
601 read_leb128
!(self, read_i128_leb128
)
605 fn read_i64(&mut self) -> i64 {
606 read_leb128
!(self, read_i64_leb128
)
610 fn read_i32(&mut self) -> i32 {
611 read_leb128
!(self, read_i32_leb128
)
615 fn read_i16(&mut self) -> i16 {
616 let bytes
= [self.data
[self.position
], self.data
[self.position
+ 1]];
617 let value
= i16::from_le_bytes(bytes
);
623 fn read_i8(&mut self) -> i8 {
624 let value
= self.data
[self.position
];
630 fn read_isize(&mut self) -> isize {
631 read_leb128
!(self, read_isize_leb128
)
635 fn read_bool(&mut self) -> bool
{
636 let value
= self.read_u8();
641 fn read_f64(&mut self) -> f64 {
642 let bits
= self.read_u64();
647 fn read_f32(&mut self) -> f32 {
648 let bits
= self.read_u32();
653 fn read_char(&mut self) -> char {
654 let bits
= self.read_u32();
655 std
::char::from_u32(bits
).unwrap()
659 fn read_str(&mut self) -> &'a
str {
660 let len
= self.read_usize();
661 let sentinel
= self.data
[self.position
+ len
];
662 assert
!(sentinel
== STR_SENTINEL
);
664 std
::str::from_utf8_unchecked(&self.data
[self.position
..self.position
+ len
])
666 self.position
+= len
+ 1;
671 fn read_raw_bytes(&mut self, bytes
: usize) -> &'a
[u8] {
672 let start
= self.position
;
673 self.position
+= bytes
;
674 &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 Encodable
<MemEncoder
> for [u8] {
686 fn encode(&self, e
: &mut MemEncoder
) {
687 Encoder
::emit_usize(e
, self.len());
688 e
.emit_raw_bytes(self);
692 impl Encodable
<FileEncoder
> for [u8] {
693 fn encode(&self, e
: &mut FileEncoder
) {
694 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
> Decodable
<MemDecoder
<'a
>> for Vec
<u8> {
702 fn decode(d
: &mut MemDecoder
<'a
>) -> Self {
703 let len
= Decoder
::read_usize(d
);
704 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 Encodable
<MemEncoder
> for IntEncodedWithFixedSize
{
717 fn encode(&self, e
: &mut MemEncoder
) {
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
);
725 impl Encodable
<FileEncoder
> for IntEncodedWithFixedSize
{
727 fn encode(&self, e
: &mut FileEncoder
) {
728 let _start_pos
= e
.position();
729 e
.emit_raw_bytes(&self.0.to_le_bytes());
730 let _end_pos
= e
.position();
731 debug_assert_eq
!((_end_pos
- _start_pos
), IntEncodedWithFixedSize
::ENCODED_SIZE
);
735 impl<'a
> Decodable
<MemDecoder
<'a
>> for IntEncodedWithFixedSize
{
737 fn decode(decoder
: &mut MemDecoder
<'a
>) -> IntEncodedWithFixedSize
{
738 let _start_pos
= decoder
.position();
739 let bytes
= decoder
.read_raw_bytes(IntEncodedWithFixedSize
::ENCODED_SIZE
);
740 let value
= u64::from_le_bytes(bytes
.try_into().unwrap());
741 let _end_pos
= decoder
.position();
742 debug_assert_eq
!((_end_pos
- _start_pos
), IntEncodedWithFixedSize
::ENCODED_SIZE
);
744 IntEncodedWithFixedSize(value
)