A better Rust ATProto crate
1// Buffered reader types, pulling HEAVILY from the std library and tokio
2// Single reader supporting sync and async traits
3
4use core::pin::Pin;
5
6pub use buffer::BorrowedCursor;
7pub use buffer::Buffer;
8
9use crate::io::{self, BufRead, DEFAULT_BUF_SIZE, ErrorKind, Read, Seek, SeekFrom};
10
11/// The `BufReader` struct adds buffering to any reader.
12///
13/// It can be excessively inefficient to work directly with a [`AsyncRead`]
14/// instance. A `BufReader` performs large, infrequent reads on the underlying
15/// [`AsyncRead`] and maintains an in-memory buffer of the results.
16///
17/// `BufReader` can improve the speed of programs that make *small* and
18/// *repeated* read calls to the same file or network socket. It does not
19/// help when reading very large amounts at once, or reading just one or a few
20/// times. It also provides no advantage when reading from a source that is
21/// already in memory, like a `Vec<u8>`.
22///
23/// When the `BufReader` is dropped, the contents of its buffer will be
24/// discarded. Creating multiple instances of a `BufReader` on the same
25/// stream can cause data loss.
26#[pin_project::pin_project]
27pub struct BufReader<R: ?Sized> {
28 // for async buffered read ops
29 pub(super) seek_state: SeekState,
30 pub(super) buf: Buffer,
31 #[pin]
32 pub(super) inner: R,
33}
34
35#[derive(Debug, Clone, Copy)]
36pub(super) enum SeekState {
37 /// `start_seek` has not been called.
38 Init,
39 /// `start_seek` has been called, but `poll_complete` has not yet been called.
40 Start(SeekFrom),
41 /// Waiting for completion of the first `poll_complete` in the `n.checked_sub(remainder).is_none()` branch.
42 PendingOverflowed(i64),
43 /// Waiting for completion of `poll_complete`.
44 Pending,
45}
46
47impl<R: Read> BufReader<R> {
48 /// Creates a new `BufReader` with a default buffer capacity. The default is currently 8 KB,
49 /// but may change in the future.
50 pub fn new(inner: R) -> Self {
51 Self::with_capacity(DEFAULT_BUF_SIZE, inner)
52 }
53
54 /// Creates a new `BufReader` with the specified buffer capacity.
55 pub fn with_capacity(capacity: usize, inner: R) -> Self {
56 Self {
57 inner,
58 buf: Buffer::with_capacity(capacity),
59 seek_state: SeekState::Init,
60 }
61 }
62}
63
64impl<R: ?Sized> BufReader<R> {
65 /// Gets a reference to the underlying reader.
66 ///
67 /// It is inadvisable to directly read from the underlying reader.
68 pub fn get_ref(&self) -> &R {
69 &self.inner
70 }
71
72 /// Gets a mutable reference to the underlying reader.
73 ///
74 /// It is inadvisable to directly read from the underlying reader.
75 pub fn get_mut(&mut self) -> &mut R {
76 &mut self.inner
77 }
78
79 /// Gets a pinned mutable reference to the underlying reader.
80 ///
81 /// It is inadvisable to directly read from the underlying reader.
82 pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> {
83 self.project().inner
84 }
85
86 /// Returns a reference to the internally buffered data.
87 ///
88 /// Unlike `fill_buf`, this will not attempt to fill the buffer if it is empty.
89 pub fn buffer(&self) -> &[u8] {
90 &self.buf.buffer()[self.buf.pos()..self.buf.capacity()]
91 }
92
93 /// Consumes this `BufReader`, returning the underlying reader.
94 ///
95 /// Note that any leftover data in the internal buffer is lost.
96 pub fn into_inner(self) -> R
97 where
98 R: Sized,
99 {
100 self.inner
101 }
102
103 /// Invalidates all data in the internal buffer.
104 #[inline]
105 fn discard_buffer_pinned(self: Pin<&mut Self>) {
106 let me = self.project();
107 me.buf.discard_buffer();
108 }
109
110 /// Invalidates all data in the internal buffer.
111 #[inline]
112 fn discard_buffer(&mut self) {
113 self.buf.discard_buffer();
114 }
115
116 pub fn capacity(&self) -> usize {
117 self.buf.capacity()
118 }
119}
120
121impl<R: Read + ?Sized> BufReader<R> {
122 pub fn peek(&mut self, n: usize) -> io::Result<&[u8]> {
123 assert!(n <= self.capacity());
124 while n > self.buf.buffer().len() {
125 if self.buf.pos() > 0 {
126 self.buf.backshift();
127 }
128 let new = self.buf.read_more(&mut self.inner)?;
129 if new == 0 {
130 // end of file, no more bytes to read
131 return Ok(&self.buf.buffer()[..]);
132 }
133 debug_assert_eq!(self.buf.pos(), 0);
134 }
135 Ok(&self.buf.buffer()[..n])
136 }
137}
138
139impl<R: Read + ?Sized> Read for BufReader<R> {
140 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
141 // If we don't have any buffered data and we're doing a massive read
142 // (larger than our internal buffer), bypass our internal buffer
143 // entirely.
144 if self.buf.pos() == self.buf.filled() && buf.len() >= self.capacity() {
145 self.discard_buffer();
146 return self.inner.read(buf);
147 }
148 let mut rem = self.fill_buf()?;
149 let nread = rem.read(buf)?;
150 self.consume(nread);
151 Ok(nread)
152 }
153
154 fn read_exact(&mut self, mut buf: &mut [u8]) -> io::Result<()> {
155 if self
156 .buf
157 .consume_with(buf.len(), |claimed| buf.copy_from_slice(claimed))
158 {
159 return Ok(());
160 }
161 while !buf.is_empty() {
162 match self.read(buf) {
163 Ok(0) => break,
164 Ok(n) => {
165 buf = &mut buf[n..];
166 }
167 Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
168 Err(e) => return Err(e),
169 }
170 }
171 if !buf.is_empty() {
172 Err(io::Error::new(
173 io::ErrorKind::UnexpectedEof,
174 "failed to fill whole buffer".into(),
175 ))
176 } else {
177 Ok(())
178 }
179 }
180}
181
182impl<R: Read + ?Sized> BufRead for BufReader<R> {
183 fn fill_buf(&mut self) -> io::Result<&[u8]> {
184 self.buf.fill_buf(&mut self.inner)
185 }
186
187 fn consume(&mut self, amt: usize) {
188 self.buf.consume(amt)
189 }
190}
191
192impl<R: ?Sized + Seek> BufReader<R> {
193 /// Seeks relative to the current position. If the new position lies within the buffer,
194 /// the buffer will not be flushed, allowing for more efficient seeks.
195 /// This method does not return the location of the underlying reader, so the caller
196 /// must track this information themselves if it is required.
197 pub fn seek_relative(&mut self, offset: i64) -> io::Result<()> {
198 let pos = self.buf.pos() as u64;
199 if offset < 0 {
200 if let Some(_) = pos.checked_sub((-offset) as u64) {
201 self.buf.unconsume((-offset) as usize);
202 return Ok(());
203 }
204 } else if let Some(new_pos) = pos.checked_add(offset as u64) {
205 if new_pos <= self.buf.filled() as u64 {
206 self.buf.consume(offset as usize);
207 return Ok(());
208 }
209 }
210
211 self.seek(SeekFrom::Current(offset)).map(drop)
212 }
213}
214
215impl<R: ?Sized + Seek> Seek for BufReader<R> {
216 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
217 let result: u64;
218 if let SeekFrom::Current(n) = pos {
219 let remainder = (self.buf.filled() - self.buf.pos()) as i64;
220 // it should be safe to assume that remainder fits within an i64 as the alternative
221 // means we managed to allocate 8 exbibytes and that's absurd.
222 // But it's not out of the realm of possibility for some weird underlying reader to
223 // support seeking by i64::MIN so we need to handle underflow when subtracting
224 // remainder.
225 if let Some(offset) = n.checked_sub(remainder) {
226 result = self.inner.seek(SeekFrom::Current(offset))?;
227 } else {
228 // seek backwards by our remainder, and then by the offset
229 self.inner.seek(SeekFrom::Current(-remainder))?;
230 self.discard_buffer();
231 result = self.inner.seek(SeekFrom::Current(n))?;
232 }
233 } else {
234 // Seeking with Start/End doesn't care about our buffer length.
235 result = self.inner.seek(pos)?;
236 }
237 self.discard_buffer();
238 Ok(result)
239 }
240
241 fn stream_position(&mut self) -> io::Result<u64> {
242 let remainder = (self.buf.filled() - self.buf.pos()) as u64;
243 self.inner.stream_position().map(|pos| {
244 pos.checked_sub(remainder).expect(
245 "overflow when subtracting remaining buffer size from inner stream position",
246 )
247 })
248 }
249}
250
251mod buffer {
252 //! An encapsulation of `BufReader`'s buffer management logic.
253 //!
254 //! This module factors out the basic functionality of `BufReader` in order to protect two core
255 //! invariants:
256 //! * `filled` bytes of `buf` are always initialized
257 //! * `pos` is always <= `filled`
258 //! Since this module encapsulates the buffer management logic, we can ensure that the range
259 //! `pos..filled` is always a valid index into the initialized region of the buffer. This means
260 //! that user code which wants to do reads from a `BufReader` via `buffer` + `consume` can do so
261 //! without encountering any runtime bounds checks.
262
263 use crate::io::{self, Read};
264 use core::mem::MaybeUninit;
265 use core::{cmp, fmt};
266
267 pub struct Buffer {
268 // The buffer.
269 buf: Box<[MaybeUninit<u8>]>,
270 // The current seek offset into `buf`, must always be <= `filled`.
271 pos: usize,
272 // Each call to `fill_buf` sets `filled` to indicate how many bytes at the start of `buf` are
273 // initialized with bytes from a read.
274 filled: usize,
275 // This is the max number of bytes returned across all `fill_buf` calls. We track this so that we
276 // can accurately tell `read_buf` how many bytes of buf are initialized, to bypass as much of its
277 // defensive initialization as possible. Note that while this often the same as `filled`, it
278 // doesn't need to be. Calls to `fill_buf` are not required to actually fill the buffer, and
279 // omitting this is a huge perf regression for `Read` impls that do not.
280 initialized: usize,
281 }
282
283 impl Buffer {
284 #[inline]
285 pub fn with_capacity(capacity: usize) -> Self {
286 let buf = Box::new_uninit_slice(capacity);
287 Self {
288 buf,
289 pos: 0,
290 filled: 0,
291 initialized: 0,
292 }
293 }
294
295 #[inline]
296 pub fn buffer(&self) -> &[u8] {
297 // SAFETY: self.pos and self.filled are valid, and self.filled >= self.pos, and
298 // that region is initialized because those are all invariants of this type.
299 unsafe {
300 self.buf
301 .get_unchecked(self.pos..self.filled)
302 .assume_init_ref()
303 }
304 }
305
306 #[inline]
307 pub fn capacity(&self) -> usize {
308 self.buf.len()
309 }
310
311 #[inline]
312 pub fn filled(&self) -> usize {
313 self.filled
314 }
315
316 #[inline]
317 pub fn pos(&self) -> usize {
318 self.pos
319 }
320
321 // This is only used by a test which asserts that the initialization-tracking is correct.
322 #[cfg(test)]
323 pub fn initialized(&self) -> usize {
324 self.initialized
325 }
326
327 #[inline]
328 pub fn discard_buffer(&mut self) {
329 self.pos = 0;
330 self.filled = 0;
331 }
332
333 #[inline]
334 pub fn consume(&mut self, amt: usize) {
335 self.pos = cmp::min(self.pos + amt, self.filled);
336 }
337
338 /// If there are `amt` bytes available in the buffer, pass a slice containing those bytes to
339 /// `visitor` and return true. If there are not enough bytes available, return false.
340 #[inline]
341 pub fn consume_with<V>(&mut self, amt: usize, mut visitor: V) -> bool
342 where
343 V: FnMut(&[u8]),
344 {
345 if let Some(claimed) = self.buffer().get(..amt) {
346 visitor(claimed);
347 // If the indexing into self.buffer() succeeds, amt must be a valid increment.
348 self.pos += amt;
349 true
350 } else {
351 false
352 }
353 }
354
355 #[inline]
356 pub fn unconsume(&mut self, amt: usize) {
357 self.pos = self.pos.saturating_sub(amt);
358 }
359
360 /// Read more bytes into the buffer without discarding any of its contents
361 pub fn read_more(&mut self, mut reader: impl Read) -> io::Result<usize> {
362 let mut buf = ReadBuf::new(&mut unsafe { self.buf.assume_init_mut() }[self.filled..]);
363 let old_init = self.initialized - self.filled;
364 unsafe {
365 buf.assume_init(old_init);
366 }
367 reader.read_buf(buf.unfilled())?;
368 self.filled += buf.len();
369 self.initialized += buf.init_len() - old_init;
370 Ok(buf.len())
371 }
372
373 /// Remove bytes that have already been read from the buffer.
374 pub fn backshift(&mut self) {
375 self.buf.copy_within(self.pos..self.filled, 0);
376 self.filled -= self.pos;
377 self.pos = 0;
378 }
379
380 #[inline]
381 pub fn fill_buf(&mut self, mut reader: impl Read) -> io::Result<&[u8]> {
382 // If we've reached the end of our internal buffer then we need to fetch
383 // some more data from the reader.
384 // Branch using `>=` instead of the more correct `==`
385 // to tell the compiler that the pos..cap slice is always valid.
386 if self.pos >= self.filled {
387 debug_assert!(self.pos == self.filled);
388
389 let mut buf = ReadBuf::new(unsafe { self.buf.assume_init_mut() });
390 // SAFETY: `self.filled` bytes will always have been initialized.
391 unsafe {
392 buf.assume_init(self.initialized);
393 }
394
395 let result = reader.read_buf(buf.unfilled());
396
397 self.pos = 0;
398 self.filled = buf.len();
399 self.initialized = buf.init_len();
400
401 result?;
402 }
403 Ok(self.buffer())
404 }
405 }
406
407 // A wrapper around a byte buffer that is incrementally filled and initialized.
408 ///
409 /// This type is a sort of "double cursor". It tracks three regions in the
410 /// buffer: a region at the beginning of the buffer that has been logically
411 /// filled with data, a region that has been initialized at some point but not
412 /// yet logically filled, and a region at the end that may be uninitialized.
413 /// The filled region is guaranteed to be a subset of the initialized region.
414 ///
415 /// In summary, the contents of the buffer can be visualized as:
416 ///
417 /// ```not_rust
418 /// [ capacity ]
419 /// [ filled | unfilled ]
420 /// [ initialized | uninitialized ]
421 /// ```
422 ///
423 /// It is undefined behavior to de-initialize any bytes from the uninitialized
424 /// region, since it is merely unknown whether this region is uninitialized or
425 /// not, and if part of it turns out to be initialized, it must stay initialized.
426 pub struct ReadBuf<'a> {
427 buf: &'a mut [MaybeUninit<u8>],
428 filled: usize,
429 initialized: usize,
430 }
431
432 impl<'a> ReadBuf<'a> {
433 /// Creates a new `ReadBuf` from a fully initialized buffer.
434 #[inline]
435 pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> {
436 let initialized = buf.len();
437 let buf = unsafe { slice_to_uninit_mut(buf) };
438 ReadBuf {
439 buf,
440 filled: 0,
441 initialized,
442 }
443 }
444
445 /// Creates a new `ReadBuf` from a buffer that may be uninitialized.
446 ///
447 /// The internal cursor will mark the entire buffer as uninitialized. If
448 /// the buffer is known to be partially initialized, then use `assume_init`
449 /// to move the internal cursor.
450 #[inline]
451 pub fn uninit(buf: &'a mut [MaybeUninit<u8>]) -> ReadBuf<'a> {
452 ReadBuf {
453 buf,
454 filled: 0,
455 initialized: 0,
456 }
457 }
458
459 #[inline]
460 pub fn len(&self) -> usize {
461 self.filled
462 }
463
464 /// Returns the length of the initialized part of the buffer.
465 #[inline]
466 pub fn init_len(&self) -> usize {
467 self.initialized
468 }
469
470 /// Returns the total capacity of the buffer.
471 #[inline]
472 pub fn capacity(&self) -> usize {
473 self.buf.len()
474 }
475
476 /// Returns a shared reference to the filled portion of the buffer.
477 #[inline]
478 pub fn filled(&self) -> &[u8] {
479 let slice = &self.buf[..self.filled];
480 // safety: filled describes how far into the buffer that the
481 // user has filled with bytes, so it's been initialized.
482 unsafe { slice_assume_init(slice) }
483 }
484
485 /// Returns a mutable reference to the filled portion of the buffer.
486 #[inline]
487 pub fn filled_mut(&mut self) -> &mut [u8] {
488 let slice = &mut self.buf[..self.filled];
489 // safety: filled describes how far into the buffer that the
490 // user has filled with bytes, so it's been initialized.
491 unsafe { slice_assume_init_mut(slice) }
492 }
493
494 /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`.
495 #[inline]
496 pub fn take(&mut self, n: usize) -> ReadBuf<'_> {
497 let max = std::cmp::min(self.remaining(), n);
498 // Safety: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`.
499 unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) }
500 }
501
502 /// Returns a shared reference to the initialized portion of the buffer.
503 ///
504 /// This includes the filled portion.
505 #[inline]
506 pub fn initialized(&self) -> &[u8] {
507 let slice = &self.buf[..self.initialized];
508 // safety: initialized describes how far into the buffer that the
509 // user has at some point initialized with bytes.
510 unsafe { slice_assume_init(slice) }
511 }
512
513 /// Returns a mutable reference to the initialized portion of the buffer.
514 ///
515 /// This includes the filled portion.
516 #[inline]
517 pub fn initialized_mut(&mut self) -> &mut [u8] {
518 let slice = &mut self.buf[..self.initialized];
519 // safety: initialized describes how far into the buffer that the
520 // user has at some point initialized with bytes.
521 unsafe { slice_assume_init_mut(slice) }
522 }
523
524 /// Returns a mutable reference to the entire buffer, without ensuring that it has been fully
525 /// initialized.
526 ///
527 /// The elements between 0 and `self.filled().len()` are filled, and those between 0 and
528 /// `self.initialized().len()` are initialized (and so can be converted to a `&mut [u8]`).
529 ///
530 /// The caller of this method must ensure that these invariants are upheld. For example, if the
531 /// caller initializes some of the uninitialized section of the buffer, it must call
532 /// [`assume_init`](Self::assume_init) with the number of bytes initialized.
533 ///
534 /// # Safety
535 ///
536 /// The caller must not de-initialize portions of the buffer that have already been initialized.
537 /// This includes any bytes in the region marked as uninitialized by `ReadBuf`.
538 #[inline]
539 pub unsafe fn inner_mut(&mut self) -> &mut [MaybeUninit<u8>] {
540 self.buf
541 }
542
543 /// Returns a mutable reference to the unfilled part of the buffer without ensuring that it has been fully
544 /// initialized.
545 ///
546 /// # Safety
547 ///
548 /// The caller must not de-initialize portions of the buffer that have already been initialized.
549 /// This includes any bytes in the region marked as uninitialized by `ReadBuf`.
550 #[inline]
551 pub unsafe fn unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>] {
552 &mut self.buf[self.filled..]
553 }
554
555 /// Returns a mutable reference to the unfilled part of the buffer, ensuring it is fully initialized.
556 ///
557 /// Since `ReadBuf` tracks the region of the buffer that has been initialized, this is effectively "free" after
558 /// the first use.
559 #[inline]
560 pub fn initialize_unfilled(&mut self) -> &mut [u8] {
561 self.initialize_unfilled_to(self.remaining())
562 }
563
564 /// Returns a mutable reference to the first `n` bytes of the unfilled part of the buffer, ensuring it is
565 /// fully initialized.
566 ///
567 /// # Panics
568 ///
569 /// Panics if `self.remaining()` is less than `n`.
570 #[inline]
571 #[track_caller]
572 pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] {
573 assert!(self.remaining() >= n, "n overflows remaining");
574
575 // This can't overflow, otherwise the assert above would have failed.
576 let end = self.filled + n;
577
578 if self.initialized < end {
579 unsafe {
580 self.buf[self.initialized..end]
581 .as_mut_ptr()
582 .write_bytes(0, end - self.initialized);
583 }
584 self.initialized = end;
585 }
586
587 let slice = &mut self.buf[self.filled..end];
588 // safety: just above, we checked that the end of the buf has
589 // been initialized to some value.
590 unsafe { slice_assume_init_mut(slice) }
591 }
592
593 /// Returns a cursor over the unfilled part of the buffer.
594 #[inline]
595 pub fn unfilled<'this>(&'this mut self) -> BorrowedCursor<'this> {
596 BorrowedCursor {
597 // SAFETY: we never assign into `BorrowedCursor::buf`, so treating its
598 // lifetime covariantly is safe.
599 buf: unsafe {
600 core::mem::transmute::<&'this mut ReadBuf<'a>, &'this mut ReadBuf<'this>>(self)
601 },
602 }
603 }
604
605 /// Returns the number of bytes at the end of the slice that have not yet been filled.
606 #[inline]
607 pub fn remaining(&self) -> usize {
608 self.capacity() - self.filled
609 }
610
611 /// Clears the buffer, resetting the filled region to empty.
612 ///
613 /// The number of initialized bytes is not changed, and the contents of the buffer are not modified.
614 #[inline]
615 pub fn clear(&mut self) {
616 self.filled = 0;
617 }
618
619 /// Advances the size of the filled region of the buffer.
620 ///
621 /// The number of initialized bytes is not changed.
622 ///
623 /// # Panics
624 ///
625 /// Panics if the filled region of the buffer would become larger than the initialized region.
626 #[inline]
627 #[track_caller]
628 pub fn advance(&mut self, n: usize) {
629 let new = self.filled.checked_add(n).expect("filled overflow");
630 self.set_filled(new);
631 }
632
633 /// Sets the size of the filled region of the buffer.
634 ///
635 /// The number of initialized bytes is not changed.
636 ///
637 /// Note that this can be used to *shrink* the filled region of the buffer in addition to growing it (for
638 /// example, by a `AsyncRead` implementation that compresses data in-place).
639 ///
640 /// # Panics
641 ///
642 /// Panics if the filled region of the buffer would become larger than the initialized region.
643 #[inline]
644 #[track_caller]
645 pub fn set_filled(&mut self, n: usize) {
646 assert!(
647 n <= self.initialized,
648 "filled must not become larger than initialized"
649 );
650 self.filled = n;
651 }
652
653 /// Asserts that the first `n` unfilled bytes of the buffer are initialized.
654 ///
655 /// `ReadBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer
656 /// bytes than are already known to be initialized.
657 ///
658 /// # Safety
659 ///
660 /// The caller must ensure that `n` unfilled bytes of the buffer have already been initialized.
661 #[inline]
662 pub unsafe fn assume_init(&mut self, n: usize) {
663 let new = self.filled + n;
664 if new > self.initialized {
665 self.initialized = new;
666 }
667 }
668
669 /// Appends data to the buffer, advancing the written position and possibly also the initialized position.
670 ///
671 /// # Panics
672 ///
673 /// Panics if `self.remaining()` is less than `buf.len()`.
674 #[inline]
675 #[track_caller]
676 pub fn put_slice(&mut self, buf: &[u8]) {
677 assert!(
678 self.remaining() >= buf.len(),
679 "buf.len() must fit in remaining(); buf.len() = {}, remaining() = {}",
680 buf.len(),
681 self.remaining()
682 );
683
684 let amt = buf.len();
685 // Cannot overflow, asserted above
686 let end = self.filled + amt;
687
688 // Safety: the length is asserted above
689 unsafe {
690 self.buf[self.filled..end]
691 .as_mut_ptr()
692 .cast::<u8>()
693 .copy_from_nonoverlapping(buf.as_ptr(), amt);
694 }
695
696 if self.initialized < end {
697 self.initialized = end;
698 }
699 self.filled = end;
700 }
701 }
702
703 impl fmt::Debug for ReadBuf<'_> {
704 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
705 f.debug_struct("ReadBuf")
706 .field("filled", &self.filled)
707 .field("initialized", &self.initialized)
708 .field("capacity", &self.capacity())
709 .finish()
710 }
711 }
712
713 /// # Safety
714 ///
715 /// The caller must ensure that `slice` is fully initialized
716 /// and never writes uninitialized bytes to the returned slice.
717 unsafe fn slice_to_uninit_mut(slice: &mut [u8]) -> &mut [MaybeUninit<u8>] {
718 // SAFETY: `MaybeUninit<u8>` has the same memory layout as u8, and the caller
719 // promises to not write uninitialized bytes to the returned slice.
720 unsafe { &mut *(slice as *mut [u8] as *mut [MaybeUninit<u8>]) }
721 }
722
723 /// # Safety
724 ///
725 /// The caller must ensure that `slice` is fully initialized.
726 // TODO: This could use `MaybeUninit::slice_assume_init` when it is stable.
727 unsafe fn slice_assume_init(slice: &[MaybeUninit<u8>]) -> &[u8] {
728 // SAFETY: `MaybeUninit<u8>` has the same memory layout as u8, and the caller
729 // promises that `slice` is fully initialized.
730 unsafe { &*(slice as *const [MaybeUninit<u8>] as *const [u8]) }
731 }
732
733 /// # Safety
734 ///
735 /// The caller must ensure that `slice` is fully initialized.
736 // TODO: This could use `MaybeUninit::slice_assume_init_mut` when it is stable.
737 unsafe fn slice_assume_init_mut(slice: &mut [MaybeUninit<u8>]) -> &mut [u8] {
738 // SAFETY: `MaybeUninit<u8>` has the same memory layout as `u8`, and the caller
739 // promises that `slice` is fully initialized.
740 unsafe { &mut *(slice as *mut [MaybeUninit<u8>] as *mut [u8]) }
741 }
742
743 #[derive(Debug)]
744 pub struct BorrowedCursor<'a> {
745 /// The underlying buffer.
746 // Safety invariant: we treat the type of buf as covariant in the lifetime of `BorrowedBuf` when
747 // we create a `BorrowedCursor`. This is only safe if we never replace `buf` by assigning into
748 // it, so don't do that!
749 buf: &'a mut ReadBuf<'a>,
750 }
751
752 impl<'a> BorrowedCursor<'a> {
753 /// Reborrows this cursor by cloning it with a smaller lifetime.
754 ///
755 /// Since a cursor maintains unique access to its underlying buffer, the borrowed cursor is
756 /// not accessible while the new cursor exists.
757 #[inline]
758 pub fn reborrow<'this>(&'this mut self) -> BorrowedCursor<'this> {
759 BorrowedCursor {
760 // SAFETY: we never assign into `BorrowedCursor::buf`, so treating its
761 // lifetime covariantly is safe.
762 buf: unsafe {
763 core::mem::transmute::<&'this mut ReadBuf<'a>, &'this mut ReadBuf<'this>>(
764 self.buf,
765 )
766 },
767 }
768 }
769
770 /// Returns the available space in the cursor.
771 #[inline]
772 pub fn capacity(&self) -> usize {
773 self.buf.capacity() - self.buf.filled
774 }
775
776 /// Returns the number of bytes written to the `BorrowedBuf` this cursor was created from.
777 ///
778 /// In particular, the count returned is shared by all reborrows of the cursor.
779 #[inline]
780 pub fn written(&self) -> usize {
781 self.buf.filled
782 }
783
784 /// Returns a mutable reference to the initialized portion of the cursor.
785 #[inline]
786 pub fn init_mut(&mut self) -> &mut [u8] {
787 // SAFETY: We only slice the initialized part of the buffer, which is always valid
788 unsafe {
789 let buf = self
790 .buf
791 .buf
792 .get_unchecked_mut(self.buf.filled..self.buf.initialized);
793 buf.assume_init_mut()
794 }
795 }
796
797 /// Returns a mutable reference to the whole cursor.
798 ///
799 /// # Safety
800 ///
801 /// The caller must not uninitialize any bytes in the initialized portion of the cursor.
802 #[inline]
803 pub unsafe fn as_mut(&mut self) -> &mut [MaybeUninit<u8>] {
804 // SAFETY: always in bounds
805 unsafe { self.buf.buf.get_unchecked_mut(self.buf.filled..) }
806 }
807
808 /// Advances the cursor by asserting that `n` bytes have been filled.
809 ///
810 /// After advancing, the `n` bytes are no longer accessible via the cursor and can only be
811 /// accessed via the underlying buffer. I.e., the buffer's filled portion grows by `n` elements
812 /// and its unfilled portion (and the capacity of this cursor) shrinks by `n` elements.
813 ///
814 /// If less than `n` bytes initialized (by the cursor's point of view), `set_init` should be
815 /// called first.
816 ///
817 /// # Panics
818 ///
819 /// Panics if there are less than `n` bytes initialized.
820 #[inline]
821 pub fn advance(&mut self, n: usize) -> &mut Self {
822 // The subtraction cannot underflow by invariant of this type.
823 assert!(n <= self.buf.initialized - self.buf.filled);
824
825 self.buf.filled += n;
826 self
827 }
828
829 /// Advances the cursor by asserting that `n` bytes have been filled.
830 ///
831 /// After advancing, the `n` bytes are no longer accessible via the cursor and can only be
832 /// accessed via the underlying buffer. I.e., the buffer's filled portion grows by `n` elements
833 /// and its unfilled portion (and the capacity of this cursor) shrinks by `n` elements.
834 ///
835 /// # Safety
836 ///
837 /// The caller must ensure that the first `n` bytes of the cursor have been properly
838 /// initialised.
839 #[inline]
840 pub unsafe fn advance_unchecked(&mut self, n: usize) -> &mut Self {
841 self.buf.filled += n;
842 self.buf.initialized = cmp::max(self.buf.initialized, self.buf.filled);
843 self
844 }
845
846 /// Initializes all bytes in the cursor.
847 #[inline]
848 pub fn ensure_init(&mut self) -> &mut Self {
849 // SAFETY: always in bounds and we never uninitialize these bytes.
850 let uninit = unsafe { self.buf.buf.get_unchecked_mut(self.buf.initialized..) };
851
852 // SAFETY: 0 is a valid value for MaybeUninit<u8> and the length matches the allocation
853 // since it is comes from a slice reference.
854 unsafe {
855 core::ptr::write_bytes(uninit.as_mut_ptr(), 0, uninit.len());
856 }
857 self.buf.initialized = self.buf.capacity();
858
859 self
860 }
861
862 /// Asserts that the first `n` unfilled bytes of the cursor are initialized.
863 ///
864 /// `BorrowedBuf` assumes that bytes are never de-initialized, so this method does nothing when
865 /// called with fewer bytes than are already known to be initialized.
866 ///
867 /// # Safety
868 ///
869 /// The caller must ensure that the first `n` bytes of the buffer have already been initialized.
870 #[inline]
871 pub unsafe fn set_init(&mut self, n: usize) -> &mut Self {
872 self.buf.initialized = cmp::max(self.buf.initialized, self.buf.filled + n);
873 self
874 }
875
876 /// Appends data to the cursor, advancing position within its buffer.
877 ///
878 /// # Panics
879 ///
880 /// Panics if `self.capacity()` is less than `buf.len()`.
881 #[inline]
882 pub fn append(&mut self, buf: &[u8]) {
883 assert!(self.capacity() >= buf.len());
884
885 // SAFETY: we do not de-initialize any of the elements of the slice
886 unsafe {
887 self.as_mut()[..buf.len()].write_copy_of_slice(buf);
888 }
889
890 // SAFETY: We just added the entire contents of buf to the filled section.
891 unsafe {
892 self.set_init(buf.len());
893 }
894 self.buf.filled += buf.len();
895 }
896
897 /// Runs the given closure with a `BorrowedBuf` containing the unfilled part
898 /// of the cursor.
899 ///
900 /// This enables inspecting what was written to the cursor.
901 ///
902 /// # Panics
903 ///
904 /// Panics if the `BorrowedBuf` given to the closure is replaced by another
905 /// one.
906 pub fn with_unfilled_buf<T>(&mut self, f: impl FnOnce(&mut ReadBuf<'_>) -> T) -> T {
907 let mut buf = ReadBuf::from(self.reborrow());
908 let prev_ptr = buf.buf as *const _;
909 let res = f(&mut buf);
910
911 // Check that the caller didn't replace the `BorrowedBuf`.
912 // This is necessary for the safety of the code below: if the check wasn't
913 // there, one could mark some bytes as initialized even though there aren't.
914 assert!(core::ptr::addr_eq(prev_ptr, buf.buf));
915
916 let filled = buf.filled;
917 let init = buf.initialized;
918
919 // Update `init` and `filled` fields with what was written to the buffer.
920 // `self.buf.filled` was the starting length of the `BorrowedBuf`.
921 //
922 // SAFETY: These amounts of bytes were initialized/filled in the `BorrowedBuf`,
923 // and therefore they are initialized/filled in the cursor too, because the
924 // buffer wasn't replaced.
925 self.buf.initialized = self.buf.filled + init;
926 self.buf.filled += filled;
927
928 res
929 }
930 }
931
932 impl<'data> From<BorrowedCursor<'data>> for ReadBuf<'data> {
933 #[inline]
934 fn from(mut buf: BorrowedCursor<'data>) -> ReadBuf<'data> {
935 let initialized = buf.init_mut().len();
936 ReadBuf {
937 // SAFETY: no initialized byte is ever uninitialized as per
938 // `BorrowedBuf`'s invariant
939 buf: unsafe { buf.buf.buf.get_unchecked_mut(buf.buf.filled..) },
940 filled: 0,
941 initialized,
942 }
943 }
944 }
945}