1use std::fs::File;
2use std::io::{self, Write};
3use std::marker::PhantomData;
4use std::ops::Range;
5use std::path::{Path, PathBuf};
6
7use crate::int_overflow::DebugStrictAdd;
10use crate::leb128;
11use crate::serialize::{Decodable, Decoder, Encodable, Encoder};
12
13pub mod mem_encoder;
14
15pub type FileEncodeResult = Result<usize, (PathBuf, io::Error)>;
20
21pub const MAGIC_END_BYTES: &[u8] = b"rust-end-file";
22
23const BUF_SIZE: usize = 64 * 1024;
25
26pub struct FileEncoder {
33 buf: Box<[u8; BUF_SIZE]>,
39 buffered: usize,
40 flushed: usize,
41 file: File,
42 res: Result<(), io::Error>,
45 path: PathBuf,
46 #[cfg(debug_assertions)]
47 finished: bool,
48}
49
50impl FileEncoder {
51 pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
52 let file =
56 File::options().read(true).write(true).create(true).truncate(true).open(&path)?;
57
58 Ok(FileEncoder {
59 buf: vec![0u8; BUF_SIZE].into_boxed_slice().try_into().unwrap(),
60 path: path.as_ref().into(),
61 buffered: 0,
62 flushed: 0,
63 file,
64 res: Ok(()),
65 #[cfg(debug_assertions)]
66 finished: false,
67 })
68 }
69
70 #[inline]
71 pub fn position(&self) -> usize {
72 self.flushed.debug_strict_add(self.buffered)
76 }
77
78 #[cold]
79 #[inline(never)]
80 pub fn flush(&mut self) {
81 #[cfg(debug_assertions)]
82 {
83 self.finished = false;
84 }
85 if self.res.is_ok() {
86 self.res = self.file.write_all(&self.buf[..self.buffered]);
87 }
88 self.flushed += self.buffered;
89 self.buffered = 0;
90 }
91
92 pub fn file(&self) -> &File {
93 &self.file
94 }
95
96 pub fn path(&self) -> &Path {
97 &self.path
98 }
99
100 #[inline]
101 fn buffer_empty(&mut self) -> &mut [u8] {
102 unsafe { self.buf.get_unchecked_mut(self.buffered..) }
104 }
105
106 #[cold]
107 #[inline(never)]
108 fn write_all_cold_path(&mut self, buf: &[u8]) {
109 self.flush();
110 if let Some(dest) = self.buf.get_mut(..buf.len()) {
111 dest.copy_from_slice(buf);
112 self.buffered += buf.len();
113 } else {
114 if self.res.is_ok() {
115 self.res = self.file.write_all(buf);
116 }
117 self.flushed += buf.len();
118 }
119 }
120
121 #[inline]
122 fn write_all(&mut self, buf: &[u8]) {
123 #[cfg(debug_assertions)]
124 {
125 self.finished = false;
126 }
127 if let Some(dest) = self.buffer_empty().get_mut(..buf.len()) {
128 dest.copy_from_slice(buf);
129 self.buffered = self.buffered.debug_strict_add(buf.len());
130 } else {
131 self.write_all_cold_path(buf);
132 }
133 }
134
135 #[inline]
150 pub fn write_with<const N: usize>(&mut self, visitor: impl FnOnce(&mut [u8; N]) -> usize) {
151 #[cfg(debug_assertions)]
152 {
153 self.finished = false;
154 }
155 let flush_threshold = const { BUF_SIZE.checked_sub(N).unwrap() };
156 if std::intrinsics::unlikely(self.buffered > flush_threshold) {
157 self.flush();
158 }
159 let buf = unsafe { self.buffer_empty().first_chunk_mut::<N>().unwrap_unchecked() };
163 let written = visitor(buf);
164 if written > N {
166 Self::panic_invalid_write::<N>(written);
167 }
168 self.buffered = self.buffered.debug_strict_add(written);
169 }
170
171 #[cold]
172 #[inline(never)]
173 fn panic_invalid_write<const N: usize>(written: usize) {
174 panic!("FileEncoder::write_with::<{N}> cannot be used to write {written} bytes");
175 }
176
177 #[inline]
179 pub fn write_array<const N: usize>(&mut self, buf: [u8; N]) {
180 self.write_with(|dest| {
181 *dest = buf;
182 N
183 })
184 }
185
186 pub fn finish(&mut self) -> FileEncodeResult {
187 self.write_all(MAGIC_END_BYTES);
188 self.flush();
189 #[cfg(debug_assertions)]
190 {
191 self.finished = true;
192 }
193 match std::mem::replace(&mut self.res, Ok(())) {
194 Ok(()) => Ok(self.position()),
195 Err(e) => Err((self.path.clone(), e)),
196 }
197 }
198}
199
200#[cfg(debug_assertions)]
201impl Drop for FileEncoder {
202 fn drop(&mut self) {
203 if !std::thread::panicking() {
204 assert!(self.finished);
205 }
206 }
207}
208
209macro_rules! write_leb128 {
210 ($this_fn:ident, $int_ty:ty, $write_leb_fn:ident) => {
211 #[inline]
212 fn $this_fn(&mut self, v: $int_ty) {
213 self.write_with(|buf| leb128::$write_leb_fn(buf, v))
214 }
215 };
216}
217
218impl Encoder for FileEncoder {
219 write_leb128!(emit_usize, usize, write_usize_leb128);
220 write_leb128!(emit_u128, u128, write_u128_leb128);
221 write_leb128!(emit_u64, u64, write_u64_leb128);
222 write_leb128!(emit_u32, u32, write_u32_leb128);
223
224 #[inline]
225 fn emit_u16(&mut self, v: u16) {
226 self.write_array(v.to_le_bytes());
227 }
228
229 #[inline]
230 fn emit_u8(&mut self, v: u8) {
231 self.write_array([v]);
232 }
233
234 write_leb128!(emit_isize, isize, write_isize_leb128);
235 write_leb128!(emit_i128, i128, write_i128_leb128);
236 write_leb128!(emit_i64, i64, write_i64_leb128);
237 write_leb128!(emit_i32, i32, write_i32_leb128);
238
239 #[inline]
240 fn emit_i16(&mut self, v: i16) {
241 self.write_array(v.to_le_bytes());
242 }
243
244 #[inline]
245 fn emit_raw_bytes(&mut self, s: &[u8]) {
246 self.write_all(s);
247 }
248}
249
250pub struct MemDecoder<'a> {
260 start: *const u8,
261 current: *const u8,
262 end: *const u8,
263 _marker: PhantomData<&'a u8>,
264}
265
266impl<'a> MemDecoder<'a> {
267 #[inline]
268 pub fn new(data: &'a [u8], position: usize) -> Result<MemDecoder<'a>, ()> {
269 let data = data.strip_suffix(MAGIC_END_BYTES).ok_or(())?;
270 let Range { start, end } = data.as_ptr_range();
271 Ok(MemDecoder { start, current: data[position..].as_ptr(), end, _marker: PhantomData })
272 }
273
274 #[inline]
275 pub fn split_at(&self, position: usize) -> MemDecoder<'a> {
276 assert!(position <= self.len());
277 let current = unsafe { self.start.add(position) };
279 MemDecoder { start: self.start, current, end: self.end, _marker: PhantomData }
280 }
281
282 #[inline]
283 pub fn len(&self) -> usize {
284 unsafe { self.end.offset_from_unsigned(self.start) }
286 }
287
288 #[inline]
289 pub fn remaining(&self) -> usize {
290 unsafe { self.end.offset_from_unsigned(self.current) }
292 }
293
294 #[cold]
295 #[inline(never)]
296 fn decoder_exhausted() -> ! {
297 panic!("MemDecoder exhausted")
298 }
299
300 #[inline]
301 pub fn read_array<const N: usize>(&mut self) -> [u8; N] {
302 self.read_raw_bytes(N).try_into().unwrap()
303 }
304
305 #[inline]
309 pub fn with_position<F, T>(&mut self, pos: usize, func: F) -> T
310 where
311 F: Fn(&mut MemDecoder<'a>) -> T,
312 {
313 struct SetOnDrop<'a, 'guarded> {
314 decoder: &'guarded mut MemDecoder<'a>,
315 current: *const u8,
316 }
317 impl Drop for SetOnDrop<'_, '_> {
318 fn drop(&mut self) {
319 self.decoder.current = self.current;
320 }
321 }
322
323 if pos >= self.len() {
324 Self::decoder_exhausted();
325 }
326 let previous = self.current;
327 unsafe {
329 self.current = self.start.add(pos);
330 }
331 let guard = SetOnDrop { current: previous, decoder: self };
332 func(guard.decoder)
333 }
334}
335
336macro_rules! read_leb128 {
337 ($this_fn:ident, $int_ty:ty, $read_leb_fn:ident) => {
338 #[inline]
339 fn $this_fn(&mut self) -> $int_ty {
340 leb128::$read_leb_fn(self)
341 }
342 };
343}
344
345impl<'a> Decoder for MemDecoder<'a> {
346 read_leb128!(read_usize, usize, read_usize_leb128);
347 read_leb128!(read_u128, u128, read_u128_leb128);
348 read_leb128!(read_u64, u64, read_u64_leb128);
349 read_leb128!(read_u32, u32, read_u32_leb128);
350
351 #[inline]
352 fn read_u16(&mut self) -> u16 {
353 u16::from_le_bytes(self.read_array())
354 }
355
356 #[inline]
357 fn read_u8(&mut self) -> u8 {
358 if self.current == self.end {
359 Self::decoder_exhausted();
360 }
361 unsafe {
363 let byte = *self.current;
364 self.current = self.current.add(1);
365 byte
366 }
367 }
368
369 read_leb128!(read_isize, isize, read_isize_leb128);
370 read_leb128!(read_i128, i128, read_i128_leb128);
371 read_leb128!(read_i64, i64, read_i64_leb128);
372 read_leb128!(read_i32, i32, read_i32_leb128);
373
374 #[inline]
375 fn read_i16(&mut self) -> i16 {
376 i16::from_le_bytes(self.read_array())
377 }
378
379 #[inline]
380 fn read_raw_bytes(&mut self, bytes: usize) -> &'a [u8] {
381 if bytes > self.remaining() {
382 Self::decoder_exhausted();
383 }
384 unsafe {
386 let slice = std::slice::from_raw_parts(self.current, bytes);
387 self.current = self.current.add(bytes);
388 slice
389 }
390 }
391
392 #[inline]
393 fn peek_byte(&self) -> u8 {
394 if self.current == self.end {
395 Self::decoder_exhausted();
396 }
397 unsafe { *self.current }
400 }
401
402 #[inline]
403 fn position(&self) -> usize {
404 unsafe { self.current.offset_from_unsigned(self.start) }
406 }
407}
408
409impl Encodable<FileEncoder> for [u8] {
417 fn encode(&self, e: &mut FileEncoder) {
418 Encoder::emit_usize(e, self.len());
419 e.emit_raw_bytes(self);
420 }
421}
422
423impl<'a> Decodable<MemDecoder<'a>> for Vec<u8> {
426 fn decode(d: &mut MemDecoder<'a>) -> Self {
427 let len = Decoder::read_usize(d);
428 d.read_raw_bytes(len).to_owned()
429 }
430}
431
432pub struct IntEncodedWithFixedSize(pub u64);
434
435impl IntEncodedWithFixedSize {
436 pub const ENCODED_SIZE: usize = 8;
437}
438
439impl Encodable<FileEncoder> for IntEncodedWithFixedSize {
440 #[inline]
441 fn encode(&self, e: &mut FileEncoder) {
442 let start_pos = e.position();
443 e.write_array(self.0.to_le_bytes());
444 let end_pos = e.position();
445 debug_assert_eq!((end_pos - start_pos), IntEncodedWithFixedSize::ENCODED_SIZE);
446 }
447}
448
449impl<'a> Decodable<MemDecoder<'a>> for IntEncodedWithFixedSize {
450 #[inline]
451 fn decode(decoder: &mut MemDecoder<'a>) -> IntEncodedWithFixedSize {
452 let bytes = decoder.read_array::<{ IntEncodedWithFixedSize::ENCODED_SIZE }>();
453 IntEncodedWithFixedSize(u64::from_le_bytes(bytes))
454 }
455}
456
457#[cfg(test)]
458mod tests;