+ // Tracking position this way instead of having a `self.position` field
+ // means that we don't have to update the position on every write call.
+ self.flushed + self.buffered
+ }
+
+ #[inline]
+ pub fn emit_raw_bytes(&mut self, s: &[u8]) -> FileEncodeResult {
+ self.write_all(s)
+ }
+
+ pub fn flush(&mut self) -> FileEncodeResult {
+ // This is basically a copy of `BufWriter::flush`. If `BufWriter` ever
+ // offers a raw buffer access API, we can use it, and remove this.
+
+ /// Helper struct to ensure the buffer is updated after all the writes
+ /// are complete. It tracks the number of written bytes and drains them
+ /// all from the front of the buffer when dropped.
+ struct BufGuard<'a> {
+ buffer: &'a mut [u8],
+ encoder_buffered: &'a mut usize,
+ encoder_flushed: &'a mut usize,
+ flushed: usize,
+ }
+
+ impl<'a> BufGuard<'a> {
+ fn new(
+ buffer: &'a mut [u8],
+ encoder_buffered: &'a mut usize,
+ encoder_flushed: &'a mut usize,
+ ) -> Self {
+ assert_eq!(buffer.len(), *encoder_buffered);
+ Self { buffer, encoder_buffered, encoder_flushed, flushed: 0 }
+ }
+
+ /// The unwritten part of the buffer
+ fn remaining(&self) -> &[u8] {
+ &self.buffer[self.flushed..]
+ }
+
+ /// Flag some bytes as removed from the front of the buffer
+ fn consume(&mut self, amt: usize) {
+ self.flushed += amt;
+ }
+
+ /// true if all of the bytes have been written
+ fn done(&self) -> bool {
+ self.flushed >= *self.encoder_buffered
+ }
+ }
+
+ impl Drop for BufGuard<'_> {
+ fn drop(&mut self) {
+ if self.flushed > 0 {
+ if self.done() {
+ *self.encoder_flushed += *self.encoder_buffered;
+ *self.encoder_buffered = 0;
+ } else {
+ self.buffer.copy_within(self.flushed.., 0);
+ *self.encoder_flushed += self.flushed;
+ *self.encoder_buffered -= self.flushed;
+ }
+ }
+ }
+ }
+
+ let mut guard = BufGuard::new(
+ unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[..self.buffered]) },
+ &mut self.buffered,
+ &mut self.flushed,
+ );
+
+ while !guard.done() {
+ match self.file.write(guard.remaining()) {
+ Ok(0) => {
+ return Err(io::Error::new(
+ io::ErrorKind::WriteZero,
+ "failed to write the buffered data",
+ ));
+ }
+ Ok(n) => guard.consume(n),
+ Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
+ Err(e) => return Err(e),
+ }
+ }
+
+ Ok(())
+ }
+
+ #[inline]
+ fn capacity(&self) -> usize {
+ self.buf.len()
+ }
+
+ #[inline]
+ fn write_one(&mut self, value: u8) -> FileEncodeResult {
+ // We ensure this during `FileEncoder` construction.
+ debug_assert!(self.capacity() >= 1);
+
+ let mut buffered = self.buffered;
+
+ if std::intrinsics::unlikely(buffered >= self.capacity()) {
+ self.flush()?;
+ buffered = 0;
+ }
+
+ // SAFETY: The above check and `flush` ensures that there is enough
+ // room to write the input to the buffer.
+ unsafe {
+ *MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered) = value;
+ }
+
+ self.buffered = buffered + 1;
+
+ Ok(())
+ }
+
+ #[inline]
+ fn write_all(&mut self, buf: &[u8]) -> FileEncodeResult {
+ let capacity = self.capacity();
+ let buf_len = buf.len();
+
+ if std::intrinsics::likely(buf_len <= capacity) {
+ let mut buffered = self.buffered;
+
+ if std::intrinsics::unlikely(buf_len > capacity - buffered) {
+ self.flush()?;
+ buffered = 0;
+ }
+
+ // SAFETY: The above check and `flush` ensures that there is enough
+ // room to write the input to the buffer.
+ unsafe {
+ let src = buf.as_ptr();
+ let dst = MaybeUninit::slice_as_mut_ptr(&mut self.buf).add(buffered);
+ ptr::copy_nonoverlapping(src, dst, buf_len);
+ }
+
+ self.buffered = buffered + buf_len;
+
+ Ok(())
+ } else {
+ self.write_all_unbuffered(buf)
+ }
+ }
+
+ fn write_all_unbuffered(&mut self, mut buf: &[u8]) -> FileEncodeResult {
+ if self.buffered > 0 {
+ self.flush()?;
+ }
+
+ // This is basically a copy of `Write::write_all` but also updates our
+ // `self.flushed`. It's necessary because `Write::write_all` does not
+ // return the number of bytes written when an error is encountered, and
+ // without that, we cannot accurately update `self.flushed` on error.
+ while !buf.is_empty() {
+ match self.file.write(buf) {
+ Ok(0) => {
+ return Err(io::Error::new(
+ io::ErrorKind::WriteZero,
+ "failed to write whole buffer",
+ ));
+ }
+ Ok(n) => {
+ buf = &buf[n..];
+ self.flushed += n;
+ }
+ Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
+ Err(e) => return Err(e),
+ }
+ }
+
+ Ok(())
+ }
+}
+
+impl Drop for FileEncoder {
+ fn drop(&mut self) {
+ let _result = self.flush();
+ }
+}
+
+macro_rules! file_encoder_write_leb128 {
+ ($enc:expr, $value:expr, $int_ty:ty, $fun:ident) => {{
+ const MAX_ENCODED_LEN: usize = max_leb128_len!($int_ty);
+
+ // We ensure this during `FileEncoder` construction.
+ debug_assert!($enc.capacity() >= MAX_ENCODED_LEN);
+
+ let mut buffered = $enc.buffered;
+
+ // This can't overflow. See assertion in `FileEncoder::with_capacity`.
+ if std::intrinsics::unlikely(buffered + MAX_ENCODED_LEN > $enc.capacity()) {
+ $enc.flush()?;
+ buffered = 0;
+ }
+
+ // SAFETY: The above check and flush ensures that there is enough
+ // room to write the encoded value to the buffer.
+ let buf = unsafe {
+ &mut *($enc.buf.as_mut_ptr().add(buffered) as *mut [MaybeUninit<u8>; MAX_ENCODED_LEN])
+ };
+
+ let encoded = leb128::$fun(buf, $value);
+ $enc.buffered = buffered + encoded.len();
+
+ Ok(())
+ }};
+}
+
+impl serialize::Encoder for FileEncoder {
+ type Error = io::Error;
+
+ #[inline]
+ fn emit_unit(&mut self) -> FileEncodeResult {
+ Ok(())
+ }
+
+ #[inline]
+ fn emit_usize(&mut self, v: usize) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, usize, write_usize_leb128)
+ }
+
+ #[inline]
+ fn emit_u128(&mut self, v: u128) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, u128, write_u128_leb128)
+ }
+
+ #[inline]
+ fn emit_u64(&mut self, v: u64) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, u64, write_u64_leb128)
+ }
+
+ #[inline]
+ fn emit_u32(&mut self, v: u32) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, u32, write_u32_leb128)
+ }
+
+ #[inline]
+ fn emit_u16(&mut self, v: u16) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, u16, write_u16_leb128)
+ }
+
+ #[inline]
+ fn emit_u8(&mut self, v: u8) -> FileEncodeResult {
+ self.write_one(v)
+ }
+
+ #[inline]
+ fn emit_isize(&mut self, v: isize) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, isize, write_isize_leb128)
+ }
+
+ #[inline]
+ fn emit_i128(&mut self, v: i128) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, i128, write_i128_leb128)
+ }
+
+ #[inline]
+ fn emit_i64(&mut self, v: i64) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, i64, write_i64_leb128)
+ }
+
+ #[inline]
+ fn emit_i32(&mut self, v: i32) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, i32, write_i32_leb128)
+ }
+
+ #[inline]
+ fn emit_i16(&mut self, v: i16) -> FileEncodeResult {
+ file_encoder_write_leb128!(self, v, i16, write_i16_leb128)
+ }
+
+ #[inline]
+ fn emit_i8(&mut self, v: i8) -> FileEncodeResult {
+ let as_u8: u8 = unsafe { std::mem::transmute(v) };
+ self.emit_u8(as_u8)
+ }
+
+ #[inline]
+ fn emit_bool(&mut self, v: bool) -> FileEncodeResult {
+ self.emit_u8(if v { 1 } else { 0 })
+ }
+
+ #[inline]
+ fn emit_f64(&mut self, v: f64) -> FileEncodeResult {
+ let as_u64: u64 = v.to_bits();
+ self.emit_u64(as_u64)
+ }
+
+ #[inline]
+ fn emit_f32(&mut self, v: f32) -> FileEncodeResult {
+ let as_u32: u32 = v.to_bits();
+ self.emit_u32(as_u32)
+ }
+
+ #[inline]
+ fn emit_char(&mut self, v: char) -> FileEncodeResult {
+ self.emit_u32(v as u32)
+ }
+
+ #[inline]
+ fn emit_str(&mut self, v: &str) -> FileEncodeResult {
+ self.emit_usize(v.len())?;
+ self.emit_raw_bytes(v.as_bytes())